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

update 3.8 v2 (#2112)

Browse Source 
* update 3.8 v2

* update 3.8

---------

Co-authored-by: yuzhai <yuzhai@nvidia.com>
This commit is contained in:
Yujia Zhai
2025-02-19 19:03:14 -08:00
committed by GitHub
parent e9627ce55b
commit b84e9802d8
166 changed files with 3986 additions and 4037 deletions
Download Patch File Download Diff File Expand all files Collapse all files

25
tools/library/include/cutlass/library/descriptions.h
View File

@@ -35,6 +35,8 @@
#include <cutlass/blas3_types.h>
#include <cutlass/gemm_coord.h>
#include <optional>
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
@@ -300,13 +302,34 @@ struct GemmDescription : public OperationDescription {
transform_B(transform_B) {}
};
struct BlockScaleDescription {
/// Describes the SFA operand
TensorDescription SFA;
/// Describes the SFB operand
TensorDescription SFB;
/// Describes the SFD operand
TensorDescription SFD;
/// Describes the input ScaleFactor VectorSize
int SFVecSize;
/// Describes the Output ScaleFactor VectorSize
int EpilogueSFVecSize;
};
struct GroupedGemmDescription : public OperationDescription {
GemmDescription gemm;
std::optional<BlockScaleDescription> block_scales;
};
/// Description of all GEMM computations
struct BlockScaledGemmDescription : public OperationDescription {
/// Indicates the kind of GEMM performed
GemmKind gemm_kind;
/// Describes the A operand
TensorDescription A;

32
tools/library/include/cutlass/library/library.h
View File

@@ -336,11 +336,10 @@ struct GemmUniversalArguments {
void *packed_Scale{nullptr}; // Packed scale for int4 * fp8
int device_index{0};
bool use_pdl{false};
};
/// Block Scaled GEMM
//
// OperationKind: kBlockScaledGemm
@@ -495,29 +494,31 @@ struct GemmGroupedConfiguration {
int64_t* lda;
int64_t* ldb;
int64_t* ldc;
cute::Shape<int, int, int>* problem_sizes_3x_host;
};
struct GemmGroupedArguments {
int problem_count{};
gemm::GemmCoord* problem_sizes{nullptr};
void * ptr_A{nullptr};
void * ptr_B{nullptr};
void * ptr_C{nullptr};
void * ptr_D{nullptr};
void* ptr_A{nullptr};
void* ptr_B{nullptr};
void* ptr_C{nullptr};
void* ptr_D{nullptr};
int64_t *lda{nullptr};
int64_t *ldb{nullptr};
int64_t *ldc{nullptr};
int64_t *ldd{nullptr};
int64_t* lda{nullptr};
int64_t* ldb{nullptr};
int64_t* ldc{nullptr};
int64_t* ldd{nullptr};
void const *alpha{nullptr};
void const *beta{nullptr};
ScalarPointerMode pointer_mode{};
bool use_pdl{false};
gemm::GemmCoord cluster_shape{};
gemm::GemmCoord cluster_shape_fallback{};
gemm::GemmCoord cluster_shape{};
gemm::GemmCoord cluster_shape_fallback{};
// these should really be in the configuration but staying consistent with GEMM
int sm_count{0};
@@ -529,6 +530,13 @@ struct GemmGroupedArguments {
cute::Shape<int, int, int>* problem_sizes_3x_host;
};
struct GroupedGemmBlockScaledArguments : GemmGroupedArguments {
void* SFA{nullptr};
void* SFB{nullptr};
void* SFD{nullptr};
void* norm_constant{nullptr};
};
/////////////////////////////////////////////////////////////////////////////////////////////////
//
// OperationKind: kSparseGemm

13
tools/library/src/gemm_operation.h
View File

@@ -1200,13 +1200,24 @@ public:
GemmGroupedOperation(char const *name = "unknown_gemm"):
GemmOperationBase<Operator_>(name) {
this->description_.gemm_kind = GemmKind::kGrouped;
this->description_.kind = OperationKind::kGroupedGemm;
this->description_.provider = Provider::kCUTLASS;
this->threadblock_count = Operator::sufficient();
this->description_.gemm = GemmOperationBase<Operator_>::description_;
this->description_.gemm.gemm_kind = GemmKind::kGrouped;
this->description_.tile_description = this->description_.gemm.tile_description;
}
/// Returns the description of the GroupedGEMM operation
virtual OperationDescription const & description() const override final {
return description_;
}
private:
int threadblock_count;
GroupedGemmDescription description_;
protected:

515
tools/library/src/grouped_gemm_operation_3x.hpp
View File

@@ -41,17 +41,11 @@
#include "cutlass/library/util.h"
#include "gemm_operation_3x.hpp"
#include "library_internal.h"
#include <unordered_map>
///////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass::library {
/// **** CAUTION ****
/// Unlike other operations, initialize() must be called when
/// certain arguments change. See initialize() for details.
template <typename Operator_>
class GroupedGemmUniversal3xOperation : public GemmOperation3xBase<Operator_> {
class GroupedGemmOperation3xBase : public GemmOperation3xBase<Operator_> {
public:
using Operator = Operator_;
using OperatorArguments = typename Operator::Arguments;
@@ -70,20 +64,15 @@ public:
using CollectiveEpilogue = typename Operator::CollectiveEpilogue;
using ThreadEpilogueOp = typename CollectiveEpilogue::ThreadEpilogueOp;
private:
mutable CudaBuffer strideA_device;
mutable CudaBuffer strideB_device;
mutable CudaBuffer strideC_device;
mutable CudaBuffer strideD_device;
mutable std::vector<typename Operator::GemmKernel::InternalStrideA> strideA_host;
mutable std::vector<typename Operator::GemmKernel::InternalStrideB> strideB_host;
mutable std::vector<typename Operator::GemmKernel::InternalStrideC> strideC_host;
mutable std::vector<typename Operator::GemmKernel::InternalStrideD> strideD_host;
public:
GroupedGemmUniversal3xOperation(char const* name = "unknown_gemm")
GroupedGemmOperation3xBase(char const* name = "unknown_gemm")
: GemmOperation3xBase<Operator_>(name, GemmKind::kGrouped) {
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;
if constexpr (Operator::ArchTag::kMinComputeCapability >= 90) {
dim3 cluster_dims(
cute::size<0>(typename Operator::GemmKernel::ClusterShape{}),
@@ -96,8 +85,157 @@ public:
threads_per_block,
kernel_ptr);
}
};
public:
mutable CudaBuffer strideA_device;
mutable CudaBuffer strideB_device;
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_;
int max_active_clusters;
Status initialize_strides(GemmGroupedConfiguration const& config) const {
auto const num_groups = config.problem_count;
this->strideA_device =
CudaBuffer(sizeof(typename Operator::GemmKernel::InternalStrideA) * num_groups);
this->strideB_device =
CudaBuffer(sizeof(typename Operator::GemmKernel::InternalStrideB) * num_groups);
this->strideC_device =
CudaBuffer(sizeof(typename Operator::GemmKernel::InternalStrideC) * num_groups);
this->strideD_device =
CudaBuffer(sizeof(typename Operator::GemmKernel::InternalStrideD) * num_groups);
std::vector<typename Operator::GemmKernel::InternalStrideA> strideA_host(num_groups);
std::vector<typename Operator::GemmKernel::InternalStrideB> strideB_host(num_groups);
std::vector<typename Operator::GemmKernel::InternalStrideC> strideC_host(num_groups);
std::vector<typename Operator::GemmKernel::InternalStrideD> strideD_host(num_groups);
for (int group_idx = 0; group_idx < num_groups; group_idx++) {
strideA_host[group_idx] =
cute::make_int_tuple_from<typename Operator::GemmKernel::InternalStrideA>(
config.lda[group_idx]);
strideB_host[group_idx] =
cute::make_int_tuple_from<typename Operator::GemmKernel::InternalStrideB>(
config.ldb[group_idx]);
strideC_host[group_idx] =
cute::make_int_tuple_from<typename Operator::GemmKernel::InternalStrideC>(
config.ldc[group_idx]);
strideD_host[group_idx] =
cute::make_int_tuple_from<typename Operator::GemmKernel::InternalStrideD>(
config.ldc[group_idx]);
}
CUDA_CHECK(cudaMemcpy(
this->strideA_device.data(),
strideA_host.data(),
sizeof(typename Operator::GemmKernel::InternalStrideA) * num_groups,
cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(
this->strideB_device.data(),
strideB_host.data(),
sizeof(typename Operator::GemmKernel::InternalStrideB) * num_groups,
cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(
this->strideC_device.data(),
strideC_host.data(),
sizeof(typename Operator::GemmKernel::InternalStrideC) * num_groups,
cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(
this->strideD_device.data(),
strideD_host.data(),
sizeof(typename Operator::GemmKernel::InternalStrideD) * num_groups,
cudaMemcpyHostToDevice));
return Status::kSuccess;
}
/// 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;
operator_args.problem_shape = {
arguments.problem_count,
arguments.problem_sizes_3x,
arguments.pointer_mode == ScalarPointerMode::kHost ? arguments.problem_sizes_3x_host
: nullptr};
operator_args.mainloop.ptr_A = static_cast<ElementA const**>(arguments.ptr_A);
operator_args.mainloop.ptr_B = static_cast<ElementB const**>(arguments.ptr_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.mainloop.dA =
static_cast<typename Operator::GemmKernel::InternalStrideA*>(this->strideA_device.data());
operator_args.mainloop.dB =
static_cast<typename Operator::GemmKernel::InternalStrideB*>(this->strideB_device.data());
operator_args.epilogue.dC =
static_cast<typename Operator::GemmKernel::InternalStrideC*>(this->strideC_device.data());
operator_args.epilogue.dD =
static_cast<typename Operator::GemmKernel::InternalStrideD*>(this->strideD_device.data());
operator_args.hw_info.sm_count = arguments.sm_count;
if constexpr (Operator::ArchTag::kMinComputeCapability >= 90) {
operator_args.hw_info.max_active_clusters = max_active_clusters;
}
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;
fusion_args.alpha_ptr_array = nullptr;
fusion_args.beta_ptr_array = nullptr;
// Single alpha and beta for all groups
fusion_args.dAlpha = {cute::_0{}, cute::_0{}, 0};
fusion_args.dBeta = {cute::_0{}, cute::_0{}, 0};
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;
}
}
};
/// **** CAUTION ****
/// Unlike other operations, initialize() must be called when
/// certain arguments change. See initialize() for details.
template <typename Operator_>
class GroupedGemmUniversal3xOperation : public GroupedGemmOperation3xBase<Operator_> {
public:
using Operator = Operator_;
using OperatorArguments = typename Operator::Arguments;
public:
GroupedGemmUniversal3xOperation(char const* name = "unknown_gemm")
: GroupedGemmOperation3xBase<Operator_>(name) {}
~GroupedGemmUniversal3xOperation() override = default;
private:
@@ -115,29 +253,7 @@ protected:
template <class FusionArgs>
struct UpdateFusionArgs<FusionArgs, cute::void_t<decltype(FusionArgs{}.alpha)>> {
static Status update_(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;
fusion_args.alpha_ptr_array = nullptr;
fusion_args.beta_ptr_array = nullptr;
// Single alpha and beta for all groups
fusion_args.dAlpha = {cute::_0{}, cute::_0{}, 0};
fusion_args.dBeta = {cute::_0{}, cute::_0{}, 0};
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;
}
return GroupedGemmOperation3xBase<Operator>::update_fusion_args(fusion_args, arguments);
}
};
@@ -152,46 +268,7 @@ protected:
return status;
}
operator_args.mode = cutlass::gemm::GemmUniversalMode::kGrouped;
operator_args.problem_shape = {
arguments->problem_count,
arguments->problem_sizes_3x,
arguments->pointer_mode == ScalarPointerMode::kHost ? arguments->problem_sizes_3x_host
: nullptr};
operator_args.mainloop.ptr_A =
static_cast<const typename Operator::ElementA**>(arguments->ptr_A);
operator_args.mainloop.ptr_B =
static_cast<const typename Operator::ElementB**>(arguments->ptr_B);
operator_args.epilogue.ptr_C =
static_cast<const typename Operator::ElementC**>(arguments->ptr_C);
operator_args.epilogue.ptr_D = static_cast<typename Operator::ElementD**>(arguments->ptr_D);
operator_args.mainloop.dA =
static_cast<typename Operator::GemmKernel::InternalStrideA*>(strideA_device.data());
operator_args.mainloop.dB =
static_cast<typename Operator::GemmKernel::InternalStrideB*>(strideB_device.data());
operator_args.epilogue.dC =
static_cast<typename Operator::GemmKernel::InternalStrideC*>(strideC_device.data());
operator_args.epilogue.dD =
static_cast<typename Operator::GemmKernel::InternalStrideD*>(strideD_device.data());
operator_args.hw_info.sm_count = arguments->sm_count;
if constexpr (Operator::ArchTag::kMinComputeCapability >= 90) {
operator_args.hw_info.max_active_clusters = max_active_clusters;
}
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());
}
status = this->update_arguments_base(operator_args, *arguments);
return status;
}
@@ -201,7 +278,6 @@ public:
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;
@@ -211,11 +287,6 @@ public:
return status;
}
/// Gets the host-side workspace
uint64_t get_host_workspace_size(void const* configuration) const override {
return sizeof(Operator);
}
/// Gets the device-side workspace
uint64_t get_device_workspace_size(void const* configuration_ptr, void const* arguments_ptr)
const override {
@@ -246,59 +317,10 @@ public:
void* device_workspace,
cudaStream_t stream = nullptr) const override {
auto const& config = *static_cast<GemmGroupedConfiguration const*>(configuration_ptr);
auto num_groups = config.problem_count;
strideA_device =
CudaBuffer(sizeof(typename Operator::GemmKernel::InternalStrideA) * num_groups);
strideB_device =
CudaBuffer(sizeof(typename Operator::GemmKernel::InternalStrideB) * num_groups);
strideC_device =
CudaBuffer(sizeof(typename Operator::GemmKernel::InternalStrideC) * num_groups);
strideD_device =
CudaBuffer(sizeof(typename Operator::GemmKernel::InternalStrideD) * num_groups);
strideA_host.resize(num_groups);
strideB_host.resize(num_groups);
strideC_host.resize(num_groups);
strideD_host.resize(num_groups);
for (int group_idx = 0; group_idx < num_groups; group_idx++) {
strideA_host[group_idx] =
cute::make_int_tuple_from<typename Operator::GemmKernel::InternalStrideA>(
config.lda[group_idx]);
strideB_host[group_idx] =
cute::make_int_tuple_from<typename Operator::GemmKernel::InternalStrideB>(
config.ldb[group_idx]);
strideC_host[group_idx] =
cute::make_int_tuple_from<typename Operator::GemmKernel::InternalStrideC>(
config.ldc[group_idx]);
strideD_host[group_idx] =
cute::make_int_tuple_from<typename Operator::GemmKernel::InternalStrideD>(
config.ldc[group_idx]);
}
CUDA_CHECK(cudaMemcpy(
strideA_device.data(),
strideA_host.data(),
sizeof(typename Operator::GemmKernel::InternalStrideA) * num_groups,
cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(
strideB_device.data(),
strideB_host.data(),
sizeof(typename Operator::GemmKernel::InternalStrideB) * num_groups,
cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(
strideC_device.data(),
strideC_host.data(),
sizeof(typename Operator::GemmKernel::InternalStrideC) * num_groups,
cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(
strideD_device.data(),
strideD_host.data(),
sizeof(typename Operator::GemmKernel::InternalStrideD) * num_groups,
cudaMemcpyHostToDevice));
Operator* op = new (host_workspace) Operator;
return Status::kSuccess;
auto const& config = *static_cast<GemmGroupedConfiguration const*>(configuration_ptr);
return this->initialize_strides(config);
}
/// **** CAUTION ****
@@ -323,8 +345,215 @@ public:
return status;
}
};
///////////////////////////////////////////////////////////////////////////////////////////////////
template <typename Operator_>
class GroupedBlockScaledGemmUniversal3xOperation : public GroupedGemmOperation3xBase<Operator_> {
public:
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 ElementSFA = typename Operator::CollectiveMainloop::ElementSF;
using ElementSFB = 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>;
GroupedBlockScaledGemmUniversal3xOperation(char const* name = "unknown_gemm")
: GroupedGemmOperation3xBase<Operator_>(name) {
BlockScaleDescription block_scaled_desc{};
block_scaled_desc.SFA.element = NumericTypeMap<ElementSFA>::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<ElementSFB>::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.SFVecSize = SFVecSize;
block_scaled_desc.SFD = make_TensorDescription<ElementSFD, LayoutSFD>(128);
block_scaled_desc.EpilogueSFVecSize = SFD_VectorSize;
this->description_.block_scales = block_scaled_desc;
}
~GroupedBlockScaledGemmUniversal3xOperation() override = default;
mutable CudaBuffer layout_SFA_device;
mutable CudaBuffer layout_SFB_device;
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, 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 GroupedGemmOperation3xBase<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_SFA =
static_cast<const typename Operator::GemmKernel::ElementSF**>(arguments->SFA);
operator_args.mainloop.ptr_SFB =
static_cast<const typename Operator::GemmKernel::ElementSF**>(arguments->SFB);
operator_args.mainloop.layout_SFA =
static_cast<typename CollectiveMainloop::InternalLayoutSFA*>(this->layout_SFA_device.data());
operator_args.mainloop.layout_SFB =
static_cast<typename CollectiveMainloop::InternalLayoutSFB*>(this->layout_SFB_device.data());
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 lda, ldb, 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 {
auto const& config = *static_cast<GemmGroupedConfiguration const*>(configuration_ptr);
auto status = this->initialize_strides(config);
if (status != Status::kSuccess) {
return status;
}
auto num_groups = config.problem_count;
this->layout_SFA_device =
CudaBuffer(sizeof(typename CollectiveMainloop::InternalLayoutSFA) * num_groups);
this->layout_SFB_device =
CudaBuffer(sizeof(typename CollectiveMainloop::InternalLayoutSFB) * num_groups);
auto layout_SFA_host = std::vector<typename CollectiveMainloop::InternalLayoutSFA>(num_groups);
auto layout_SFB_host = std::vector<typename CollectiveMainloop::InternalLayoutSFB>(num_groups);
for (int group_idx = 0; group_idx < num_groups; group_idx++) {
auto const& shape = config.problem_sizes_3x_host[group_idx];
auto M = get<0>(shape);
auto N = get<1>(shape);
auto K = get<2>(shape);
auto layout_SFA = CollectiveMainloop::Sm100BlkScaledConfig::tile_atom_to_shape_SFA(cute::make_shape(M, N, K, 1));
auto layout_SFB = CollectiveMainloop::Sm100BlkScaledConfig::tile_atom_to_shape_SFB(cute::make_shape(M, N, K, 1));
layout_SFA_host[group_idx] = layout_SFA;
layout_SFB_host[group_idx] = layout_SFB;
}
CUDA_CHECK(cudaMemcpy(
this->layout_SFA_device.data(),
layout_SFA_host.data(),
sizeof(typename CollectiveMainloop::InternalLayoutSFA) * num_groups,
cudaMemcpyHostToDevice));
CUDA_CHECK(cudaMemcpy(
this->layout_SFB_device.data(),
layout_SFB_host.data(),
sizeof(typename CollectiveMainloop::InternalLayoutSFB) * num_groups,
cudaMemcpyHostToDevice));
Operator* op = new (host_workspace) Operator;
return status;
}
/// **** 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<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
///////////////////////////////////////////////////////////////////////////////////////////////////

87
tools/profiler/include/cutlass/profiler/grouped_gemm_operation_profiler.h
View File

@@ -85,7 +85,7 @@ public:
/// Parses the problem
Status parse(
library::GemmDescription const& operation_desc,
library::GroupedGemmDescription const& operation_desc,
ProblemSpace const& problem_space,
ProblemSpace::Problem const& problem);
@@ -94,27 +94,50 @@ public:
int64_t k(int group_idx) const { return problem_sizes[group_idx].k(); };
/// Total number of bytes loaded
int64_t bytes(library::GemmDescription const& operation_desc) const;
int64_t bytes(library::GroupedGemmDescription const& operation_desc) const;
/// Total number of flops computed
int64_t flops(library::GemmDescription const& operation_desc) const;
int64_t flops(library::GroupedGemmDescription const& operation_desc) const;
/// Initializes a performance result
void initialize_result(
PerformanceResult& result,
library::GemmDescription const& operation_desc,
library::GroupedGemmDescription const& operation_desc,
ProblemSpace const& problem_space);
};
struct BlockScalingWorkspace {
// host vector (per L2 workspace) of device vectors (per group) of device pointers
std::vector<DeviceAllocation*> SFA_ptr_array_device;
std::vector<DeviceAllocation*> SFB_ptr_array_device;
std::vector<DeviceAllocation*> SFC_ptr_array_device;
std::vector<DeviceAllocation*> SFD_ptr_array_device;
// host vector (per group) of device tensors
// (where each batch of device allocation is for a L2 workspace)
std::vector<DeviceAllocation*> SFA_ptr_array_host;
std::vector<DeviceAllocation*> SFB_ptr_array_host;
std::vector<DeviceAllocation*> SFC_ptr_array_host;
std::vector<DeviceAllocation*> SFD_ptr_array_host;
std::vector<DeviceAllocation*> SFD_reference_ptr_array_host;
// matrix wide constant, not per-batch or per-group
DeviceAllocation* norm_constant;
};
// workspace contains the allocated blocks, arguments just contain the raw
// pointers
struct GroupedGemmWorkspace {
// host vector (per L2 workspace) of device vectors (per group) of device pointers
std::vector<DeviceAllocation*> A_ptr_array_device;
std::vector<DeviceAllocation*> B_ptr_array_device;
std::vector<DeviceAllocation*> C_ptr_array_device;
std::vector<DeviceAllocation*> D_ptr_array_device;
std::vector<DeviceAllocation*> reference_ptr_array_host;
// host vector (per group) of device tensors
// (where each batch of device allocation is for a L2 workspace)
std::vector<DeviceAllocation*> A_ptr_array_host;
std::vector<DeviceAllocation*> B_ptr_array_host;
std::vector<DeviceAllocation*> C_ptr_array_host;
@@ -122,7 +145,7 @@ public:
/// Number of copies of the problem workspace which are visited sequentially during
/// profiling to avoid camping in the last level cache.
/// *NOT* the number of groups in the grouped GEMM
/// *NOT* the number of groups in the grouped GEMM (we use `num_groups` in the profiler)
int problem_count{1};
DeviceAllocation* problem_sizes_array_device{nullptr};
@@ -132,8 +155,10 @@ public:
DeviceAllocation* ldc_array_device{nullptr};
DeviceAllocation* ldd_array_device{nullptr};
std::optional<BlockScalingWorkspace> block_scales;
library::GemmGroupedConfiguration configuration;
library::GemmGroupedArguments arguments;
library::GroupedGemmBlockScaledArguments arguments;
std::vector<uint8_t> host_workspace;
DeviceAllocation device_workspace;
@@ -141,28 +166,38 @@ public:
private:
void init_arguments(Options const& options) {
gemm_workspace_.arguments.ptr_A = gemm_workspace_.A_ptr_array_device[0]->data();
gemm_workspace_.arguments.ptr_B = gemm_workspace_.B_ptr_array_device[0]->data();
gemm_workspace_.arguments.ptr_C = gemm_workspace_.C_ptr_array_device[0]->data();
gemm_workspace_.arguments.ptr_D = gemm_workspace_.D_ptr_array_device[0]->data();
gemm_workspace_.arguments.alpha = problem_.alpha.data();
gemm_workspace_.arguments.beta = problem_.beta.data();
gemm_workspace_.arguments.pointer_mode = library::ScalarPointerMode::kHost;
gemm_workspace_.arguments.lda = static_cast<int64_t*>(gemm_workspace_.lda_array_device->data());
gemm_workspace_.arguments.ldb = static_cast<int64_t*>(gemm_workspace_.ldb_array_device->data());
gemm_workspace_.arguments.ldc = static_cast<int64_t*>(gemm_workspace_.ldc_array_device->data());
gemm_workspace_.arguments.ldd = static_cast<int64_t*>(gemm_workspace_.ldc_array_device->data());
gemm_workspace_.arguments.problem_sizes =
auto& arguments = gemm_workspace_.arguments;
// these get updated in each profiler run to ensure L2 cycling
arguments.ptr_A = gemm_workspace_.A_ptr_array_device[0]->data();
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();
arguments.alpha = problem_.alpha.data();
arguments.beta = problem_.beta.data();
arguments.pointer_mode = library::ScalarPointerMode::kHost;
arguments.lda = static_cast<int64_t*>(gemm_workspace_.lda_array_device->data());
arguments.ldb = static_cast<int64_t*>(gemm_workspace_.ldb_array_device->data());
arguments.ldc = static_cast<int64_t*>(gemm_workspace_.ldc_array_device->data());
arguments.ldd = static_cast<int64_t*>(gemm_workspace_.ldc_array_device->data());
arguments.problem_sizes =
static_cast<gemm::GemmCoord*>(gemm_workspace_.problem_sizes_array_device->data());
gemm_workspace_.arguments.problem_sizes_3x = static_cast<cute::Shape<int, int, int>*>(
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)};
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)};
/* Query device SM count to pass onto the kernel as an argument, where needed */
gemm_workspace_.arguments.sm_count = options.device.properties[0].multiProcessorCount;
arguments.sm_count = options.device.properties[0].multiProcessorCount;
if (is_block_scaled) {
auto& block_scaled_ws = gemm_workspace_.block_scales.value();
arguments.SFA = block_scaled_ws.SFA_ptr_array_device[0]->data();
arguments.SFB = block_scaled_ws.SFB_ptr_array_device[0]->data();
arguments.SFD = block_scaled_ws.SFD_ptr_array_device[0]->data();
arguments.norm_constant = block_scaled_ws.norm_constant->data();
}
}
protected:
@@ -172,6 +207,8 @@ protected:
/// Device memory allocations
GroupedGemmWorkspace gemm_workspace_;
bool is_block_scaled{false};
public:
GroupedGemmOperationProfiler(Options const& options);
@@ -226,7 +263,7 @@ protected:
void initialize_result_(
PerformanceResult& result,
Options const& options,
library::GemmDescription const& operation_desc,
library::GroupedGemmDescription const& operation_desc,
ProblemSpace const& problem_space);
/// Verifies CUTLASS against host and device references
@@ -249,10 +286,6 @@ protected:
void* host_workspace,
void* device_workspace) override;
/// Initialize reduction problem dimensions and library::Operation
bool initialize_reduction_configuration_(
library::Operation const* operation,
ProblemSpace::Problem const& problem);
};
/////////////////////////////////////////////////////////////////////////////////////////////////

