leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 26 Packages Projects Releases Wiki Activity

v4.3 tag release update. (#2789)

Browse Source
This commit is contained in:
Junkai-Wu
2025-11-20 20:49:44 -05:00
committed by GitHub
parent 406e078b29
commit 8cd5bef43a
225 changed files with 23229 additions and 2813 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/library/include/cutlass/library/descriptions.h
+1
View File
@@ -328,6 +328,7 @@ struct BlockScaleDescription {
struct GroupedGemmDescription : public OperationDescription {
GemmDescription gemm;
std::optional<BlockScaleDescription> block_scales;
bool is_moe{false};
};
/// Description of all GEMM computations
tools/library/include/cutlass/library/library.h
+3
View File
@@ -592,6 +592,9 @@ struct GemmGroupedArguments {
// underlying operation uses the one it needs.
cute::Shape<int, int, int>* problem_sizes_3x;
cute::Shape<int, int, int>* problem_sizes_3x_host;
std::vector<int32_t> max_problem_size_3x;
int32_t* tokens_per_expert{nullptr};
int32_t* tokens_per_expert_host{nullptr};
};
struct GroupedGemmBlockScaledArguments : GemmGroupedArguments {
tools/library/src/grouped_gemm_operation_3x.hpp
+500
View File
@@ -869,5 +869,505 @@ public:
}
};
template <typename Operator_>
class MoeGroupedGemmOperation3xBase : public GemmOperation3xBase<Operator_> {
public:
using Operator = Operator_;
using OperatorArguments = typename Operator::Arguments;
using ElementA = typename Operator::ElementA;
using LayoutA = typename Operator::LayoutA;
using ElementB = typename Operator::ElementB;
using LayoutB = typename Operator::LayoutB;
using ElementC = typename Operator::ElementC;
using LayoutC = typename Operator::LayoutC;
using ElementD = typename Operator::ElementD;
using LayoutD = typename Operator::LayoutD;
using ElementAccumulator = typename Operator::ElementAccumulator;
using ElementCompute = typename Operator::EpilogueOutputOp::ElementCompute;
using ArrayElementA = typename Operator::GemmKernel::CollectiveMainloop::ArrayElementA;
using ArrayElementB = typename Operator::GemmKernel::CollectiveMainloop::ArrayElementB;
using CollectiveMainloop = typename Operator::CollectiveMainloop;
using CollectiveEpilogue = typename Operator::CollectiveEpilogue;
using ThreadEpilogueOp = typename CollectiveEpilogue::ThreadEpilogueOp;
using StrideC = typename Operator::GemmKernel::StrideC;
using StrideD = typename Operator::GemmKernel::StrideD;
static constexpr bool IsRuntimeDataTypeA = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4<ElementA>();
static constexpr bool IsRuntimeDataTypeB = cutlass::gemm::collective::detail::is_sm10x_runtime_f8f6f4<ElementB>();
static_assert((IsRuntimeDataTypeA && IsRuntimeDataTypeB) ||
(!IsRuntimeDataTypeA && !IsRuntimeDataTypeB),
"ElementA and ElementB in a GEMM kernel should be both runtime or both static.");
static constexpr bool IsRuntimeDataType = IsRuntimeDataTypeA && IsRuntimeDataTypeB;
MoeGroupedGemmOperation3xBase(char const* name = "unknown_gemm")
: GemmOperation3xBase<Operator_>(name, GemmKind::kGrouped) {
this->description_.is_moe = true;
this->description_.kind = OperationKind::kGroupedGemm;
this->description_.name = name;
this->description_.provider = Provider::kCUTLASS;
this->description_.gemm = GemmOperation3xBase<Operator_>::description_;
this->description_.tile_description = this->description_.gemm.tile_description;
};
public:
// mutable CudaBuffer strideC_device;
// mutable CudaBuffer strideD_device;
/// Returns the description of the GEMM operation
virtual OperationDescription const& description() const override final { return description_; }
/// Gets the host-side workspace
uint64_t get_host_workspace_size(void const* configuration) const override final {
return sizeof(Operator);
}
protected:
library::GroupedGemmDescription description_;
/// Constructs the arguments structure given the configuration and arguments
Status update_arguments_base(
OperatorArguments& operator_args,
GemmGroupedArguments const& arguments) const {
operator_args.mode = cutlass::gemm::GemmUniversalMode::kGrouped;
int M= arguments.max_problem_size_3x[0];
int N = arguments.max_problem_size_3x[1];
int K = arguments.max_problem_size_3x[2];
int L = arguments.problem_count;
operator_args.problem_shape = {
M,
N,
K,
L,
arguments.tokens_per_expert,
arguments.tokens_per_expert_host
};
if constexpr (IsRuntimeDataType) {
using RuntimeDataTypeA = typename Operator::GemmKernel::CollectiveMainloop::RuntimeDataTypeA;
using RuntimeDataTypeB = typename Operator::GemmKernel::CollectiveMainloop::RuntimeDataTypeB;
static_assert(cute::is_same_v<RuntimeDataTypeA, RuntimeDataTypeB>,
"RuntimeDataTypeA/B should be identical, either MXF8F6F4Format or MXF4Format");
using RuntimeDatatypeArg = RuntimeDataTypeA;
auto mapping = [](RuntimeDatatype type) {
if constexpr (cute::is_same_v<RuntimeDatatypeArg, cute::UMMA::MXF8F6F4Format>) {
if (type == RuntimeDatatype::kE5M2) {
return cute::UMMA::MXF8F6F4Format::E5M2;
}
else if (type == RuntimeDatatype::kE4M3) {
return cute::UMMA::MXF8F6F4Format::E4M3;
}
else if (type == RuntimeDatatype::kE3M2) {
return cute::UMMA::MXF8F6F4Format::E3M2;
}
else if (type == RuntimeDatatype::kE2M3) {
return cute::UMMA::MXF8F6F4Format::E2M3;
}
else if (type == RuntimeDatatype::kE2M1) {
return cute::UMMA::MXF8F6F4Format::E2M1;
}
else {
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && CUTLASS_DEBUG_TRACE_LEVEL >= 1
std::cerr << "Invalid input datatype specified. Running with e4m3." << std::endl;
#endif
return cute::UMMA::MXF8F6F4Format::E4M3;
}
}
else if constexpr (cute::is_same_v<RuntimeDatatypeArg, cute::UMMA::MXF4Format>) {
if (type == RuntimeDatatype::kE2M1) {
return cute::UMMA::MXF4Format::E2M1;
}
else {
#if defined(CUTLASS_DEBUG_TRACE_LEVEL) && CUTLASS_DEBUG_TRACE_LEVEL >= 1
std::cerr << "Invalid input datatype specified. Running with e2m1." << std::endl;
#endif
return cute::UMMA::MXF4Format::E2M1;
}
}
// BlockScaled kernels receive either MXF4Format or MXF8F6F4Format runtime datatype
CUTE_GCC_UNREACHABLE;
};
operator_args.mainloop.runtime_data_type_a = mapping(arguments.runtime_input_datatype_a);
operator_args.mainloop.runtime_data_type_b = mapping(arguments.runtime_input_datatype_b);
}
operator_args.epilogue.ptr_C = static_cast<ElementC const*>(arguments.ptr_C);
operator_args.epilogue.ptr_D = static_cast<ElementD*>(arguments.ptr_D);
operator_args.epilogue.dC = cutlass::make_cute_packed_stride(StrideC{}, cute::make_shape(M, N, L));
operator_args.epilogue.dD = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(M, N, L));
/* Query device SM count and max active clusters to pass onto the kernel as an argument, where needed */
operator_args.hw_info.sm_count = arguments.sm_count;
if constexpr (Operator::ArchTag::kMinComputeCapability >= 90) {
operator_args.hw_info.max_active_clusters = arguments.max_active_clusters;
}
if constexpr (!std::is_const_v<decltype(operator_args.scheduler.max_swizzle_size)>) {
operator_args.scheduler.max_swizzle_size = arguments.swizzle_size;
}
if constexpr (!std::is_const_v<decltype(operator_args.scheduler.raster_order)>) {
using Enum_t = decltype(operator_args.scheduler.raster_order);
switch (arguments.raster_order) {
case RasterOrder::kAlongN:
operator_args.scheduler.raster_order = Enum_t::AlongN;
break;
case RasterOrder::kAlongM:
operator_args.scheduler.raster_order = Enum_t::AlongM;
break;
default:
operator_args.scheduler.raster_order = Enum_t::Heuristic;
}
}
if constexpr (Operator::ArchTag::kMinComputeCapability >= 100) {
operator_args.hw_info.cluster_shape =
dim3(arguments.cluster_shape.m(), arguments.cluster_shape.n(), arguments.cluster_shape.k());
operator_args.hw_info.cluster_shape_fallback = dim3(
arguments.cluster_shape_fallback.m(),
arguments.cluster_shape_fallback.n(),
arguments.cluster_shape_fallback.k());
}
return Status::kSuccess;
}
template <typename FusionArgs>
static Status update_fusion_args(FusionArgs& fusion_args, GemmGroupedArguments const& arguments) {
if (arguments.pointer_mode == ScalarPointerMode::kHost) {
fusion_args.alpha = *static_cast<ElementCompute const*>(arguments.alpha);
fusion_args.beta = *static_cast<ElementCompute const*>(arguments.beta);
fusion_args.alpha_ptr = nullptr;
fusion_args.beta_ptr = nullptr;
return Status::kSuccess;
}
else if (arguments.pointer_mode == ScalarPointerMode::kDevice) {
fusion_args.alpha = 0;
fusion_args.beta = 0;
fusion_args.alpha_ptr = static_cast<ElementCompute const*>(arguments.alpha);
fusion_args.beta_ptr = static_cast<ElementCompute const*>(arguments.beta);
return Status::kSuccess;
}
else {
return Status::kErrorInvalidProblem;
}
}
};
template<typename Operator_>
class MoeGroupedGemmUniversal3xOperation : public MoeGroupedGemmOperation3xBase<Operator_> {
public:
using Base = MoeGroupedGemmOperation3xBase<Operator_>;
using Operator = Operator_;
using OperatorArguments = typename Operator::Arguments;
MoeGroupedGemmUniversal3xOperation(char const* name = "unknown_gemm")
: MoeGroupedGemmOperation3xBase<Operator_>(name) {
}
~MoeGroupedGemmUniversal3xOperation() override = default;
protected:
template <class FusionArgs, class = void> struct UpdateFusionArgs {
static Status update_(FusionArgs const& fusion_args, GemmGroupedArguments const& arguments) {
// If a custom EVT is instantiated then it is the users's responsibility
// to ensure alpha and beta are updated appropriately
return Status::kSuccess;
}
};
template <class FusionArgs>
struct UpdateFusionArgs<FusionArgs, cute::void_t<decltype(FusionArgs{}.alpha)>> {
static Status update_(FusionArgs& fusion_args, GemmGroupedArguments const& arguments) {
return MoeGroupedGemmOperation3xBase<Operator>::update_fusion_args(fusion_args, arguments);
}
};
/// Constructs the arguments structure given the configuration and arguments
Status
update_arguments_(OperatorArguments& operator_args, GemmGroupedArguments const* arguments) const {
Status status = UpdateFusionArgs<decltype(operator_args.epilogue.thread)>::update_(
operator_args.epilogue.thread,
*arguments);
if (status != Status::kSuccess) {
return status;
}
status = this->update_arguments_base(operator_args, *arguments);
operator_args.mainloop.ptr_A = static_cast<typename Base::ArrayElementA const*>(arguments->ptr_A);
operator_args.mainloop.ptr_B = static_cast<typename Base::ArrayElementB const*>(arguments->ptr_B);
return status;
}
public:
/// Returns success if the operation can proceed
Status can_implement([[maybe_unused]] void const* configuration_ptr, void const* arguments_ptr)
const override {
GemmGroupedArguments const* arguments = static_cast<GemmGroupedArguments const*>(arguments_ptr);
OperatorArguments args;
auto status = update_arguments_(args, arguments);
if (status != Status::kSuccess) {
return status;
}
status = Operator::can_implement(args);
return status;
}
/// Gets the device-side workspace
uint64_t get_device_workspace_size(void const* configuration_ptr, void const* arguments_ptr)
const override {
OperatorArguments args;
auto status = update_arguments_(args, static_cast<GemmGroupedArguments const*>(arguments_ptr));
if (status != Status::kSuccess) {
return 0;
}
uint64_t size = Operator::get_workspace_size(args);
return size;
}
/// Initializes the workspace
/// **** CAUTION ****
/// Must be called when ldc, or ldd change.
/// The CUTLASS library stores the operations in a type-
/// erased manifest. Therefore, only this class knows
/// the type of strideC, and strideD.
/// Since grouped GEMM needs to allocate storage for
/// the strides on device, the concrete type of the stride
/// must be known in order to copy in the correct memory
/// layout on device.
Status initialize(
void const* configuration_ptr,
void* host_workspace,
void* device_workspace,
cudaStream_t stream = nullptr) const override {
Operator* op = new (host_workspace) Operator;
return Status::kSuccess;
}
/// **** CAUTION ****
/// initialize() must be called if lda, ldb, ldc, or ldd change.
Status run(
void const* arguments_ptr,
void* host_workspace,
void* device_workspace = nullptr,
cudaStream_t stream = nullptr) const override {
OperatorArguments operator_args;
auto const& args = *static_cast<GemmGroupedArguments const*>(arguments_ptr);
Status status = update_arguments_(operator_args, &args);
if (status != Status::kSuccess) {
return status;
}
Operator* op = static_cast<Operator*>(host_workspace);
// We need to call initialize() since we have to rebuild TMA desc for every new set of args
status = op->run(operator_args, device_workspace, stream, nullptr, args.use_pdl);
return status;
}
// Set arguments that should only be set once before verifying or profiling the kernel.
// This should encompass any expensive operations that don't vary from run to run
// (e.g., max_active_clusters).
Status initialize_with_arguments(void* arguments_ptr) const override {
if constexpr (Operator::ArchTag::kMinComputeCapability < 90) {
return Status::kSuccess;
}
GemmGroupedArguments* args = static_cast<GemmGroupedArguments*>(arguments_ptr);
dim3 cluster_dims;
if constexpr (cute::is_static_v<typename Operator::GemmKernel::ClusterShape>) {
cluster_dims = dim3(
cute::size<0>(typename Operator::GemmKernel::ClusterShape{}),
cute::size<1>(typename Operator::GemmKernel::ClusterShape{}),
cute::size<2>(typename Operator::GemmKernel::ClusterShape{})
);
}
else {
cluster_dims = dim3(
args->cluster_shape.m(),
args->cluster_shape.n(),
args->cluster_shape.k()
);
}
uint32_t threads_per_block = Operator::GemmKernel::MaxThreadsPerBlock;
void const* kernel_ptr = (void*)(device_kernel<typename Operator::GemmKernel>);
args->max_active_clusters = cutlass::KernelHardwareInfo::query_device_max_active_clusters(
cluster_dims,
threads_per_block,
kernel_ptr);
if (args->max_active_clusters == 0) {
std::cerr << "Max Active Clusters could not be queried. "
<< "Falling back to heuristics mode (static cluster shape) or preferred cluster mode.\n";
}
return Status::kSuccess;
}
};
template<typename Operator_>
class BlockScaledMoeGroupedGemmUniversal3xOperation : public MoeGroupedGemmOperation3xBase<Operator_> {
public:
using Base = MoeGroupedGemmOperation3xBase<Operator_>;
using Operator = Operator_;
using OperatorArguments = typename Operator::Arguments;
using ElementD = typename Operator::ElementD;
using LayoutD = typename Operator::LayoutD;
using ElementAccumulator = typename Operator::ElementAccumulator;
using ElementCompute = typename Operator::EpilogueOutputOp::ElementCompute;
using CollectiveMainloop = typename Operator::CollectiveMainloop;
using CollectiveEpilogue = typename Operator::CollectiveEpilogue;
using ThreadEpilogueOp = typename CollectiveEpilogue::ThreadEpilogueOp;
using ElementSF = typename Operator::CollectiveMainloop::ElementSF;
using TiledMma = typename Operator::CollectiveMainloop::TiledMma;
constexpr static int SFVecSize = TiledMma::SFVecSize;
static constexpr bool epilogue_scalefactor_generation = not cute::is_same_v<typename ThreadEpilogueOp::ElementBlockScaleFactor, void>;
static constexpr int32_t SFD_VectorSize = epilogue_scalefactor_generation ? ThreadEpilogueOp::SFVecSize : SFVecSize;
using ElementSFD = cute::conditional_t<epilogue_scalefactor_generation, typename ThreadEpilogueOp::ElementBlockScaleFactor, void>;
using LayoutSFD = cute::conditional_t<epilogue_scalefactor_generation, typename ThreadEpilogueOp::GmemLayoutTagScalefactor, LayoutD>;
BlockScaledMoeGroupedGemmUniversal3xOperation(char const* name = "unknown_gemm")
: MoeGroupedGemmOperation3xBase<Operator_>(name) {
BlockScaleDescription block_scaled_desc{};
block_scaled_desc.kind = OperationKind::kBlockScaledGemm;
block_scaled_desc.SFA.element = NumericTypeMap<ElementSF>::kId;
block_scaled_desc.SFA.layout = LayoutTypeID::kRowMajor;
block_scaled_desc.SFA.alignment = 128;
block_scaled_desc.SFA.log_extent_range = 32;
block_scaled_desc.SFA.log_stride_range = 32;
block_scaled_desc.SFB.element = NumericTypeMap<ElementSF>::kId;
block_scaled_desc.SFB.layout = LayoutTypeID::kRowMajor;
block_scaled_desc.SFB.alignment = 128;
block_scaled_desc.SFB.log_extent_range = 32;
block_scaled_desc.SFB.log_stride_range = 32;
block_scaled_desc.SFMVecSize = 1;
block_scaled_desc.SFNVecSize = 1;
block_scaled_desc.SFKVecSize = SFVecSize;
block_scaled_desc.SFD = make_TensorDescription<ElementSFD, LayoutSFD>(128);
block_scaled_desc.EpilogueSFVecSize = SFD_VectorSize;
this->description_.block_scales = block_scaled_desc;
}
~BlockScaledMoeGroupedGemmUniversal3xOperation() override = default;
protected:
template <class FusionArgs, class = void> struct UpdateFusionArgs {
static Status update_(FusionArgs const& fusion_args, GroupedGemmBlockScaledArguments const& arguments) {
// If a custom EVT is instantiated then it is the users's responsibility
// to ensure alpha and beta are updated appropriately
return Status::kSuccess;
}
};
template <class FusionArgs>
struct UpdateFusionArgs<FusionArgs, cute::void_t<decltype(FusionArgs{}.alpha)>> {
static Status
update_(FusionArgs& fusion_args, GroupedGemmBlockScaledArguments const& arguments) {
if constexpr (epilogue_scalefactor_generation) {
fusion_args.block_scale_factor_ptr = static_cast<ElementSFD*>(arguments.SFD);
fusion_args.norm_constant_ptr = static_cast<ElementCompute const*>(arguments.norm_constant);
}
return MoeGroupedGemmOperation3xBase<Operator>::update_fusion_args(fusion_args, arguments);
}
};
public:
/// Returns success if the operation can proceed
Status can_implement([[maybe_unused]] void const* configuration_ptr, void const* arguments_ptr)
const override {
GroupedGemmBlockScaledArguments const* arguments =
static_cast<GroupedGemmBlockScaledArguments const*>(arguments_ptr);
OperatorArguments args;
auto status = update_arguments_(args, arguments);
if (status != Status::kSuccess) {
return status;
}
status = Operator::can_implement(args);
return status;
}
Status update_arguments_(
OperatorArguments& operator_args,
GroupedGemmBlockScaledArguments const* arguments) const {
Status status = UpdateFusionArgs<decltype(operator_args.epilogue.thread)>::update_(
operator_args.epilogue.thread,
*arguments);
if (status != Status::kSuccess) {
return status;
}
operator_args.mainloop.ptr_A = static_cast<typename Base::ArrayElementA const*>(arguments->ptr_A);
operator_args.mainloop.ptr_B = static_cast<typename Base::ArrayElementB const*>(arguments->ptr_B);
operator_args.mainloop.ptr_SFA = static_cast<const ElementSF*>(arguments->SFA);
operator_args.mainloop.ptr_SFB = static_cast<const ElementSF*>(arguments->SFB);
return this->update_arguments_base(operator_args, *arguments);
}
uint64_t get_device_workspace_size(void const* configuration_ptr, void const* arguments_ptr)
const override {
OperatorArguments args;
auto status =
update_arguments_(args, static_cast<GroupedGemmBlockScaledArguments const*>(arguments_ptr));
if (status != Status::kSuccess) {
return 0;
}
uint64_t size = Operator::get_workspace_size(args);
return size;
}
/// Initializes the workspace
/// **** CAUTION ****
/// Must be called when ldc, or ldd change.
/// The CUTLASS library stores the operations in a type-
/// erased manifest. Therefore, only this class knows
/// the type of strideA, strideB, strideC, and strideD.
/// Since grouped GEMM needs to allocate storage for
/// the strides on device, the concrete type of the stride
/// must be known in order to copy in the correct memory
/// layout on device.
Status initialize(
void const* configuration_ptr,
void* host_workspace,
void* device_workspace,
cudaStream_t stream = nullptr) const override {
Operator* op = new (host_workspace) Operator;
return Status::kSuccess;
}
/// **** CAUTION ****
/// initialize() must be called if ldc, or ldd change.
Status run(
void const* arguments_ptr,
void* host_workspace,
void* device_workspace = nullptr,
cudaStream_t stream = nullptr) const override {
OperatorArguments operator_args;
auto const& args = *static_cast<GroupedGemmBlockScaledArguments const*>(arguments_ptr);
Status status = update_arguments_(operator_args, &args);
if (status != Status::kSuccess) {
return status;
}
Operator* op = static_cast<Operator*>(host_workspace);
status = op->run(operator_args, device_workspace, stream, nullptr);
return status;
}
};
} // namespace cutlass::library
tools/library/src/handle.cu
+3
View File
@@ -279,6 +279,9 @@ static int gemm_problem_alignment(
int max_element_alignment = 0;
for (NumericTypeID type_id : elements) {
if (library::sizeof_bits(type_id)==0) {
continue;
}
int element_alignment = max_alignment_in_bytes * 8 / library::sizeof_bits(type_id);
max_element_alignment = std::max(max_element_alignment, element_alignment);
}
tools/library/src/reference/block_scaled_gemm_fp4a_vs16.cu
+5
View File
@@ -74,6 +74,11 @@ void initialize_block_scaled_gemm_reference_operations_fp4a_vs16(Manifest &manif
16 /*EpilogueSFVecSize*/
>(manifest);
make_block_scaled_gemm_tn<
float_e2m1_t /*A*/, float_ue4m3_t /*SFA*/, float_e2m1_t /*B*/, float_ue4m3_t /*SFB*/,
void /*C*/, float /*Compute*/, void /*SFD*/, float /*Accum*/, bfloat16_t /*D*/, 16 /*SFVecSize*/
>(manifest);
make_block_scaled_gemm_tn<
float_e2m1_t /*A*/, float_ue4m3_t /*SFA*/, float_e2m1_t /*B*/, float_ue4m3_t /*SFB*/,
half_t /*C*/, float /*Compute*/, float_ue8m0_t /*SFD*/, float /*Accum*/, float_e2m1_t /*D*/, 16 /*SFVecSize*/,
tools/library/src/reference/block_scaled_gemm_fp4a_vs32.cu
+9
View File
@@ -73,6 +73,15 @@ void initialize_block_scaled_gemm_reference_operations_fp4a_vs32(Manifest &manif
half_t /*C*/, float /*Compute*/, void /*SFD*/, float /*Accum*/, float_e3m2_t /*D*/, 32 /*SFVecSize*/
>(manifest);
make_block_scaled_gemm<
float_e2m1_t /*A*/, float_ue8m0_t /*SFA*/, float_e2m1_t /*B*/, float_ue8m0_t /*SFB*/,
bfloat16_t /*C*/, float /*Compute*/, void /*SFD*/, float /*Accum*/, bfloat16_t /*D*/, 32 /*SFVecSize*/
>(manifest);
make_block_scaled_gemm<
float_e2m1_t /*A*/, float_ue8m0_t /*SFA*/, float_e2m1_t /*B*/, float_ue8m0_t /*SFB*/,
void /*C*/, float /*Compute*/, void /*SFD*/, float /*Accum*/, bfloat16_t /*D*/, 32 /*SFVecSize*/
>(manifest);
// With SF generation reference
make_block_scaled_gemm<
float_e2m1_t /*A*/, float_ue8m0_t /*SFA*/, float_e2m1_t /*B*/, float_ue8m0_t /*SFB*/,
tools/library/src/reference/gemm_e4m3a_e4m3out.cu
+18
View File
@@ -56,6 +56,15 @@ void initialize_gemm_reference_operations_e4m3a_e4m3out(Manifest &manifest) {
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
@@ -65,6 +74,15 @@ void initialize_gemm_reference_operations_e4m3a_e4m3out(Manifest &manifest) {
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
tools/library/src/reference/gemm_e4m3a_e5m2out.cu
+18
View File
@@ -56,6 +56,15 @@ void initialize_gemm_reference_operations_e4m3a_e5m2out(Manifest &manifest) {
float_e5m2_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
@@ -65,6 +74,15 @@ void initialize_gemm_reference_operations_e4m3a_e5m2out(Manifest &manifest) {
float_e5m2_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
tools/library/src/reference/gemm_e5m2a_e4m3out.cu
+18
View File
@@ -56,6 +56,15 @@ void initialize_gemm_reference_operations_e5m2a_e4m3out(Manifest &manifest) {
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
@@ -65,6 +74,15 @@ void initialize_gemm_reference_operations_e5m2a_e4m3out(Manifest &manifest) {
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
tools/library/src/reference/gemm_e5m2a_e5m2out.cu
+9
View File
@@ -56,6 +56,15 @@ void initialize_gemm_reference_operations_e5m2a_e5m2out(Manifest &manifest) {
float_e5m2_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
tools/library/src/reference/gemm_f4_f4_f32.cu
+39
View File
@@ -97,7 +97,46 @@ void initialize_gemm_reference_operations_f4_f4_f32(Manifest &manifest) {
float, // ElementAccumulator
float // ElementD
>(manifest);
// 1.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e2m1_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
// 2.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e2m1_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
// 3.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e2m1_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
// 4.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e2m1_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_f4_f6_f32.cu
+39
View File
@@ -99,6 +99,45 @@ void initialize_gemm_reference_operations_f4_f6_f32(Manifest &manifest) {
float // ElementD
>(manifest);
// 1.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e3m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
// 2.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e3m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
// 3.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e3m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
// 4.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e3m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_f4_f8_f32.cu
+79
View File
@@ -99,6 +99,46 @@ void initialize_gemm_reference_operations_f4_f8_f32(Manifest &manifest) {
float // ElementD
>(manifest);
// 1.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
// 2.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
// 3.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
// 4.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
// 1.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
@@ -139,6 +179,45 @@ void initialize_gemm_reference_operations_f4_f8_f32(Manifest &manifest) {
float // ElementD
>(manifest);
// 1.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e5m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
// 2.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e5m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
// 3.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e5m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
// 4.
make_gemm_real_canonical_layouts<
float_e2m1_t, // ElementA
float_e5m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_f6_f4_f32.cu
+41
View File
@@ -99,6 +99,47 @@ void initialize_gemm_reference_operations_f6_f4_f32(Manifest &manifest) {
float // ElementD
>(manifest);
// 1.
make_gemm_real_canonical_layouts<
float_e3m2_t, // ElementA
float_e2m1_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
// 2.
make_gemm_real_canonical_layouts<
float_e3m2_t, // ElementA
float_e2m1_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
// 3.
make_gemm_real_canonical_layouts<
float_e3m2_t, // ElementA
float_e2m1_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
// 4.
make_gemm_real_canonical_layouts<
float_e3m2_t, // ElementA
float_e2m1_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_f6_f6_f32.cu
+40
View File
@@ -98,6 +98,46 @@ void initialize_gemm_reference_operations_f6_f6_f32(Manifest &manifest) {
float // ElementD
>(manifest);
// 1.
make_gemm_real_canonical_layouts<
float_e3m2_t, // ElementA
float_e3m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
// 2.
make_gemm_real_canonical_layouts<
float_e3m2_t, // ElementA
float_e3m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
// 3.
make_gemm_real_canonical_layouts<
float_e3m2_t, // ElementA
float_e3m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
// 4.
make_gemm_real_canonical_layouts<
float_e3m2_t, // ElementA
float_e3m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_f6_f8_f32.cu
+39
View File
@@ -99,6 +99,45 @@ void initialize_gemm_reference_operations_f6_f8_f32(Manifest &manifest) {
float // ElementD
>(manifest);
// 1.
make_gemm_real_canonical_layouts<
float_e3m2_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
// 2.
make_gemm_real_canonical_layouts<
float_e3m2_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
// 3.
make_gemm_real_canonical_layouts<
float_e3m2_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
// 4.
make_gemm_real_canonical_layouts<
float_e3m2_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_f8_f4_f32.cu
+40
View File
@@ -99,6 +99,46 @@ void initialize_gemm_reference_operations_f8_f4_f32(Manifest &manifest) {
float // ElementD
>(manifest);
// 1.
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e2m1_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
// 2.
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e2m1_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
// 3.
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e2m1_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
// 4.
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e2m1_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_f8_f6_f32.cu
+40
View File
@@ -99,6 +99,46 @@ void initialize_gemm_reference_operations_f8_f6_f32(Manifest &manifest) {
float // ElementD
>(manifest);
// 1.
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e3m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e4m3_t // ElementD
>(manifest);
// 2.
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e3m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float_e5m2_t // ElementD
>(manifest);
// 3.
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e3m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
// 4.
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e3m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_fp32out.cu
+61
View File
@@ -54,6 +54,15 @@ void initialize_gemm_reference_operations_fp32out(Manifest &manifest) {
float // ElementAccumulator
>(manifest);
make_gemm_real_canonical_layouts<
float, // ElementA
float, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float
>(manifest);
make_gemm_real_canonical_layouts<
tfloat32_t,
tfloat32_t,
@@ -62,6 +71,15 @@ void initialize_gemm_reference_operations_fp32out(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
tfloat32_t,
tfloat32_t,
void,
float,
float,
float
>(manifest);
make_gemm_real_canonical_layouts<
tfloat32_t,
tfloat32_t,
@@ -69,6 +87,14 @@ void initialize_gemm_reference_operations_fp32out(Manifest &manifest) {
float,
float
>(manifest);
make_gemm_real_canonical_layouts<
tfloat32_t,
tfloat32_t,
void,
float,
float,
tfloat32_t
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
@@ -78,6 +104,15 @@ void initialize_gemm_reference_operations_fp32out(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
half_t,
void,
float,
float,
half_t
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
half_t,
@@ -86,6 +121,14 @@ void initialize_gemm_reference_operations_fp32out(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
half_t,
void,
float,
float,
float
>(manifest);
make_gemm_real_canonical_layouts<
bfloat16_t,
bfloat16_t,
@@ -94,6 +137,15 @@ void initialize_gemm_reference_operations_fp32out(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
bfloat16_t,
bfloat16_t,
void,
float,
float,
float
>(manifest);
make_gemm_real_canonical_layouts<
bfloat16_t,
bfloat16_t,
@@ -101,6 +153,15 @@ void initialize_gemm_reference_operations_fp32out(Manifest &manifest) {
float,
float
>(manifest);
make_gemm_real_canonical_layouts<
bfloat16_t,
bfloat16_t,
void,
float,
float,
bfloat16_t
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_fp8in_bf16out.cu
+36
View File
@@ -82,6 +82,42 @@ void initialize_gemm_reference_operations_fp8in_bf16out(Manifest &manifest) {
float, // ElementAccumulator
bfloat16_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
bfloat16_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
bfloat16_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
bfloat16_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
bfloat16_t // ElementD
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_fp8in_fp16out.cu
+36
View File
@@ -82,6 +82,42 @@ void initialize_gemm_reference_operations_fp8in_fp16out(Manifest &manifest) {
float, // ElementAccumulator
half_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
half_t // ElementD
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_fp8in_fp32out.cu
+36
View File
@@ -82,6 +82,42 @@ void initialize_gemm_reference_operations_fp8in_fp32out(Manifest &manifest) {
float, // ElementAccumulator
float // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e4m3_t, // ElementA
float_e5m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e4m3_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
make_gemm_real_canonical_layouts<
float_e5m2_t, // ElementA
float_e5m2_t, // ElementB
void, // ElementC
float, // ElementScalar
float, // ElementAccumulator
float // ElementD
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_fp_mixed_input.cu
+183
View File
@@ -54,6 +54,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
int8_t,
half_t,
void,
float,
float,
float
>(manifest);
make_gemm_real_canonical_layouts<
uint8_t,
half_t,
@@ -61,6 +70,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
uint8_t,
half_t,
void,
float,
float,
float
>(manifest);
make_gemm_real_canonical_layouts<
int8_t,
half_t,
@@ -68,6 +86,16 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
int8_t,
half_t,
void,
float,
float,
half_t
>(manifest);
make_gemm_real_canonical_layouts<
uint8_t,
half_t,
@@ -75,6 +103,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
uint8_t,
half_t,
void,
float,
float,
half_t
>(manifest);
make_gemm_real_canonical_layouts<
int8_t,
half_t,
@@ -82,6 +119,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
half_t
>(manifest);
make_gemm_real_canonical_layouts<
int8_t,
half_t,
void,
half_t,
half_t,
half_t
>(manifest);
make_gemm_real_canonical_layouts<
uint8_t,
half_t,
@@ -89,6 +135,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
half_t
>(manifest);
make_gemm_real_canonical_layouts<
uint8_t,
half_t,
void,
half_t,
half_t,
half_t
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
int8_t,
@@ -96,6 +151,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
int8_t,
void,
float,
float,
float
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
uint8_t,
@@ -103,6 +167,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
uint8_t,
void,
float,
float,
float
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
int8_t,
@@ -110,6 +183,16 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
half_t
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
int8_t,
void,
half_t,
half_t,
half_t
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
uint8_t,
@@ -117,6 +200,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
half_t
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
uint8_t,
void,
half_t,
half_t,
half_t
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
int8_t,
@@ -124,6 +216,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
int8_t,
void,
float,
float,
half_t
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
uint8_t,
@@ -131,6 +232,16 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
uint8_t,
void,
float,
float,
half_t
>(manifest);
// bfloat16_t mixed with 8-bit integer input
make_gemm_real_canonical_layouts<
int8_t,
@@ -139,6 +250,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
int8_t,
bfloat16_t,
void,
float,
float,
float
>(manifest);
make_gemm_real_canonical_layouts<
uint8_t,
bfloat16_t,
@@ -146,6 +266,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
uint8_t,
bfloat16_t,
void,
float,
float,
float
>(manifest);
make_gemm_real_canonical_layouts<
int8_t,
bfloat16_t,
@@ -153,6 +282,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
int8_t,
bfloat16_t,
void,
float,
float,
bfloat16_t
>(manifest);
make_gemm_real_canonical_layouts<
uint8_t,
bfloat16_t,
@@ -160,6 +298,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
uint8_t,
bfloat16_t,
void,
float,
float,
bfloat16_t
>(manifest);
make_gemm_real_canonical_layouts<
bfloat16_t,
int8_t,
@@ -167,6 +314,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
bfloat16_t,
int8_t,
void,
float,
float,
float
>(manifest);
make_gemm_real_canonical_layouts<
bfloat16_t,
uint8_t,
@@ -174,6 +330,15 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
bfloat16_t,
uint8_t,
void,
float,
float,
float
>(manifest);
make_gemm_real_canonical_layouts<
bfloat16_t,
int8_t,
@@ -181,12 +346,30 @@ void initialize_gemm_reference_operations_fp_mixed_input(Manifest &manifest) {
float
>(manifest);
make_gemm_real_canonical_layouts<
bfloat16_t,
int8_t,
void,
float,
float,
bfloat16_t
>(manifest);
make_gemm_real_canonical_layouts<
bfloat16_t,
uint8_t,
bfloat16_t,
float
>(manifest);
make_gemm_real_canonical_layouts<
bfloat16_t,
uint8_t,
void,
float,
float,
bfloat16_t
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_fp_other.cu
+54
View File
@@ -54,6 +54,15 @@ void initialize_gemm_reference_operations_fp_other(Manifest &manifest) {
half_t
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
half_t,
void,
half_t,
half_t,
half_t
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
half_t,
@@ -62,6 +71,24 @@ void initialize_gemm_reference_operations_fp_other(Manifest &manifest) {
half_t
>(manifest);
make_gemm_real_canonical_layouts<
half_t,
half_t,
float,
float,
float,
half_t
>(manifest);
make_gemm_real_canonical_layouts<
bfloat16_t,
bfloat16_t,
float,
float,
float,
bfloat16_t
>(manifest);
make_gemm_real_canonical_layouts<
double,
double,
@@ -70,6 +97,15 @@ void initialize_gemm_reference_operations_fp_other(Manifest &manifest) {
double
>(manifest);
make_gemm_real_canonical_layouts<
double,
double,
void,
double,
double,
double
>(manifest);
make_gemm_complex_canonical_layouts<
complex<float>,
complex<float>,
@@ -78,6 +114,15 @@ void initialize_gemm_reference_operations_fp_other(Manifest &manifest) {
complex<float>
>(manifest);
make_gemm_complex_canonical_layouts<
complex<float>,
complex<float>,
void,
complex<float>,
complex<float>,
complex<float>
>(manifest);
make_gemm_complex_canonical_layouts<
complex<double>,
complex<double>,
@@ -85,6 +130,15 @@ void initialize_gemm_reference_operations_fp_other(Manifest &manifest) {
complex<double>,
complex<double>
>(manifest);
make_gemm_complex_canonical_layouts<
complex<double>,
complex<double>,
void,
complex<double>,
complex<double>,
complex<double>
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_int4.cu
+75
View File
@@ -59,7 +59,28 @@ void initialize_gemm_reference_operations_int4(Manifest &manifest) {
64,
int4b_t,
int4b_t,
void,
int32_t,
int32_t,
int32_t
>(manifest);
make_gemm_interleaved_layouts<
64,
int4b_t,
int4b_t,
int32_t,
float,
int32_t,
int32_t,
NumericConverterClamp<int32_t, float>
>(manifest);
make_gemm_interleaved_layouts<
64,
int4b_t,
int4b_t,
void,
float,
int32_t,
int32_t,
@@ -77,6 +98,17 @@ void initialize_gemm_reference_operations_int4(Manifest &manifest) {
NumericConverterClamp<int4b_t, float>
>(manifest);
make_gemm_interleaved_layouts<
64,
int4b_t,
int4b_t,
void,
float,
int32_t,
int4b_t,
NumericConverterClamp<int4b_t, float>
>(manifest);
make_gemm_interleaved_layouts<
64,
uint4b_t,
@@ -90,7 +122,28 @@ void initialize_gemm_reference_operations_int4(Manifest &manifest) {
64,
uint4b_t,
uint4b_t,
void,
int32_t,
int32_t,
int32_t
>(manifest);
make_gemm_interleaved_layouts<
64,
uint4b_t,
uint4b_t,
int32_t,
float,
int32_t,
int32_t,
NumericConverterClamp<int32_t, float>
>(manifest);
make_gemm_interleaved_layouts<
64,
uint4b_t,
uint4b_t,
void,
float,
int32_t,
int32_t,
@@ -108,6 +161,17 @@ void initialize_gemm_reference_operations_int4(Manifest &manifest) {
NumericConverterClamp<uint4b_t, float>
>(manifest);
make_gemm_interleaved_layouts<
64,
uint4b_t,
uint4b_t,
void,
float,
int32_t,
uint4b_t,
NumericConverterClamp<uint4b_t, float>
>(manifest);
make_gemm_interleaved_layouts<
64,
uint4b_t,
@@ -118,6 +182,17 @@ void initialize_gemm_reference_operations_int4(Manifest &manifest) {
int4b_t,
NumericConverterClamp<int4b_t, float>
>(manifest);
make_gemm_interleaved_layouts<
64,
uint4b_t,
uint4b_t,
void,
float,
int32_t,
int4b_t,
NumericConverterClamp<int4b_t, float>
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_int8_interleaved_32.cu
+75
View File
@@ -59,7 +59,28 @@ void initialize_gemm_reference_operations_int8_interleaved_32(Manifest &manifest
32,
int8_t,
int8_t,
void,
int32_t,
int32_t,
int32_t
>(manifest);
make_gemm_interleaved_layouts<
32,
int8_t,
int8_t,
int32_t,
float,
int32_t,
int32_t,
NumericConverterClamp<int32_t, float>
>(manifest);
make_gemm_interleaved_layouts<
32,
int8_t,
int8_t,
void,
float,
int32_t,
int32_t,
@@ -77,6 +98,17 @@ void initialize_gemm_reference_operations_int8_interleaved_32(Manifest &manifest
NumericConverterClamp<int8_t, float>
>(manifest);
make_gemm_interleaved_layouts<
32,
int8_t,
int8_t,
void,
float,
int32_t,
int8_t,
NumericConverterClamp<int8_t, float>
>(manifest);
make_gemm_interleaved_layouts<
32,
uint8_t,
@@ -90,7 +122,28 @@ void initialize_gemm_reference_operations_int8_interleaved_32(Manifest &manifest
32,
uint8_t,
uint8_t,
void,
int32_t,
int32_t,
int32_t
>(manifest);
make_gemm_interleaved_layouts<
32,
uint8_t,
uint8_t,
int32_t,
float,
int32_t,
int32_t,
NumericConverterClamp<int32_t, float>
>(manifest);
make_gemm_interleaved_layouts<
32,
uint8_t,
uint8_t,
void,
float,
int32_t,
int32_t,
@@ -108,6 +161,17 @@ void initialize_gemm_reference_operations_int8_interleaved_32(Manifest &manifest
NumericConverterClamp<uint8_t, float>
>(manifest);
make_gemm_interleaved_layouts<
32,
uint8_t,
uint8_t,
void,
float,
int32_t,
uint8_t,
NumericConverterClamp<uint8_t, float>
>(manifest);
make_gemm_interleaved_layouts<
32,
uint8_t,
@@ -118,6 +182,17 @@ void initialize_gemm_reference_operations_int8_interleaved_32(Manifest &manifest
int8_t,
NumericConverterClamp<int8_t, float>
>(manifest);
make_gemm_interleaved_layouts<
32,
uint8_t,
uint8_t,
void,
float,
int32_t,
int8_t,
NumericConverterClamp<int8_t, float>
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_int8_interleaved_64.cu
+76
View File
@@ -59,7 +59,28 @@ void initialize_gemm_reference_operations_int8_interleaved_64(Manifest &manifest
64,
int4b_t,
int4b_t,
void,
int32_t,
int32_t,
int32_t
>(manifest);
make_gemm_interleaved_layouts<
64,
int4b_t,
int4b_t,
int32_t,
float,
int32_t,
int32_t,
NumericConverterClamp<int32_t, float>
>(manifest);
make_gemm_interleaved_layouts<
64,
int4b_t,
int4b_t,
void,
float,
int32_t,
int32_t,
@@ -77,6 +98,18 @@ void initialize_gemm_reference_operations_int8_interleaved_64(Manifest &manifest
NumericConverterClamp<int4b_t, float>
>(manifest);
make_gemm_interleaved_layouts<
64,
int4b_t,
int4b_t,
void,
float,
int32_t,
int4b_t,
NumericConverterClamp<int4b_t, float>
>(manifest);
make_gemm_interleaved_layouts<
64,
uint4b_t,
@@ -90,7 +123,28 @@ void initialize_gemm_reference_operations_int8_interleaved_64(Manifest &manifest
64,
uint4b_t,
uint4b_t,
void,
int32_t,
int32_t,
int32_t
>(manifest);
make_gemm_interleaved_layouts<
64,
uint4b_t,
uint4b_t,
int32_t,
float,
int32_t,
int32_t,
NumericConverterClamp<int32_t, float>
>(manifest);
make_gemm_interleaved_layouts<
64,
uint4b_t,
uint4b_t,
void,
float,
int32_t,
int32_t,
@@ -108,6 +162,17 @@ void initialize_gemm_reference_operations_int8_interleaved_64(Manifest &manifest
NumericConverterClamp<uint4b_t, float>
>(manifest);
make_gemm_interleaved_layouts<
64,
uint4b_t,
uint4b_t,
void,
float,
int32_t,
uint4b_t,
NumericConverterClamp<uint4b_t, float>
>(manifest);
make_gemm_interleaved_layouts<
64,
uint4b_t,
@@ -118,6 +183,17 @@ void initialize_gemm_reference_operations_int8_interleaved_64(Manifest &manifest
int4b_t,
NumericConverterClamp<int4b_t, float>
>(manifest);
make_gemm_interleaved_layouts<
64,
uint4b_t,
uint4b_t,
void,
float,
int32_t,
int4b_t,
NumericConverterClamp<int4b_t, float>
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_int_mixed_input.cu
+78
View File
@@ -57,7 +57,26 @@ void initialize_gemm_reference_operations_int_mixed_input(Manifest &manifest) {
make_gemm_real_canonical_layouts<
int4b_t,
int8_t,
void,
int32_t,
int32_t,
int32_t
>(manifest);
make_gemm_real_canonical_layouts<
int4b_t,
int8_t,
int8_t,
int32_t,
int32_t,
int8_t,
NumericConverterClamp<int8_t, int32_t>
>(manifest);
make_gemm_real_canonical_layouts<
int4b_t,
int8_t,
void,
int32_t,
int32_t,
int8_t,
@@ -77,7 +96,27 @@ void initialize_gemm_reference_operations_int_mixed_input(Manifest &manifest) {
make_gemm_real_canonical_layouts<
int4b_t,
int8_t,
void,
float,
int32_t,
int32_t,
NumericConverterClamp<int32_t, float>
>(manifest);
make_gemm_real_canonical_layouts<
int4b_t,
int8_t,
int8_t,
float,
int32_t,
int8_t,
NumericConverterClamp<int8_t, float>
>(manifest);
make_gemm_real_canonical_layouts<
int4b_t,
int8_t,
void,
float,
int32_t,
int8_t,
@@ -91,6 +130,15 @@ void initialize_gemm_reference_operations_int_mixed_input(Manifest &manifest) {
int32_t
>(manifest);
make_gemm_real_canonical_layouts<
int8_t,
int4b_t,
void,
int32_t,
int32_t,
int32_t
>(manifest);
make_gemm_real_canonical_layouts<
int8_t,
int4b_t,
@@ -104,7 +152,27 @@ void initialize_gemm_reference_operations_int_mixed_input(Manifest &manifest) {
make_gemm_real_canonical_layouts<
int8_t,
int4b_t,
void,
int32_t,
int32_t,
int8_t,
NumericConverterClamp<int8_t, int32_t>
>(manifest);
make_gemm_real_canonical_layouts<
int8_t,
int4b_t,
int32_t,
float,
int32_t,
int32_t,
NumericConverterClamp<int32_t, float>
>(manifest);
make_gemm_real_canonical_layouts<
int8_t,
int4b_t,
void,
float,
int32_t,
int32_t,
@@ -120,6 +188,16 @@ void initialize_gemm_reference_operations_int_mixed_input(Manifest &manifest) {
int8_t,
NumericConverterClamp<int8_t, float>
>(manifest);
make_gemm_real_canonical_layouts<
int8_t,
int4b_t,
void,
float,
int32_t,
int8_t,
NumericConverterClamp<int8_t, float>
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_reference_operation.h
+6 -5
View File
@@ -86,7 +86,8 @@ public:
using ElementC = ElementC_;
using LayoutC = LayoutC_;
using ElementD = ElementD_;
using TensorRefC = TensorRef<ElementC, LayoutC>;
using NonVoidElementC = cute::conditional_t<cute::is_void_v<ElementC>, ElementD, ElementC>;
using TensorRefC = TensorRef<NonVoidElementC, LayoutC>;
using TensorRefD = TensorRef<ElementD, LayoutC>;
using ElementCompute = ElementCompute_;
using ElementAccumulator = ElementAccumulator_;
@@ -199,7 +200,7 @@ public:
TensorRefA ref_A{static_cast<ElementA *>(const_cast<void *>(args.A)), LayoutA(int(config.lda))};
TensorRefB ref_B{static_cast<ElementB *>(const_cast<void *>(args.B)), LayoutB(int(config.ldb))};
TensorRefC ref_C{static_cast<ElementC *>(const_cast<void *>(args.C)), LayoutC(int(config.ldc))};
TensorRefC ref_C{static_cast<NonVoidElementC *>(const_cast<void *>(args.C)), LayoutC(int(config.ldc))};
TensorRefD ref_D{static_cast<ElementD *>(args.D), LayoutC(int(config.ldd))};
if (kProvider == Provider::kReferenceHost) {
@@ -209,7 +210,7 @@ public:
LayoutA,
ElementB,
LayoutB,
ElementC,
NonVoidElementC,
LayoutC,
ElementCompute,
ElementAccumulator,
@@ -243,7 +244,7 @@ public:
LayoutA,
ElementB,
LayoutB,
ElementC,
NonVoidElementC,
LayoutC,
ElementCompute,
ElementAccumulator,
@@ -434,7 +435,7 @@ template <
typename ElementCompute_,
typename ElementAccumulator_ = ElementCompute_,
typename ElementD_ = ElementC_,
typename ConvertOp_ = NumericConverter<ElementC_, ElementCompute_>,
typename ConvertOp_ = NumericConverter<ElementD_, ElementCompute_>,
typename InnerProductOp_ = multiply_add<ElementAccumulator_>
>
void make_gemm_interleaved_layouts(Manifest &manifest) {
tools/library/src/reference/gemm_s8_s8_s32.cu
+74
View File
@@ -135,6 +135,80 @@ void initialize_gemm_reference_operations_s8_s8_s32(Manifest &manifest) {
half_t, // ElementD
NumericConverterClamp<half_t, float> // From Scalar to D
>(manifest);
// 1.
make_gemm_real_canonical_layouts<
int8_t, // ElementA
int8_t, // ElementB
void, // ElementC
int32_t, // ElementScalar / ElementCompute
int32_t, // ElementAccumulator
int32_t // ElementD
>(manifest);
// 2.
make_gemm_real_canonical_layouts<
int8_t, // ElementA
int8_t, // ElementB
void, // ElementC
float, // ElementScalar / ElementCompute
int32_t, // ElementAccumulator
int32_t // ElementD
>(manifest);
// 3.
make_gemm_real_canonical_layouts<
int8_t, // ElementA
int8_t, // ElementB
void, // ElementC
float, // ElementScalar / ElementCompute
int32_t, // ElementAccumulator
int8_t, // ElementD
NumericConverterClamp<int8_t, float> // From Scalar to D
>(manifest);
// 4.
make_gemm_real_canonical_layouts<
int8_t, // ElementA
int8_t, // ElementB
void, // ElementC
int32_t, // ElementScalar / ElementCompute
int32_t, // ElementAccumulator
int8_t, // ElementD
NumericConverterClamp<int8_t, int32_t> // From Scalar to D
>(manifest);
// 5.
make_gemm_real_canonical_layouts<
int8_t, // ElementA
int8_t, // ElementB
void, // ElementC
float, // ElementScalar / ElementCompute
int32_t, // ElementAccumulator
int8_t, // ElementD
NumericConverterClamp<int8_t, float> // From Scalar to D
>(manifest);
// 6.
make_gemm_real_canonical_layouts<
int8_t, // ElementA
int8_t, // ElementB
void, // ElementC
float, // ElementScalar / ElementCompute
int32_t, // ElementAccumulator
float // ElementD
>(manifest);
// 7.
make_gemm_real_canonical_layouts<
int8_t, // ElementA
int8_t, // ElementB
void, // ElementC
float, // ElementScalar / ElementCompute
int32_t, // ElementAccumulator
half_t, // ElementD
NumericConverterClamp<half_t, float> // From Scalar to D
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/library/src/reference/gemm_u8_u8_s32.cu
+43
View File
@@ -98,6 +98,49 @@ void initialize_gemm_reference_operations_u8_u8_s32(Manifest &manifest) {
NumericConverterClamp<uint8_t, float> // From Scalar to D
>(manifest);
// 1.
make_gemm_real_canonical_layouts<
uint8_t, // ElementA
uint8_t, // ElementB
void, // ElementC
int32_t, // ElementScalar / ElementCompute
int32_t, // ElementAccumulator
int32_t // ElementD
>(manifest);
// 2.
make_gemm_real_canonical_layouts<
uint8_t, // ElementA
uint8_t, // ElementB
void, // ElementC
float, // ElementScalar / ElementCompute
int32_t, // ElementAccumulator
int32_t, // ElementD
NumericConverterClamp<int32_t, float> // From Scalar to D
>(manifest);
// 3.
make_gemm_real_canonical_layouts<
uint8_t, // ElementA
uint8_t, // ElementB
void, // ElementC
float, // ElementScalar / ElementCompute
int32_t, // ElementAccumulator
int8_t, // ElementD
NumericConverterClamp<int8_t, float> // From Scalar to D
>(manifest);
// 4.
make_gemm_real_canonical_layouts<
uint8_t, // ElementA
uint8_t, // ElementB
void, // ElementC
float, // ElementScalar / ElementCompute
int32_t, // ElementAccumulator
uint8_t, // ElementD
NumericConverterClamp<uint8_t, float> // From Scalar to D
>(manifest);
}
///////////////////////////////////////////////////////////////////////////////////////////////////
tools/profiler/include/cutlass/profiler/device_allocation.h
+12
View File
@@ -84,6 +84,8 @@ private:
/// The device ID where the allocation is made
int device_;
/// Whether to free the memory when the object is destroyed
bool free_memory_{true};
public:
//
// Static member functions
@@ -140,6 +142,16 @@ public:
int batch_count = 1,
int device = -1);
DeviceAllocation(
library::NumericTypeID type,
library::LayoutTypeID layout_id,
std::vector<int> const &extent,
std::vector<int64_t> const &stride,
void* ref_pointer_,
int batch_count,
int device
);
~DeviceAllocation();
DeviceAllocation &reset();
tools/profiler/include/cutlass/profiler/device_context.h
+12
View File
@@ -90,6 +90,18 @@ public:
int batch_count,
size_t device_index);
/// creates a reference tensor of existing memory, a given type, capacity (elements), and name
DeviceAllocation *create_ref_tensor(
Options const &options,
std::string const &name,
library::NumericTypeID type,
library::LayoutTypeID layout_id,
std::vector<int> const &extent,
std::vector<int64_t> const &stride,
void* ref_pointer_,
int batch_count,
size_t device_index);
/// Allocates memory of a given type, capacity (elements), and name
DeviceAllocation *allocate_and_initialize_tensor(
Options const &options,
tools/profiler/include/cutlass/profiler/grouped_gemm_operation_profiler.h
+17 -7
View File
@@ -68,6 +68,7 @@ public:
std::vector<gemm::GemmCoord> problem_sizes;
std::vector<cute::Shape<int, int, int>> problem_sizes_3x;
std::vector<int32_t> max_problem_size_3x;
/// For exploration purposes
std::vector<std::array<int64_t, 3>> preferred_clusters;
@@ -85,10 +86,13 @@ public:
std::vector<int64_t> lda{0};
std::vector<int64_t> ldb{0};
std::vector<int64_t> ldc{0};
int64_t max_lda{0};
int64_t max_ldb{0};
int64_t max_ldc{0};
std::vector<uint8_t> alpha;
std::vector<uint8_t> beta;
cutlass::library::RasterOrder raster_order{cutlass::library::RasterOrder::kHeuristic};
int swizzle_size{1};
@@ -168,7 +172,8 @@ public:
DeviceAllocation* ldb_array_device{nullptr};
DeviceAllocation* ldc_array_device{nullptr};
DeviceAllocation* ldd_array_device{nullptr};
std::vector<int32_t> tokens_per_expert_host;
DeviceAllocation* tokens_per_expert_device{nullptr};
std::optional<BlockScalingWorkspace> block_scales;
library::GemmGroupedConfiguration configuration;
@@ -188,6 +193,10 @@ private:
arguments.ptr_B = gemm_workspace_.B_ptr_array_device[0]->data();
arguments.ptr_C = gemm_workspace_.C_ptr_array_device[0]->data();
arguments.ptr_D = gemm_workspace_.D_ptr_array_device[0]->data();
if (is_moe) {
arguments.tokens_per_expert_host = gemm_workspace_.tokens_per_expert_host.data();
arguments.tokens_per_expert = static_cast<int32_t*>(gemm_workspace_.tokens_per_expert_device->data());
}
arguments.alpha = problem_.alpha.data();
arguments.beta = problem_.beta.data();
@@ -200,10 +209,11 @@ private:
static_cast<gemm::GemmCoord*>(gemm_workspace_.problem_sizes_array_device->data());
arguments.problem_sizes_3x = static_cast<cute::Shape<int, int, int>*>(
gemm_workspace_.problem_sizes_3x_array_device->data());
gemm_workspace_.arguments.problem_sizes_3x_host = problem_.problem_sizes_3x.data();
gemm_workspace_.arguments.problem_count = problem_.problem_sizes.size();
gemm_workspace_.arguments.cluster_shape = {int(problem_.cluster_m), int(problem_.cluster_n), int(problem_.cluster_k)};
gemm_workspace_.arguments.cluster_shape_fallback = {int(problem_.cluster_m_fallback), int(problem_.cluster_n_fallback), int(problem_.cluster_k_fallback)};
arguments.problem_sizes_3x_host = problem_.problem_sizes_3x.data();
arguments.max_problem_size_3x = problem_.max_problem_size_3x;
arguments.problem_count = problem_.problem_sizes.size();
arguments.cluster_shape = {int(problem_.cluster_m), int(problem_.cluster_n), int(problem_.cluster_k)};
arguments.cluster_shape_fallback = {int(problem_.cluster_m_fallback), int(problem_.cluster_n_fallback), int(problem_.cluster_k_fallback)};
/* Query device SM count to pass onto the kernel as an argument, where needed */
arguments.sm_count = options.device.get_sm_count(0);
@@ -230,7 +240,7 @@ protected:
bool is_block_scaled{false};
bool is_blockwise{false};
bool is_moe{false};
public:
GroupedGemmOperationProfiler(Options const& options);
tools/profiler/src/device_allocation.cu
+33 -2
View File
@@ -326,8 +326,38 @@ DeviceAllocation::DeviceAllocation(
reset(type, layout_id, extent, stride, batch_count);
}
DeviceAllocation::DeviceAllocation(
library::NumericTypeID type,
library::LayoutTypeID layout_id,
std::vector<int> const &extent,
std::vector<int64_t> const &stride,
void* ref_pointer_,
int batch_count,
int device
):
type_(type), batch_stride_(size_t(0)), capacity_(size_t(0)),
pointer_(ref_pointer_), batch_count_(1), device_(device), free_memory_(false) {
tensor_ref_buffer_.resize(sizeof(pointer_) + (sizeof(int64_t) * library::get_layout_stride_rank(layout_id)), 0);
type_ = type;
layout_ = layout_id;
stride_ = stride;
extent_ = extent;
batch_count_ = batch_count;
batch_stride_ = construct_layout(
tensor_ref_buffer_.data() + sizeof(pointer_),
layout_id,
extent,
stride_);
capacity_ = batch_stride_ * batch_count_;
}
DeviceAllocation::~DeviceAllocation() {
if (pointer_) {
if (pointer_ and free_memory_) {
int current_device;
cudaGetDevice(&current_device);
@@ -343,7 +373,7 @@ DeviceAllocation::~DeviceAllocation() {
}
DeviceAllocation &DeviceAllocation::reset() {
if (pointer_) {
if (pointer_ and free_memory_) {
int current_device;
cudaGetDevice(&current_device);
@@ -366,6 +396,7 @@ DeviceAllocation &DeviceAllocation::reset() {
extent_.clear();
tensor_ref_buffer_.clear();
batch_count_ = 1;
free_memory_ = true;
return *this;
}
tools/profiler/src/device_context.cu
+21
View File
@@ -75,6 +75,27 @@ DeviceAllocation *DeviceContext::allocate_tensor(
return allocation;
}
/// creates a reference tensor of existing ptr, a given type, capacity (elements), and name
DeviceAllocation *DeviceContext::create_ref_tensor(
Options const &options,
std::string const &name,
library::NumericTypeID type,
library::LayoutTypeID layout_id,
std::vector<int> const &extent,
std::vector<int64_t> const &stride,
void* ref_pointer_,
int batch_count,
size_t device_index) {
int device = options.device.device_id(device_index);
device_memory_.emplace_back(type, layout_id, extent, stride, ref_pointer_, batch_count,
device);
DeviceAllocation *allocation = &device_memory_.back();
allocations_[name] = allocation;
return allocation;
}
static void initialize_allocation_with_data_distribution(
Options const &options,
int seed_shift,
tools/profiler/src/gemm_operation_profiler.cu
+1 -1
View File
@@ -1264,7 +1264,7 @@ bool GemmOperationProfiler::verify_cutlass(
}
}
// if verification.required is set, then return success iff at least one ref-check was run
// if verification.required is set, then return success if at least one ref-check was run
if (options.verification.required) {
bool did_any_verification_run = false;
for (auto provider : options.verification.providers) {
tools/profiler/src/grouped_gemm_operation_profiler.cu
+515 -165
View File
@@ -177,7 +177,7 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse(
library::GroupedGemmDescription const& operation_desc,
ProblemSpace const& problem_space,
ProblemSpace::Problem const& problem) {
bool is_moe = operation_desc.is_moe;
this->mode = library::GemmUniversalMode::kGrouped;
std::bitset<3> args_exist;
@@ -189,7 +189,7 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse(
arg_as_int(k, "k", problem_space, problem);
std::string problem_file;
args_exist[2] = arg_as_string(problem_file, "problem-sizes-file", problem_space, problem);
int max_m = 0, max_n = 0, max_k = 0;
if (args_exist.count() == 0) {
int num_groups = 8;
problem_sizes.resize(num_groups);
@@ -204,6 +204,7 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse(
problem_sizes[i] = {m, n, k};
problem_sizes_3x[i] = {m, n, k};
}
max_problem_size_3x = {m0 * num_groups, n0 * num_groups, k0 * num_groups};
}
else if (args_exist.count() > 1) {
std::cerr
@@ -220,9 +221,13 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse(
auto m = problems[i][0];
auto n = problems[i][1];
auto k = problems[i][2];
max_m = std::max(max_m, m);
max_n = std::max(max_n, n);
max_k = std::max(max_k, k);
problem_sizes[i] = {m, n, k};
problem_sizes_3x[i] = {m, n, k};
}
max_problem_size_3x = {max_m, max_n, max_k};
}
// m, n, k path
else if (args_exist[1]) {
@@ -237,6 +242,7 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse(
problem_sizes[i] = {m, n, k};
problem_sizes_3x[i] = {m, n, k};
}
max_problem_size_3x = {m, n, k};
}
// --problem-sizes-file path
else if (args_exist[2]) {
@@ -247,7 +253,6 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse(
// clear the problem sizes and 3x problem sizes from previous operation
problem_sizes.clear();
problem_sizes_3x.clear();
for (std::string line; std::getline(file, line);) {
std::istringstream iss(line);
@@ -258,14 +263,27 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse(
if (iss >> m >> sep1 >> n >> sep2 >> k && sep1 == 'x' && sep2 == 'x' && !(iss >> remaining)) {
problem_sizes.emplace_back(m, n, k);
problem_sizes_3x.emplace_back(m, n, k);
max_m = std::max(max_m, m);
max_n = std::max(max_n, n);
max_k = std::max(max_k, k);
}
else {
throw std::runtime_error(
"Invalid format in line: " + line + ". Each line in file expected to be 'mxnxk'.");
}
}
max_problem_size_3x = {max_m, max_n, max_k};
}
if (is_moe) {
for(size_t group_idx = 0; group_idx < problem_sizes.size(); group_idx++) {
if (problem_sizes[group_idx].m() != max_problem_size_3x[0] ||
problem_sizes[group_idx].k() != max_problem_size_3x[2]) {
std::cerr << "Problem size M:"<< problem_sizes[group_idx].m() << "K:" << problem_sizes[group_idx].k() << " for group " << group_idx << "should be equal to "
<< "Max problem size M:" << max_problem_size_3x[0] << "K:" << max_problem_size_3x[2] << " in MoE Grouped GEMM" << std::endl;
return Status::kErrorInvalidProblem;
}
}
}
if (!arg_as_int(this->cluster_m, "cluster_m", problem_space, problem)) {
// default value
this->cluster_m = std::string(operation_desc.gemm.name).find("_2sm") != std::string::npos ? 2 : 1;
@@ -382,6 +400,9 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse(
operation_desc.gemm.C.layout,
{int(this->m(group_idx)), int(this->n(group_idx))})
.front();
this->max_lda = std::max(this->max_lda, this->lda[group_idx]);
this->max_ldb = std::max(this->max_ldb, this->ldb[group_idx]);
this->max_ldc = std::max(this->max_ldc, this->ldc[group_idx]);
}
// instantiation for exploration profiling
@@ -609,7 +630,7 @@ Status GroupedGemmOperationProfiler::initialize_configuration(
is_block_scaled = false;
gemm_workspace_.block_scales = std::nullopt;
}
is_moe = operation_desc.is_moe;
if (operation_desc.gemm.gemm_kind != library::GemmKind::kGrouped) {
return Status::kErrorInvalidProblem;
}
@@ -761,232 +782,561 @@ Status GroupedGemmOperationProfiler::initialize_workspace(
gemm_workspace_.reference_ptr_array_host.resize(num_groups);
int seed_shift = 0;
for (size_t group_idx = 0; group_idx < num_groups; group_idx++) {
auto group_str = std::to_string(group_idx);
gemm_workspace_.A_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor(
if (not operation_desc.is_moe) {
for (size_t group_idx = 0; group_idx < num_groups; group_idx++) {
auto group_str = std::to_string(group_idx);
gemm_workspace_.A_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor(
options,
"A_" + group_str,
operation_desc.gemm.A.element,
operation_desc.gemm.A.layout,
{int(problem_.m(group_idx)), int(problem_.k(group_idx))},
{int(problem_.lda[group_idx])},
gemm_workspace_.problem_count,
seed_shift++,
0);
gemm_workspace_.B_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor(
options,
"B_" + group_str,
operation_desc.gemm.B.element,
operation_desc.gemm.B.layout,
{int(problem_.k(group_idx)), int(problem_.n(group_idx))},
{int(problem_.ldb[group_idx])},
gemm_workspace_.problem_count,
seed_shift++,
0);
gemm_workspace_.C_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor(
options,
"C_" + group_str,
operation_desc.gemm.C.element,
operation_desc.gemm.C.layout,
{int(problem_.m(group_idx)), int(problem_.n(group_idx))},
{int(problem_.ldc[group_idx])},
gemm_workspace_.problem_count,
seed_shift++,
0);
gemm_workspace_.D_ptr_array_host[group_idx] = device_context.allocate_tensor(
options,
"D_" + group_str,
operation_desc.gemm.D.element,
operation_desc.gemm.D.layout,
{int(problem_.m(group_idx)), int(problem_.n(group_idx))},
{int(problem_.ldc[group_idx])},
gemm_workspace_.problem_count,
0);
gemm_workspace_.reference_ptr_array_host[group_idx] = device_context.allocate_tensor(
options,
"Reference_" + group_str,
operation_desc.gemm.D.element,
operation_desc.gemm.D.layout,
{int(problem_.m(group_idx)), int(problem_.n(group_idx))},
{int(problem_.ldc[group_idx])},
1,
0);
if (is_block_scaled) {
auto const block_scale_desc = operation_desc.block_scales.value();
auto& block_scale_ws = gemm_workspace_.block_scales.value();
int sfa_m = round_up(int(problem_.m(group_idx)), 128);
int sfb_n = round_up(int(problem_.n(group_idx)), 128);
int sfa_sfb_k =
round_up(ceil_div(int(problem_.k(group_idx)), block_scale_desc.SFKVecSize), 4);
int sfd_m =
block_scale_desc.SFD.layout == cutlass::library::LayoutTypeID::kRowMajor
? sfa_m
: round_up(ceil_div(int(problem_.m(group_idx)), block_scale_desc.EpilogueSFVecSize), 4);
int sfd_n =
block_scale_desc.SFD.layout == cutlass::library::LayoutTypeID::kRowMajor
? round_up(ceil_div(int(problem_.n(group_idx)), block_scale_desc.EpilogueSFVecSize), 4)
: sfb_n;
block_scale_ws.SFA_ptr_array_host[group_idx] =
device_context.allocate_and_initialize_tensor(
options,
"SFA",
block_scale_desc.SFA.element,
block_scale_desc.SFA.layout,
{sfa_m, sfa_sfb_k},
{sfa_sfb_k},
gemm_workspace_.problem_count,
seed_shift++,
0);
block_scale_ws.SFB_ptr_array_host[group_idx] =
device_context.allocate_and_initialize_tensor(
options,
"SFB",
block_scale_desc.SFB.element,
block_scale_desc.SFB.layout,
{sfb_n, sfa_sfb_k},
{sfa_sfb_k},
gemm_workspace_.problem_count,
seed_shift++,
0);
block_scale_ws.SFD_ptr_array_host[group_idx] = device_context.allocate_tensor(
options,
"SFD",
block_scale_desc.SFD.element,
block_scale_desc.SFD.layout,
{sfd_m, sfd_n},
{sfd_n},
gemm_workspace_.problem_count,
0);
block_scale_ws.SFD_reference_ptr_array_host[group_idx] = device_context.allocate_tensor(
options,
"Reference_SFD",
block_scale_desc.SFD.element,
block_scale_desc.SFD.layout,
{sfd_m, sfd_n},
{sfd_n},
gemm_workspace_.problem_count,
0);
// ScaleFactor tensor results may have some holes and will not be touched by the kernel.
// If we randomly fill the two tensors, these holes may encounter refcheck errors.
if (block_scale_ws.SFD_ptr_array_host[group_idx]->type() != library::NumericTypeID::kVoid) {
block_scale_ws.SFD_reference_ptr_array_host[group_idx]->fill_device(0);
block_scale_ws.SFD_ptr_array_host[group_idx]->fill_device(0);
}
}
else if (is_blockwise) {
auto const block_scale_desc = operation_desc.block_scales.value();
auto& block_scale_ws = gemm_workspace_.block_scales.value();
int sfa_m = ceil_div(int(problem_.m(group_idx)), block_scale_desc.SFMVecSize);
int sfb_n = ceil_div(int(problem_.n(group_idx)), block_scale_desc.SFNVecSize);
int sfa_sfb_k = ceil_div(int(problem_.k(group_idx)), block_scale_desc.SFKVecSize);
block_scale_ws.SFA_ptr_array_host[group_idx] =
device_context.allocate_and_initialize_tensor(
options,
"SFA_" + std::to_string(group_idx),
block_scale_desc.SFA.element,
block_scale_desc.SFA.layout,
{sfa_m, sfa_sfb_k},
{sfa_m},
gemm_workspace_.problem_count,
seed_shift++,
0);
block_scale_ws.SFB_ptr_array_host[group_idx] =
device_context.allocate_and_initialize_tensor(
options,
"SFB_" + std::to_string(group_idx),
block_scale_desc.SFB.element,
block_scale_desc.SFB.layout,
{sfa_sfb_k, sfb_n},
{sfb_n},
gemm_workspace_.problem_count,
seed_shift++,
0);
}
}
// takes the allocated tensors and initializes an array of pointers per problem in the workspace
auto create_dev_ptr_array_all_workspace = [&](
std::vector<DeviceAllocation*>& dev_ptr_arrays,
std::vector<DeviceAllocation*> const& input,
std::string const& id) {
auto num_workspaces = gemm_workspace_.problem_count;
dev_ptr_arrays.resize(num_workspaces);
// note "problem_count" here refers to input/output count for L2 cycling
for (int i = 0; i < gemm_workspace_.problem_count; i++) {
std::string name = id + "_ptr_array_workspace" + std::to_string(i);
dev_ptr_arrays[i] =
device_context.allocate_block(options, name, library::NumericTypeID::kU64, num_groups, 0);
std::vector<void*> group_ptrs(num_groups);
for (size_t group_idx = 0; group_idx < num_groups; group_idx++) {
group_ptrs[group_idx] = input[group_idx]->batch_data(i);
}
dev_ptr_arrays[i]->copy_from_host(group_ptrs.data());
}
};
create_dev_ptr_array_all_workspace(
gemm_workspace_.A_ptr_array_device,
gemm_workspace_.A_ptr_array_host,
"A");
create_dev_ptr_array_all_workspace(
gemm_workspace_.B_ptr_array_device,
gemm_workspace_.B_ptr_array_host,
"B");
create_dev_ptr_array_all_workspace(
gemm_workspace_.C_ptr_array_device,
gemm_workspace_.C_ptr_array_host,
"C");
create_dev_ptr_array_all_workspace(
gemm_workspace_.D_ptr_array_device,
gemm_workspace_.D_ptr_array_host,
"D");
if (is_block_scaled) {
auto& block_scale_ws = gemm_workspace_.block_scales.value();
create_dev_ptr_array_all_workspace(
block_scale_ws.SFA_ptr_array_device,
block_scale_ws.SFA_ptr_array_host,
"SFA");
create_dev_ptr_array_all_workspace(
block_scale_ws.SFB_ptr_array_device,
block_scale_ws.SFB_ptr_array_host,
"SFB");
create_dev_ptr_array_all_workspace(
block_scale_ws.SFD_ptr_array_device,
block_scale_ws.SFD_ptr_array_host,
"SFD");
block_scale_ws.norm_constant = device_context.allocate_and_initialize_tensor(
options,
"norm_constant",
operation_desc.gemm.element_epilogue,
operation_desc.gemm.A.layout, // copied, but should this be D layout?
{1, 1},
{1},
1,
seed_shift++,
0 // device_index
);
}
else if (is_blockwise) {
auto& block_scale_ws = gemm_workspace_.block_scales.value();
create_dev_ptr_array_all_workspace(
block_scale_ws.SFA_ptr_array_device,
block_scale_ws.SFA_ptr_array_host,
"SFA");
create_dev_ptr_array_all_workspace(
block_scale_ws.SFB_ptr_array_device,
block_scale_ws.SFB_ptr_array_host,
"SFB");
}
} else {
int max_m = problem_.max_problem_size_3x[0];
int max_n = problem_.max_problem_size_3x[1];
int max_k = problem_.max_problem_size_3x[2];
// allocate the block tensors
DeviceAllocation* block_A;
DeviceAllocation* block_B;
DeviceAllocation* block_C;
DeviceAllocation* block_D;
DeviceAllocation* block_ref_D;
DeviceAllocation* block_ref_SFD;
DeviceAllocation* block_SFA;
DeviceAllocation* block_SFB;
DeviceAllocation* block_SFD;
gemm_workspace_.tokens_per_expert_host.resize(num_groups);
gemm_workspace_.tokens_per_expert_device = device_context.allocate_block(
options,
"A_" + group_str,
"tokens_per_expert",
library::NumericTypeID::kU32,
num_groups,
0);
block_A = device_context.allocate_and_initialize_tensor(
options,
"block_A",
operation_desc.gemm.A.element,
operation_desc.gemm.A.layout,
{int(problem_.m(group_idx)), int(problem_.k(group_idx))},
{int(problem_.lda[group_idx])},
gemm_workspace_.problem_count,
{max_m, max_k},
{int(problem_.max_lda)},
gemm_workspace_.problem_count * num_groups,
seed_shift++,
0);
gemm_workspace_.B_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor(
block_B = device_context.allocate_and_initialize_tensor(
options,
"B_" + group_str,
"block_B",
operation_desc.gemm.B.element,
operation_desc.gemm.B.layout,
{int(problem_.k(group_idx)), int(problem_.n(group_idx))},
{int(problem_.ldb[group_idx])},
gemm_workspace_.problem_count,
{max_k, max_n},
{int(problem_.max_ldb)},
gemm_workspace_.problem_count * num_groups,
seed_shift++,
0);
gemm_workspace_.C_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor(
block_C = device_context.allocate_and_initialize_tensor(
options,
"C_" + group_str,
"block_C",
operation_desc.gemm.C.element,
operation_desc.gemm.C.layout,
{int(problem_.m(group_idx)), int(problem_.n(group_idx))},
{int(problem_.ldc[group_idx])},
gemm_workspace_.problem_count,
{max_m, max_n},
{int(problem_.max_ldc)},
gemm_workspace_.problem_count * num_groups,
seed_shift++,
0);
gemm_workspace_.D_ptr_array_host[group_idx] = device_context.allocate_tensor(
block_D = device_context.allocate_tensor(
options,
"D_" + group_str,
"block_D",
operation_desc.gemm.D.element,
operation_desc.gemm.D.layout,
{int(problem_.m(group_idx)), int(problem_.n(group_idx))},
{int(problem_.ldc[group_idx])},
gemm_workspace_.problem_count,
{max_m, max_n},
{int(problem_.max_ldc)},
gemm_workspace_.problem_count * num_groups,
0);
block_ref_D = device_context.allocate_tensor(
options,
"Block_Reference",
operation_desc.gemm.D.element,
operation_desc.gemm.D.layout,
{max_m, max_n},
{int(problem_.max_ldc)},
num_groups,
0);
gemm_workspace_.reference_ptr_array_host[group_idx] = device_context.allocate_tensor(
options,
"Reference_" + group_str,
operation_desc.gemm.D.element,
operation_desc.gemm.D.layout,
{int(problem_.m(group_idx)), int(problem_.n(group_idx))},
{int(problem_.ldc[group_idx])},
1,
0);
if (is_block_scaled) {
auto const block_scale_desc = operation_desc.block_scales.value();
auto& block_scale_ws = gemm_workspace_.block_scales.value();
int sfa_m = round_up(int(problem_.m(group_idx)), 128);
int sfb_n = round_up(int(problem_.n(group_idx)), 128);
int sfa_m = round_up(problem_.max_problem_size_3x[0], 128);
int sfb_n = round_up(max_n, 128);
int sfa_sfb_k =
round_up(ceil_div(int(problem_.k(group_idx)), block_scale_desc.SFKVecSize), 4);
round_up(ceil_div(max_k, block_scale_desc.SFKVecSize), 4);
int sfd_m =
block_scale_desc.SFD.layout == cutlass::library::LayoutTypeID::kRowMajor
? sfa_m
: round_up(ceil_div(int(problem_.m(group_idx)), block_scale_desc.EpilogueSFVecSize), 4);
: round_up(ceil_div(max_m, block_scale_desc.EpilogueSFVecSize), 4);
int sfd_n =
block_scale_desc.SFD.layout == cutlass::library::LayoutTypeID::kRowMajor
? round_up(ceil_div(int(problem_.n(group_idx)), block_scale_desc.EpilogueSFVecSize), 4)
? round_up(ceil_div(max_n, block_scale_desc.EpilogueSFVecSize), 4)
: sfb_n;
block_scale_ws.SFA_ptr_array_host[group_idx] =
block_SFA =
device_context.allocate_and_initialize_tensor(
options,
"SFA",
"block_SFA",
block_scale_desc.SFA.element,
block_scale_desc.SFA.layout,
{sfa_m, sfa_sfb_k},
{sfa_sfb_k},
gemm_workspace_.problem_count,
gemm_workspace_.problem_count * num_groups,
seed_shift++,
0);
block_scale_ws.SFB_ptr_array_host[group_idx] =
block_SFB =
device_context.allocate_and_initialize_tensor(
options,
"SFB",
"block_SFB",
block_scale_desc.SFB.element,
block_scale_desc.SFB.layout,
{sfb_n, sfa_sfb_k},
{sfa_sfb_k},
gemm_workspace_.problem_count,
gemm_workspace_.problem_count * num_groups,
seed_shift++,
0);
block_scale_ws.SFD_ptr_array_host[group_idx] = device_context.allocate_tensor(
block_SFD = device_context.allocate_tensor(
options,
"SFD",
"block_SFD",
block_scale_desc.SFD.element,
block_scale_desc.SFD.layout,
{sfd_m, sfd_n},
{sfd_n},
gemm_workspace_.problem_count,
gemm_workspace_.problem_count * num_groups,
0);
block_scale_ws.SFD_reference_ptr_array_host[group_idx] = device_context.allocate_tensor(
block_ref_SFD = device_context.allocate_tensor(
options,
"Reference_SFD",
"block_Reference_SFD",
block_scale_desc.SFD.element,
block_scale_desc.SFD.layout,
{sfd_m, sfd_n},
{sfd_n},
gemm_workspace_.problem_count,
gemm_workspace_.problem_count * num_groups,
0);
// ScaleFactor tensor results may have some holes and will not be touched by the kernel.
// If we randomly fill the two tensors, these holes may encounter refcheck errors.
if (block_scale_ws.SFD_ptr_array_host[group_idx]->type() != library::NumericTypeID::kVoid) {
block_scale_ws.SFD_reference_ptr_array_host[group_idx]->fill_device(0);
block_scale_ws.SFD_ptr_array_host[group_idx]->fill_device(0);
}
}
else if (is_blockwise) {
auto const block_scale_desc = operation_desc.block_scales.value();
auto& block_scale_ws = gemm_workspace_.block_scales.value();
int sfa_m = ceil_div(int(problem_.m(group_idx)), block_scale_desc.SFMVecSize);
int sfb_n = ceil_div(int(problem_.n(group_idx)), block_scale_desc.SFNVecSize);
int sfa_sfb_k = ceil_div(int(problem_.k(group_idx)), block_scale_desc.SFKVecSize);
block_scale_ws.SFA_ptr_array_host[group_idx] =
device_context.allocate_and_initialize_tensor(
options,
"SFA_" + std::to_string(group_idx),
block_scale_desc.SFA.element,
block_scale_desc.SFA.layout,
{sfa_m, sfa_sfb_k},
{sfa_m},
gemm_workspace_.problem_count,
seed_shift++,
0);
block_scale_ws.SFB_ptr_array_host[group_idx] =
device_context.allocate_and_initialize_tensor(
options,
"SFB_" + std::to_string(group_idx),
block_scale_desc.SFB.element,
block_scale_desc.SFB.layout,
{sfa_sfb_k, sfb_n},
{sfb_n},
gemm_workspace_.problem_count,
seed_shift++,
0);
}
}
// takes the allocated tensors and initializes an array of pointers per problem in the workspace
auto create_dev_ptr_array_all_workspace = [&](
std::vector<DeviceAllocation*>& dev_ptr_arrays,
std::vector<DeviceAllocation*> const& input,
std::string const& id) {
auto num_workspaces = gemm_workspace_.problem_count;
dev_ptr_arrays.resize(num_workspaces);
// note "problem_count" here refers to input/output count for L2 cycling
for (int i = 0; i < gemm_workspace_.problem_count; i++) {
std::string name = id + "_ptr_array_workspace" + std::to_string(i);
dev_ptr_arrays[i] =
device_context.allocate_block(options, name, library::NumericTypeID::kU64, num_groups, 0);
std::vector<void*> group_ptrs(num_groups);
block_scale_ws.norm_constant = device_context.allocate_and_initialize_tensor(
options,
"norm_constant",
operation_desc.gemm.element_epilogue,
operation_desc.gemm.A.layout, // copied, but should this be D layout?
{1, 1},
{1},
1,
seed_shift++,
0 // device_index
);
gemm_workspace_.block_scales.value().SFA_ptr_array_device.resize(gemm_workspace_.problem_count);
gemm_workspace_.block_scales.value().SFB_ptr_array_device.resize(gemm_workspace_.problem_count);
gemm_workspace_.block_scales.value().SFD_ptr_array_device.resize(gemm_workspace_.problem_count);
for (size_t group_idx = 0; group_idx < num_groups; group_idx++) {
group_ptrs[group_idx] = input[group_idx]->batch_data(i);
auto group_str = std::to_string(group_idx);
block_scale_ws.SFA_ptr_array_host[group_idx] = device_context.create_ref_tensor(
options,
"block_SFA" + group_str,
block_scale_desc.SFA.element,
block_scale_desc.SFA.layout,
{sfa_m, sfa_sfb_k},
{sfa_sfb_k},
block_SFA->batch_data(group_idx),
1,
0);
block_scale_ws.SFB_ptr_array_host[group_idx] = device_context.create_ref_tensor(
options,
"block_SFB" + group_str,
block_scale_desc.SFB.element,
block_scale_desc.SFB.layout,
{sfb_n, sfa_sfb_k},
{sfa_sfb_k},
block_SFB->batch_data(group_idx),
1,
0);
block_scale_ws.SFD_ptr_array_host[group_idx] = device_context.create_ref_tensor(
options,
"block_SFD" + group_str,
block_scale_desc.SFD.element,
block_scale_desc.SFD.layout,
{sfd_m, sfd_n},
{sfd_n},
block_SFD->batch_data(group_idx),
1,
0);
block_scale_ws.SFD_reference_ptr_array_host[group_idx] = device_context.create_ref_tensor(
options,
"block_Reference_SFD" + group_str,
block_scale_desc.SFD.element,
block_scale_desc.SFD.layout,
{sfd_m, sfd_n},
{sfd_n},
block_ref_SFD->batch_data(group_idx),
1,
0);
}
for(int problem_idx = 0; problem_idx < gemm_workspace_.problem_count; problem_idx++) {
auto problem_str = std::to_string(problem_idx);
gemm_workspace_.block_scales.value().SFA_ptr_array_device[problem_idx] = device_context.create_ref_tensor(
options,
"block_SFA" + problem_str,
block_scale_desc.SFA.element,
block_scale_desc.SFA.layout,
{sfa_m, sfa_sfb_k},
{sfa_sfb_k},
block_SFA->batch_data(problem_idx*num_groups),
num_groups,
0);
gemm_workspace_.block_scales.value().SFB_ptr_array_device[problem_idx] = device_context.create_ref_tensor(
options,
"block_SFB" + problem_str,
block_scale_desc.SFB.element,
block_scale_desc.SFB.layout,
{sfb_n, sfa_sfb_k},
{sfa_sfb_k},
block_SFB->batch_data(problem_idx*num_groups),
num_groups,
0);
gemm_workspace_.block_scales.value().SFD_ptr_array_device[problem_idx] = device_context.create_ref_tensor(
options,
"block_SFD" + problem_str,
block_scale_desc.SFD.element,
block_scale_desc.SFD.layout,
{sfd_m, sfd_n},
{sfd_n},
block_SFD->batch_data(problem_idx*num_groups),
num_groups,
0);
}
dev_ptr_arrays[i]->copy_from_host(group_ptrs.data());
}
};
create_dev_ptr_array_all_workspace(
gemm_workspace_.A_ptr_array_device,
gemm_workspace_.A_ptr_array_host,
"A");
create_dev_ptr_array_all_workspace(
gemm_workspace_.B_ptr_array_device,
gemm_workspace_.B_ptr_array_host,
"B");
create_dev_ptr_array_all_workspace(
gemm_workspace_.C_ptr_array_device,
gemm_workspace_.C_ptr_array_host,
"C");
create_dev_ptr_array_all_workspace(
gemm_workspace_.D_ptr_array_device,
gemm_workspace_.D_ptr_array_host,
"D");
if (is_block_scaled) {
auto& block_scale_ws = gemm_workspace_.block_scales.value();
create_dev_ptr_array_all_workspace(
block_scale_ws.SFA_ptr_array_device,
block_scale_ws.SFA_ptr_array_host,
"SFA");
create_dev_ptr_array_all_workspace(
block_scale_ws.SFB_ptr_array_device,
block_scale_ws.SFB_ptr_array_host,
"SFB");
create_dev_ptr_array_all_workspace(
block_scale_ws.SFD_ptr_array_device,
block_scale_ws.SFD_ptr_array_host,
"SFD");
for (size_t group_idx = 0; group_idx < num_groups; group_idx++) {
gemm_workspace_.tokens_per_expert_host[group_idx] = problem_.n(group_idx);
auto group_str = std::to_string(group_idx);
block_scale_ws.norm_constant = device_context.allocate_and_initialize_tensor(
options,
"norm_constant",
operation_desc.gemm.element_epilogue,
operation_desc.gemm.A.layout, // copied, but should this be D layout?
{1, 1},
{1},
1,
seed_shift++,
0 // device_index
);
}
else if (is_blockwise) {
auto& block_scale_ws = gemm_workspace_.block_scales.value();
create_dev_ptr_array_all_workspace(
block_scale_ws.SFA_ptr_array_device,
block_scale_ws.SFA_ptr_array_host,
"SFA");
create_dev_ptr_array_all_workspace(
block_scale_ws.SFB_ptr_array_device,
block_scale_ws.SFB_ptr_array_host,
"SFB");
gemm_workspace_.A_ptr_array_host[group_idx] = device_context.create_ref_tensor(
options,
"block_A" + group_str,
operation_desc.gemm.A.element,
operation_desc.gemm.A.layout,
{max_m, max_k},
{int(problem_.max_lda)},
block_A->batch_data(group_idx),
1,
0);
gemm_workspace_.B_ptr_array_host[group_idx] = device_context.create_ref_tensor(
options,
"block_B" + group_str,
operation_desc.gemm.B.element,
operation_desc.gemm.B.layout,
{max_k, max_n},
{int(problem_.max_ldb)},
block_B->batch_data(group_idx),
1,
0);
gemm_workspace_.C_ptr_array_host[group_idx] = device_context.create_ref_tensor(
options,
"block_C" + group_str,
operation_desc.gemm.C.element,
operation_desc.gemm.C.layout,
{max_m, max_n},
{int(problem_.max_ldc)},
block_C->batch_data(group_idx),
1,
0);
gemm_workspace_.D_ptr_array_host[group_idx] = device_context.create_ref_tensor(
options,
"block_D" + group_str,
operation_desc.gemm.D.element,
operation_desc.gemm.D.layout,
{max_m, max_n},
{int(problem_.max_ldc)},
block_D->batch_data(group_idx),
1,
0);
gemm_workspace_.reference_ptr_array_host[group_idx] = device_context.create_ref_tensor(
options,
"Reference_" + group_str,
operation_desc.gemm.D.element,
operation_desc.gemm.D.layout,
{max_m, max_n},
{int(problem_.max_ldc)},
block_ref_D->batch_data(group_idx),
1,
0);
gemm_workspace_.A_ptr_array_device.resize(gemm_workspace_.problem_count);
gemm_workspace_.B_ptr_array_device.resize(gemm_workspace_.problem_count);
gemm_workspace_.C_ptr_array_device.resize(gemm_workspace_.problem_count);
gemm_workspace_.D_ptr_array_device.resize(gemm_workspace_.problem_count);
for(int problem_idx = 0; problem_idx < gemm_workspace_.problem_count; problem_idx++) {
auto problem_str = std::to_string(problem_idx);
gemm_workspace_.A_ptr_array_device[problem_idx] = device_context.create_ref_tensor(
options,
"block_A" + problem_str,
operation_desc.gemm.A.element,
operation_desc.gemm.A.layout,
{max_m, max_k},
{int(problem_.max_lda)},
block_A->batch_data(problem_idx*num_groups),
num_groups,
0);
gemm_workspace_.B_ptr_array_device[problem_idx] = device_context.create_ref_tensor(
options,
"block_B" + problem_str,
operation_desc.gemm.B.element,
operation_desc.gemm.B.layout,
{max_k, max_n},
{int(problem_.max_ldb)},
block_B->batch_data(problem_idx*num_groups),
num_groups,
0);
gemm_workspace_.C_ptr_array_device[problem_idx] = device_context.create_ref_tensor(
options,
"block_C" + problem_str,
operation_desc.gemm.C.element,
operation_desc.gemm.C.layout,
{max_m, max_n},
{int(problem_.max_ldc)},
block_C->batch_data(problem_idx*num_groups),
num_groups,
0);
gemm_workspace_.D_ptr_array_device[problem_idx] = device_context.create_ref_tensor(
options,
"block_D" + problem_str,
operation_desc.gemm.D.element,
operation_desc.gemm.D.layout,
{max_m, max_n},
{int(problem_.max_ldc)},
block_D->batch_data(problem_idx*num_groups),
num_groups,
0);
}
}
gemm_workspace_.tokens_per_expert_device->copy_from_host(gemm_workspace_.tokens_per_expert_host.data());
}
init_arguments(options);
tools/util/include/cutlass/util/reference/device/gemm_complex.h
+11 -6
View File
@@ -108,7 +108,9 @@ __global__ void GemmComplex(
tensor_a.add_pointer_offset(batch_idx * batch_stride_A);
tensor_b.add_pointer_offset(batch_idx * batch_stride_B);
tensor_c.add_pointer_offset(batch_idx * batch_stride_C);
if(tensor_c.data()) {
tensor_c.add_pointer_offset(batch_idx * batch_stride_C);
}
tensor_d.add_pointer_offset(batch_idx * batch_stride_D);
for (; batch_idx < batch_count; batch_idx += gridDim.z) {
@@ -163,17 +165,20 @@ __global__ void GemmComplex(
MatrixCoord coord = MatrixCoord(row, col);
if (row < M && col < N) {
tensor_d.at(coord) = convert_op(
alpha * ScalarType(accum[i][j]) +
beta * ScalarType(tensor_c.at(coord)));
ScalarType epilog = alpha * ScalarType(accum[i][j]);
if(tensor_c.data()) {
epilog += beta * ScalarType(tensor_c.at(coord));
}
tensor_d.at(coord) = convert_op(epilog);
}
}
}
tensor_a.add_pointer_offset(batch_stride_A * gridDim.z);
tensor_b.add_pointer_offset(batch_stride_B * gridDim.z);
tensor_c.add_pointer_offset(batch_stride_C * gridDim.z);
if(tensor_c.data()) {
tensor_c.add_pointer_offset(batch_stride_C * gridDim.z);
}
tensor_d.add_pointer_offset(batch_stride_D * gridDim.z);
} // for (batch_idx)
tools/util/include/cutlass/util/reference/host/gemm_complex.h
+8 -5
View File
@@ -154,10 +154,11 @@ void GemmComplex(
MatrixCoord coord = MatrixCoord(row, col);
if (row < M && col < N) {
tensor_d.at(coord) = convert_op(
alpha * ScalarType(accum[i][j]) +
beta * ScalarType(tensor_c.at(coord)));
ScalarType epilog = alpha * ScalarType(accum[i][j]);
if(tensor_c.data()) {
epilog += beta * ScalarType(tensor_c.at(coord));
}
tensor_d.at(coord) = convert_op(epilog);
}
}
}
@@ -167,7 +168,9 @@ void GemmComplex(
tensor_a.add_pointer_offset(batch_stride_A);
tensor_b.add_pointer_offset(batch_stride_B);
tensor_c.add_pointer_offset(batch_stride_C);
if(tensor_c.data()) {
tensor_c.add_pointer_offset(batch_stride_C);
}
tensor_d.add_pointer_offset(batch_stride_D);
} // for (batch_idx)