502
tools/profiler/src/grouped_gemm_operation_profiler.cu
View File

@@ -35,6 +35,8 @@
#include <bitset>
#include <cstdint>
#include <iostream>
#include <memory>
#include <optional>
#include <stdexcept>
#include <string>
#include <vector>
@@ -45,6 +47,8 @@
#include "cutlass/profiler/grouped_gemm_operation_profiler.h"
#include "cutlass/library/handle.h"
#include "cutlass/library/library.h"
#include "cutlass/library/operation_table.h"
#include "cutlass/library/singleton.h"
/////////////////////////////////////////////////////////////////////////////////////////////////
namespace {
@@ -161,7 +165,7 @@ void GroupedGemmOperationProfiler::print_examples(std::ostream& out) const {
}
Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse(
library::GemmDescription const& operation_desc,
library::GroupedGemmDescription const& operation_desc,
ProblemSpace const& problem_space,
ProblemSpace::Problem const& problem) {
@@ -242,7 +246,8 @@ 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);
} else {
}
else {
throw std::runtime_error(
"Invalid format in line: " + line + ". Each line in file expected to be 'mxnxk'.");
}
@@ -281,37 +286,42 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse(
this->mode = library::GemmUniversalMode::kGrouped;
if (!tensor_description_satisfies(operation_desc.A, "A", problem_space, problem)) {
if (!tensor_description_satisfies(operation_desc.gemm.A, "A", problem_space, problem)) {
return Status::kErrorInvalidProblem;
}
if (!tensor_description_satisfies(operation_desc.B, "B", problem_space, problem)) {
if (!tensor_description_satisfies(operation_desc.gemm.B, "B", problem_space, problem)) {
return Status::kErrorInvalidProblem;
}
if (!tensor_description_satisfies(operation_desc.C, "C", problem_space, problem)) {
if (!tensor_description_satisfies(operation_desc.gemm.C, "C", problem_space, problem)) {
return Status::kErrorInvalidProblem;
}
if (!tensor_description_satisfies(operation_desc.D, "D", problem_space, problem)) {
if (!tensor_description_satisfies(operation_desc.gemm.D, "D", problem_space, problem)) {
return Status::kErrorInvalidProblem;
}
if (!arg_as_scalar(
this->alpha,
operation_desc.element_epilogue,
operation_desc.gemm.element_epilogue,
"alpha",
problem_space,
problem)) {
if (!cast_from_double(this->alpha, operation_desc.element_epilogue, 1)) {
if (!cast_from_double(this->alpha, operation_desc.gemm.element_epilogue, 1)) {
return Status::kErrorInternal;
}
}
if (!arg_as_scalar(this->beta, operation_desc.element_epilogue, "beta", problem_space, problem)) {
if (!arg_as_scalar(
this->beta,
operation_desc.gemm.element_epilogue,
"beta",
problem_space,
problem)) {
if (!cast_from_double(this->beta, operation_desc.element_epilogue, 0)) {
if (!cast_from_double(this->beta, operation_desc.gemm.element_epilogue, 0)) {
return Status::kErrorInternal;
}
}
@@ -322,17 +332,17 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse(
this->ldc.resize(num_groups);
for (size_t group_idx = 0; group_idx < num_groups; group_idx++) {
this->lda[group_idx] = DeviceAllocation::get_packed_layout(
operation_desc.A.layout,
operation_desc.gemm.A.layout,
{int(this->m(group_idx)), int(this->k(group_idx))})
.front();
this->ldb[group_idx] = DeviceAllocation::get_packed_layout(
operation_desc.B.layout,
operation_desc.gemm.B.layout,
{int(this->k(group_idx)), int(this->n(group_idx))})
.front();
this->ldc[group_idx] = DeviceAllocation::get_packed_layout(
operation_desc.C.layout,
operation_desc.gemm.C.layout,
{int(this->m(group_idx)), int(this->n(group_idx))})
.front();
}
@@ -342,23 +352,23 @@ Status GroupedGemmOperationProfiler::GroupedGemmProblem::parse(
/// Total number of bytes loaded
int64_t GroupedGemmOperationProfiler::GroupedGemmProblem::bytes(
library::GemmDescription const& operation_desc) const {
library::GroupedGemmDescription const& operation_desc) const {
// Input bytes read and Output bytes written for the gemm problem
int64_t bytes = 0;
for (size_t group_idx = 0, num_groups = problem_sizes.size(); group_idx < num_groups;
group_idx++) {
bytes +=
int64_t(library::sizeof_bits(operation_desc.A.element) * m(group_idx) / 8) * k(group_idx) +
int64_t(library::sizeof_bits(operation_desc.B.element) * n(group_idx) / 8) * k(group_idx) +
int64_t(library::sizeof_bits(operation_desc.C.element) * m(group_idx) / 8) * n(group_idx);
int64_t(library::sizeof_bits(operation_desc.gemm.A.element) * m(group_idx) / 8) * k(group_idx) +
int64_t(library::sizeof_bits(operation_desc.gemm.B.element) * n(group_idx) / 8) * k(group_idx) +
int64_t(library::sizeof_bits(operation_desc.gemm.C.element) * m(group_idx) / 8) * n(group_idx);
// Set is_beta_zero true if beta is zero
bool is_beta_zero = std::all_of(beta.begin(), beta.end(), [](uint8_t i) { return i == 0; });
// Output bytes read for the gemm problem for non-zero beta values
if (!is_beta_zero) {
bytes +=
int64_t(library::sizeof_bits(operation_desc.C.element) * m(group_idx) / 8) * n(group_idx);
int64_t(library::sizeof_bits(operation_desc.gemm.C.element) * m(group_idx) / 8) * n(group_idx);
}
}
@@ -367,7 +377,7 @@ int64_t GroupedGemmOperationProfiler::GroupedGemmProblem::bytes(
/// Total number of flops computed
int64_t GroupedGemmOperationProfiler::GroupedGemmProblem::flops(
library::GemmDescription const& operation_desc) const {
library::GroupedGemmDescription const& operation_desc) const {
int64_t flops_ = 0;
for (size_t group_idx = 0, num_groups = problem_sizes.size(); group_idx < num_groups;
group_idx++) {
@@ -376,7 +386,7 @@ int64_t GroupedGemmOperationProfiler::GroupedGemmProblem::flops(
}
// complex-valued support
switch (operation_desc.tile_description.math_instruction.math_operation) {
switch (operation_desc.gemm.tile_description.math_instruction.math_operation) {
case library::MathOperationID::kMultiplyAddComplex:
case library::MathOperationID::kMultiplyAddComplexFastF32:
flops_ *= 4;
@@ -395,40 +405,44 @@ int64_t GroupedGemmOperationProfiler::GroupedGemmProblem::flops(
/// Initializes a performance result
void GroupedGemmOperationProfiler::GroupedGemmProblem::initialize_result(
PerformanceResult& result,
library::GemmDescription const& operation_desc,
library::GroupedGemmDescription const& operation_desc,
ProblemSpace const& problem_space) {
result.arguments.resize(problem_space.rank());
set_argument(result, "gemm_kind", problem_space, library::to_string(operation_desc.gemm_kind));
set_argument(
result,
"gemm_kind",
problem_space,
library::to_string(operation_desc.gemm.gemm_kind));
set_argument(
result,
"A",
problem_space,
std::string(library::to_string(operation_desc.A.element)) + ":" +
library::to_string(operation_desc.A.layout));
std::string(library::to_string(operation_desc.gemm.A.element)) + ":" +
library::to_string(operation_desc.gemm.A.layout));
set_argument(
result,
"B",
problem_space,
std::string(library::to_string(operation_desc.B.element)) + ":" +
library::to_string(operation_desc.B.layout));
std::string(library::to_string(operation_desc.gemm.B.element)) + ":" +
library::to_string(operation_desc.gemm.B.layout));
set_argument(
result,
"C",
problem_space,
std::string(library::to_string(operation_desc.C.element)) + ":" +
library::to_string(operation_desc.C.layout));
std::string(library::to_string(operation_desc.gemm.C.element)) + ":" +
library::to_string(operation_desc.gemm.C.layout));
set_argument(
result,
"D",
problem_space,
std::string(library::to_string(operation_desc.D.element)) + ":" +
library::to_string(operation_desc.D.layout));
std::string(library::to_string(operation_desc.gemm.D.element)) + ":" +
library::to_string(operation_desc.gemm.D.layout));
{
std::stringstream ss;
@@ -456,13 +470,13 @@ void GroupedGemmOperationProfiler::GroupedGemmProblem::initialize_result(
result,
"alpha",
problem_space,
library::lexical_cast(alpha, operation_desc.element_epilogue));
library::lexical_cast(alpha, operation_desc.gemm.element_epilogue));
set_argument(
result,
"beta",
problem_space,
library::lexical_cast(beta, operation_desc.element_epilogue));
library::lexical_cast(beta, operation_desc.gemm.element_epilogue));
}
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -476,10 +490,23 @@ Status GroupedGemmOperationProfiler::initialize_configuration(
ProblemSpace const& problem_space,
ProblemSpace::Problem const& problem) {
library::GemmDescription const& operation_desc =
static_cast<library::GemmDescription const&>(operation->description());
library::GroupedGemmDescription const& operation_desc =
static_cast<library::GroupedGemmDescription const&>(operation->description());
if (operation_desc.gemm_kind != library::GemmKind::kGrouped) {
// We want to share the same operation profiler for any grouped gemm operation.
// We distinguish between block scaled and non-block scaled operations by looking at the kernel
// name, which tells us what reference kernel to use, which arguments to pass to the operation
// etc. This avoids creating yet another OperationProfiler with a lot of boilerplate in it.
if (std::string(operation_desc.gemm.name).find("bstensor") != std::string::npos) {
is_block_scaled = true;
gemm_workspace_.block_scales = BlockScalingWorkspace{};
}
else {
is_block_scaled = false;
gemm_workspace_.block_scales = std::nullopt;
}
if (operation_desc.gemm.gemm_kind != library::GemmKind::kGrouped) {
return Status::kErrorInvalidProblem;
}
@@ -489,10 +516,12 @@ Status GroupedGemmOperationProfiler::initialize_configuration(
}
auto num_groups = problem_.problem_sizes.size();
gemm_workspace_.configuration.problem_count = num_groups;
gemm_workspace_.configuration.lda = problem_.lda.data();
gemm_workspace_.configuration.ldb = problem_.ldb.data();
gemm_workspace_.configuration.ldc = problem_.ldc.data();
auto& config = gemm_workspace_.configuration;
config.problem_count = num_groups;
config.lda = problem_.lda.data();
config.ldb = problem_.ldb.data();
config.ldc = problem_.ldc.data();
config.problem_sizes_3x_host = problem_.problem_sizes_3x.data();
initialize_result_(this->model_result_, options, operation_desc, problem_space);
@@ -503,13 +532,13 @@ Status GroupedGemmOperationProfiler::initialize_configuration(
void GroupedGemmOperationProfiler::initialize_result_(
PerformanceResult& result,
Options const& options,
library::GemmDescription const& operation_desc,
library::GroupedGemmDescription const& operation_desc,
ProblemSpace const& problem_space) {
result.provider = library::Provider::kCUTLASS;
result.disposition = Disposition::kNotRun;
result.status = Status::kSuccess;
result.operation_name = operation_desc.name;
result.operation_name = operation_desc.gemm.name;
problem_.initialize_result(result, operation_desc, problem_space);
@@ -542,8 +571,8 @@ Status GroupedGemmOperationProfiler::initialize_workspace(
}
library::Operation const* underlying_operation = operation;
library::GemmDescription const& operation_desc =
static_cast<library::GemmDescription const&>(operation->description());
library::GroupedGemmDescription const& operation_desc =
static_cast<library::GroupedGemmDescription const&>(operation->description());
// Compute the number of copies of the problem to avoid L2 camping.
if (!options.profiling.workspace_count) {
@@ -568,6 +597,14 @@ Status GroupedGemmOperationProfiler::initialize_workspace(
gemm_workspace_.B_ptr_array_host.resize(num_groups);
gemm_workspace_.C_ptr_array_host.resize(num_groups);
gemm_workspace_.D_ptr_array_host.resize(num_groups);
if (is_block_scaled) {
auto& block_scaling_ws = gemm_workspace_.block_scales.value();
block_scaling_ws.SFA_ptr_array_host.resize(num_groups);
block_scaling_ws.SFB_ptr_array_host.resize(num_groups);
block_scaling_ws.SFC_ptr_array_host.resize(num_groups);
block_scaling_ws.SFD_ptr_array_host.resize(num_groups);
block_scaling_ws.SFD_reference_ptr_array_host.resize(num_groups);
}
static_assert(sizeof(void*) == 8); // allocating blocks for pointers, so verify pointer size
// ldx
gemm_workspace_.lda_array_device =
@@ -608,8 +645,8 @@ Status GroupedGemmOperationProfiler::initialize_workspace(
gemm_workspace_.A_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor(
options,
"A_" + group_str,
operation_desc.A.element,
operation_desc.A.layout,
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,
@@ -618,8 +655,8 @@ Status GroupedGemmOperationProfiler::initialize_workspace(
gemm_workspace_.B_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor(
options,
"B_" + group_str,
operation_desc.B.element,
operation_desc.B.layout,
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,
@@ -628,8 +665,8 @@ Status GroupedGemmOperationProfiler::initialize_workspace(
gemm_workspace_.C_ptr_array_host[group_idx] = device_context.allocate_and_initialize_tensor(
options,
"C_" + group_str,
operation_desc.C.element,
operation_desc.C.layout,
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,
@@ -638,8 +675,8 @@ Status GroupedGemmOperationProfiler::initialize_workspace(
gemm_workspace_.D_ptr_array_host[group_idx] = device_context.allocate_tensor(
options,
"D_" + group_str,
operation_desc.D.element,
operation_desc.D.layout,
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,
@@ -648,12 +685,81 @@ Status GroupedGemmOperationProfiler::initialize_workspace(
gemm_workspace_.reference_ptr_array_host[group_idx] = device_context.allocate_tensor(
options,
"Reference_" + group_str,
operation_desc.D.element,
operation_desc.D.layout,
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,
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.SFVecSize), 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);
}
}
}
// takes the allocated tensors and initializes an array of pointers per problem in the workspace
@@ -691,8 +797,36 @@ Status GroupedGemmOperationProfiler::initialize_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
);
}
init_arguments(options);
}
init_arguments(options);
//
// Initialize the CUTLASS operation
@@ -769,7 +903,6 @@ bool GroupedGemmOperationProfiler::verify_cutlass(
if (results_.back().status != Status::kSuccess) {
results_.back().disposition = Disposition::kFailed;
throw "failed";
return false;
}
@@ -795,8 +928,8 @@ bool GroupedGemmOperationProfiler::verify_cutlass(
}
#endif // #if CUTLASS_ENABLE_CUBLAS
library::GemmDescription const& gemm_desc =
static_cast<library::GemmDescription const&>(operation->description());
auto const& desc =
static_cast<library::GroupedGemmDescription const&>(operation->description());
bool verification_status = verify_with_reference_(
options,
@@ -805,8 +938,8 @@ bool GroupedGemmOperationProfiler::verify_cutlass(
operation,
problem_space,
problem,
gemm_desc.A.element,
gemm_desc.B.element);
desc.gemm.A.element,
desc.gemm.B.element);
// Update disposition to worst case verification outcome among all
// verification providers which are supported
@@ -854,8 +987,8 @@ bool GroupedGemmOperationProfiler::verify_with_reference_(
ProblemSpace::Problem const& problem,
cutlass::library::NumericTypeID element_A,
cutlass::library::NumericTypeID element_B) {
library::GemmDescription const& gemm_desc =
static_cast<library::GemmDescription const&>(operation->description());
library::GroupedGemmDescription const& desc =
static_cast<library::GroupedGemmDescription const&>(operation->description());
for (auto provider : options.verification.providers) {
@@ -864,8 +997,15 @@ bool GroupedGemmOperationProfiler::verify_with_reference_(
continue;
}
// we only have a block scaled reference kernel implemented on the host
if (is_block_scaled && provider != library::Provider::kReferenceHost) {
continue;
}
auto status = Status::kSuccess;
auto disposition = Disposition::kFailed;
// we don't have grouped GEMM reference kernels so we loop over the groups and perform
// a regular GEMM for each group
for (size_t group_idx = 0, num_groups = problem_.problem_sizes.size(); group_idx < num_groups;
group_idx++) {
void* ptr_A = gemm_workspace_.A_ptr_array_host[group_idx]->data();
@@ -879,6 +1019,16 @@ bool GroupedGemmOperationProfiler::verify_with_reference_(
std::vector<uint8_t> host_data_B;
std::vector<uint8_t> host_data_C;
std::vector<uint8_t> host_data_D;
std::vector<uint8_t> host_data_SFA;
std::vector<uint8_t> host_data_SFB;
std::vector<uint8_t> host_data_SFC;
std::vector<uint8_t> host_data_SFD;
std::vector<uint8_t> host_data_norm_constant;
void* ptr_SFA{nullptr};
void* ptr_SFB{nullptr};
void* ptr_SFD{nullptr};
void* ptr_norm_constant{nullptr};
if (provider == library::Provider::kReferenceHost) {
host_data_A.resize(gemm_workspace_.A_ptr_array_host[group_idx]->bytes());
@@ -896,52 +1046,171 @@ bool GroupedGemmOperationProfiler::verify_with_reference_(
host_data_D.resize(gemm_workspace_.reference_ptr_array_host[group_idx]->bytes());
ptr_D = host_data_D.data();
if (is_block_scaled) {
auto const& ws = gemm_workspace_.block_scales.value();
host_data_SFA.resize(ws.SFA_ptr_array_host[group_idx]->bytes());
ptr_SFA = host_data_SFA.data();
ws.SFA_ptr_array_host[group_idx]->copy_to_host(ptr_SFA);
host_data_SFB.resize(ws.SFB_ptr_array_host[group_idx]->bytes());
ptr_SFB = host_data_SFB.data();
ws.SFB_ptr_array_host[group_idx]->copy_to_host(ptr_SFB);
host_data_SFD.resize(ws.SFD_reference_ptr_array_host[group_idx]->bytes());
ptr_SFD = host_data_SFD.data();
host_data_norm_constant.resize(ws.norm_constant->bytes());
ptr_norm_constant = host_data_norm_constant.data();
ws.norm_constant->copy_to_host(ptr_norm_constant);
}
}
library::Handle handle;
handle.set_provider(provider);
const auto &desc = static_cast<library::GroupedGemmDescription const &>(operation->description());
const auto& gemm_desc = desc.gemm;
status = handle.gemm_universal(
library::GemmUniversalMode::kGemm,
problem_.m(group_idx),
problem_.n(group_idx),
problem_.k(group_idx),
problem_.cluster_m,
problem_.cluster_n,
problem_.cluster_k,
problem_.cluster_m_fallback,
problem_.cluster_n_fallback,
problem_.cluster_k_fallback,
gemm_desc.tile_description.math_instruction.element_accumulator,
gemm_desc.element_epilogue,
problem_.alpha.data(),
element_A,
gemm_desc.A.layout,
gemm_desc.transform_A,
ptr_A,
int(problem_.lda[group_idx]),
element_B,
gemm_desc.B.layout,
gemm_desc.transform_B,
ptr_B,
int(problem_.ldb[group_idx]),
problem_.beta.data(),
gemm_desc.C.element,
gemm_desc.C.layout,
ptr_C,
int(problem_.ldc[group_idx]),
gemm_desc.D.element,
gemm_desc.D.layout,
ptr_D,
int(problem_.ldc[group_idx]),
1,
gemm_workspace_.A_ptr_array_host[group_idx]->batch_stride(),
gemm_workspace_.B_ptr_array_host[group_idx]->batch_stride(),
gemm_workspace_.C_ptr_array_host[group_idx]->batch_stride(),
gemm_workspace_.reference_ptr_array_host[group_idx]->batch_stride());
if (!is_block_scaled) {
library::Handle handle;
handle.set_provider(provider);
if (status != Status::kSuccess)
status = handle.gemm_universal(
library::GemmUniversalMode::kGemm,
problem_.m(group_idx),
problem_.n(group_idx),
problem_.k(group_idx),
problem_.cluster_m,
problem_.cluster_n,
problem_.cluster_k,
problem_.cluster_m_fallback,
problem_.cluster_n_fallback,
problem_.cluster_k_fallback,
desc.gemm.tile_description.math_instruction.element_accumulator,
desc.gemm.element_epilogue,
problem_.alpha.data(),
element_A,
desc.gemm.A.layout,
desc.gemm.transform_A,
ptr_A,
int(problem_.lda[group_idx]),
element_B,
desc.gemm.B.layout,
desc.gemm.transform_B,
ptr_B,
int(problem_.ldb[group_idx]),
problem_.beta.data(),
desc.gemm.C.element,
desc.gemm.C.layout,
ptr_C,
int(problem_.ldc[group_idx]),
desc.gemm.D.element,
desc.gemm.D.layout,
ptr_D,
int(problem_.ldc[group_idx]),
1,
gemm_workspace_.A_ptr_array_host[group_idx]->batch_stride(),
gemm_workspace_.B_ptr_array_host[group_idx]->batch_stride(),
gemm_workspace_.C_ptr_array_host[group_idx]->batch_stride(),
gemm_workspace_.reference_ptr_array_host[group_idx]->batch_stride());
}
else {
auto const& block_scale_desc = desc.block_scales.value();
auto& block_scale_ws = gemm_workspace_.block_scales.value();
library::BlockScaledGemmFunctionalKey blockScaledGemm_key(
library::Provider::kReferenceHost,
library::GemmKind::kUniversal,
library::OperationKind::kBlockScaledGemm,
gemm_desc.tile_description.math_instruction.element_accumulator,
gemm_desc.element_epilogue,
element_A,
gemm_desc.A.layout,
block_scale_desc.SFA.element,
element_B,
gemm_desc.B.layout,
block_scale_desc.SFB.element,
gemm_desc.C.element,
gemm_desc.C.layout,
gemm_desc.D.element,
gemm_desc.D.layout,
block_scale_desc.SFD.element,
block_scale_desc.SFD.layout,
block_scale_desc.SFVecSize,
block_scale_desc.EpilogueSFVecSize);
auto operators_it =
library::Singleton::get().operation_table.block_scaled_gemm_operations.find(
blockScaledGemm_key);
if (
operators_it ==
library::Singleton::get().operation_table.block_scaled_gemm_operations.end()) {
disposition = Disposition::kNotSupported;
break;
}
if (operators_it->second.empty()) {
disposition = Disposition::kNotSupported;
break;
}
auto cc_it = operators_it->second.begin();
if (cc_it == operators_it->second.end()) {
disposition = Disposition::kNotSupported;
break;
}
// host reference has only one instances in BlockScaledOperationVectorMap
library::Operation const* reference_op = cc_it->second[0];
library::BlockScaledGemmArguments arguments{
{int(problem_.m(group_idx)), int(problem_.n(group_idx)), int(problem_.k(group_idx))},
{int(problem_.cluster_m), int(problem_.cluster_n), int(problem_.cluster_k)},
{int(problem_.cluster_m_fallback), int(problem_.cluster_n_fallback), int(problem_.cluster_k_fallback)},
1, // batch count
ptr_A,
ptr_B,
ptr_SFA,
ptr_SFB,
ptr_C,
ptr_D,
ptr_SFD,
problem_.alpha.data(),
problem_.beta.data(),
library::ScalarPointerMode::kHost,
problem_.lda[group_idx],
problem_.ldb[group_idx],
problem_.ldc[group_idx],
problem_.ldc[group_idx],
gemm_workspace_.A_ptr_array_host[group_idx]->batch_stride(),
gemm_workspace_.B_ptr_array_host[group_idx]->batch_stride(),
gemm_workspace_.C_ptr_array_host[group_idx]->batch_stride(),
gemm_workspace_.reference_ptr_array_host[group_idx]->batch_stride(),
ptr_norm_constant};
library::GemmUniversalConfiguration configuration{
library::GemmUniversalMode::kGemm,
problem_.problem_sizes[group_idx],
{problem_.cluster_m, problem_.cluster_n, problem_.cluster_k},
{problem_.cluster_m_fallback, problem_.cluster_n_fallback, problem_.cluster_k_fallback},
1,
problem_.lda[group_idx],
problem_.ldb[group_idx],
problem_.ldc[group_idx],
problem_.ldc[group_idx],
1,
};
uint64_t host_workspace_size_needed = reference_op->get_host_workspace_size(&gemm_workspace_.configuration);
std::vector<char> host_workspace(host_workspace_size_needed);
status = reference_op->initialize(&configuration, host_workspace.data());
if (status != Status::kSuccess) {
break;
}
status = reference_op->run(&arguments, host_workspace.data());
block_scale_ws.SFD_reference_ptr_array_host[group_idx]->copy_from_host(ptr_SFD);
}
if (status != Status::kSuccess) {
break;
}
if (provider == library::Provider::kReferenceHost) {
gemm_workspace_.reference_ptr_array_host[group_idx]->copy_from_host(ptr_D);
@@ -952,26 +1221,40 @@ bool GroupedGemmOperationProfiler::verify_with_reference_(
*gemm_workspace_.D_ptr_array_host[group_idx],
*gemm_workspace_.reference_ptr_array_host[group_idx],
gemm_workspace_.D_ptr_array_host[group_idx]->batch_stride());
if (disposition != Disposition::kPassed)
if (disposition != Disposition::kPassed) {
break;
}
if (is_block_scaled) {
auto& ws = gemm_workspace_.block_scales.value();
auto const& block_scale_desc = desc.block_scales.value();
if (block_scale_desc.SFD.element != library::NumericTypeID::kVoid) {
disposition = compare_tensors(
options,
*ws.SFD_ptr_array_host[group_idx],
*ws.SFD_reference_ptr_array_host[group_idx],
ws.SFD_ptr_array_host[group_idx]->batch_stride());
if (disposition != Disposition::kPassed) {
break;
}
}
}
}
if (status != Status::kSuccess) {
results_.back().verification_map[provider] = Disposition::kNotRun;
results_.back().verification_map[provider] = Disposition::kNotVerified;
continue;
}
results_.back().status = status;
results_.back().verification_map[provider] = disposition;
// Save workspace if incorrect
if (
options.verification.save_workspace == SaveWorkspace::kIncorrect &&
results_.back().verification_map[provider] == Disposition::kIncorrect) {
save_workspace(device_context, options, gemm_desc, library::Provider::kCUTLASS, provider);
save_workspace(device_context, options, desc, library::Provider::kCUTLASS, provider);
}
}
return true;
return true; // continue profiling
}
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -1008,7 +1291,6 @@ Status GroupedGemmOperationProfiler::profile_cutlass_(
void* host_workspace,
void* device_workspace) {
// initialize gemm underlying operation to handle parallel reduction
library::Operation const* underlying_operation = operation;
auto func = [&](cudaStream_t stream, int iteration) {

2
tools/util/include/cutlass/util/reference/device/tensor_compare.h
View File

@@ -138,7 +138,6 @@ bool BlockCompareEqual(
if (result != cudaSuccess) {
throw std::runtime_error("Failed to query occupancy.");
}
// Limit block size. This has the effect of increasing the number of items processed by a
// single thread and reduces the impact of initialization overhead.
block_size = (block_size < 128 ? block_size : 128);
@@ -205,7 +204,6 @@ bool BlockCompareRelativelyEqual(
if (result != cudaSuccess) {
throw std::runtime_error("Failed to query occupancy.");
}
// Limit block size. This has the effect of increasing the number of items processed by a
// single thread and reduces the impact of initialization overhead.
block_size = (block_size < 128 ? block_size : 128);

2
tools/util/include/cutlass/util/reference/device/tensor_foreach.h
View File

@@ -61,7 +61,6 @@ struct TensorForEach {
if (result != cudaSuccess) {
throw std::runtime_error("Failed to query occupancy.");
}
// Limit block size. This has the effect of increasing the number of items processed by a
// single thread and reduces the impact of initialization overhead.
block_size = (block_size < 128 ? block_size : 128);
@@ -124,7 +123,6 @@ struct BlockForEach {
if (result != cudaSuccess) {
throw std::runtime_error("Failed to query occupancy.");
}
// Limit block size. This has the effect of increasing the number of items processed by a
// single thread and reduces the impact of initialization overhead.
block_size = (block_size < 128 ? block_size : 128);