v4.1 release

This commit is contained in:
Junkai-Wu
2025-07-03 20:07:53 +08:00
committed by GitHub
parent b995f93317
commit a1aaf2300a
155 changed files with 18407 additions and 6068 deletions

View File

@@ -99,7 +99,7 @@ template <typename OperatorClass> struct ArchMap<arch::Sm86, OperatorClass> {
template <typename OperatorClass> struct ArchMap<arch::Sm89, OperatorClass> {
static int const kMin = 89;
static int const kMax = 89;
static int const kMax = 100;
};
template <typename OperatorClass> struct ArchMap<arch::Sm90, OperatorClass> {

View File

@@ -59,13 +59,7 @@ namespace library {
namespace detail {
template <typename T>
auto make_iterator(T* ptr) {
using namespace cute;
if constexpr (cute::is_subbyte_v<T>) {
return subbyte_iterator<T>(ptr);
}
else {
return ptr;
}
return cute::recast_ptr<T>(ptr);
}
}
@@ -73,7 +67,7 @@ auto make_iterator(T* ptr) {
template <
Provider Provider_,
typename ElementA_,
typename ElementA_,
typename LayoutA_,
typename ElementSFA_,
typename ElementB_,
@@ -125,7 +119,7 @@ public:
/// Constructor
BlockScaledGemmReferenceOperation() {
// Basic information
description_.provider = kProvider;
description_.kind = OperationKind::kBlockScaledGemm;
@@ -139,7 +133,7 @@ public:
description_.C = make_TensorDescription<ElementC, LayoutC>();
description_.D = make_TensorDescription<ElementD, LayoutC>();
description_.SFD = make_TensorDescription<ElementSFD, LayoutSFD>();
// Epilogue compute and accumulator type description
description_.element_epilogue = NumericTypeMap<ElementCompute>::kId;
@@ -147,7 +141,7 @@ public:
NumericTypeMap<ElementAccumulator>::kId;
// Compute capability for gemm reference
description_.tile_description.minimum_compute_capability =
description_.tile_description.minimum_compute_capability =
(kProvider == Provider::kReferenceDevice ? 50 : 0);
description_.tile_description.maximum_compute_capability = 1024;
@@ -158,7 +152,7 @@ public:
// Procedural name
std::stringstream ss;
ss << "gemm"
ss << "gemm"
<< "_reference_" << to_string(description_.provider)
<< "_" << to_string(description_.A.element) << to_string(description_.A.layout)
<< "_" << to_string(description_.SFA.element) << to_string(description_.SFA.layout)
@@ -221,7 +215,7 @@ public:
BlockScaledGemmArguments const &args = *static_cast<BlockScaledGemmArguments const *>(arguments);
// Construct cute::Tensor A/B/C
// Construct cute::Tensor A/B/C
int M = args.problem_size.m();
int N = args.problem_size.n();
@@ -266,12 +260,12 @@ public:
auto D = cute::make_tensor(detail::make_iterator(static_cast<ElementD *>(args.D)),
cute::make_layout(cute::make_shape(M, N, L), stride_d));
cutlass::reference::host::GettBlockScalingMainloopParams<ElementAccumulator,
decltype(A), decltype(SfA),
decltype(B), decltype(SfB)>
cutlass::reference::host::GettBlockScalingMainloopParams<ElementAccumulator,
decltype(A), decltype(SfA),
decltype(B), decltype(SfB)>
mainloop_params{A, SfA, B, SfB};
if constexpr (not is_same_v<ElementSFD, void>) {
if constexpr (not is_same_v<ElementSFD, void>) {
using Sm1xxBlockScaledOutputConfig= cutlass::detail::Sm1xxBlockScaledOutputConfig<
EpilogueSFVecSize
@@ -289,7 +283,7 @@ public:
else {
// W/O SF generation
auto SfD = cute::make_tensor(static_cast<ElementSFA *>(nullptr),
cute::make_layout(cute::make_shape(M, N, L))); // not used.
cute::make_layout(cute::make_shape(M, N, L))); // not used.
cutlass::reference::host::GettBlockScalingEpilogueParams<
ElementCompute, ElementAccumulator, ElementCompute,
decltype(C), decltype(D), decltype(SfD)>
@@ -362,7 +356,7 @@ template <
typename InnerProductOp_ = multiply_add<ElementAccumulator_>
>
void make_block_scaled_gemm(Manifest &manifest) {
///
///
/// A is Row , B is Col
///
manifest.append(new BlockScaledGemmReferenceOperation<
@@ -405,7 +399,7 @@ void make_block_scaled_gemm(Manifest &manifest) {
ConvertOp_,
InnerProductOp_
>);
///
///
/// A is Col , B is Row
///
manifest.append(new BlockScaledGemmReferenceOperation<

View File

@@ -105,8 +105,6 @@ public:
cutlass::library::RasterOrder raster_order{cutlass::library::RasterOrder::kHeuristic};
int swizzle_size{1};
cutlass::library::RuntimeDatatype runtime_input_datatype_a{};
cutlass::library::RuntimeDatatype runtime_input_datatype_b{};
@@ -264,7 +262,8 @@ protected:
std::array<int64_t, 3> const &preferred_cluster,
std::array<int64_t, 3> const &fallback_cluster,
cutlass::library::RasterOrder const &raster_order,
int swizzle_size);
int swizzle_size,
bool is_dynamic_cluster_enabled);
/// Update performance result configuration according to flexible user setups
void update_result_(
@@ -275,7 +274,8 @@ protected:
cutlass::library::RasterOrder const &raster_order,
std::array<int64_t, 3> const &preferred_cluster,
std::array<int64_t, 3> const &fallback_cluster,
int swizzle_size);
int swizzle_size,
bool is_dynamic_cluster_enabled);
/// Initializes the performance result
void initialize_result_(

View File

@@ -103,8 +103,6 @@ public:
cutlass::library::RasterOrder raster_order{cutlass::library::RasterOrder::kHeuristic};
int swizzle_size{1};
cutlass::library::RuntimeDatatype runtime_input_datatype_a{};
cutlass::library::RuntimeDatatype runtime_input_datatype_b{};
@@ -268,7 +266,8 @@ protected:
std::array<int64_t, 3> const &preferred_cluster,
std::array<int64_t, 3> const &fallback_cluster,
cutlass::library::RasterOrder const &raster_order,
int swizzle_size);
int swizzle_size,
bool is_dynamic_cluster_enabled);
/// Update performance result configuration according to flexible user setups
void update_result_(
@@ -279,7 +278,8 @@ protected:
cutlass::library::RasterOrder const &raster_order,
std::array<int64_t, 3> const &preferred_cluster,
std::array<int64_t, 3> const &fallback_cluster,
int swizzle_size);
int swizzle_size,
bool is_dynamic_cluster_enabled);
/// Initializes the performance result
void initialize_result_(

View File

@@ -286,14 +286,25 @@ protected:
library::GroupedGemmDescription const& operation_desc,
ProblemSpace const& problem_space);
/// Update workspace configuration according to flexible user setups
void update_workspace_(
GroupedGemmWorkspace &gemm_workspace,
std::array<int64_t, 3> const &preferred_cluster,
std::array<int64_t, 3> const &fallback_cluster,
cutlass::library::RasterOrder const &raster_order,
int swizzle_size,
bool is_dynamic_cluster_enabled);
/// Update performance result configuration for exploration parameters
void update_result_(
void update_workspace_and_result_(
GroupedGemmWorkspace &gemm_workspace,
PerformanceResult &result,
ProblemSpace const &problem_space,
cutlass::library::RasterOrder const &raster_order,
std::array<int64_t, 3> const &preferred_cluster,
std::array<int64_t, 3> const &fallback_cluster,
int swizzle_size);
int swizzle_size,
bool is_dynamic_cluster_enabled);
/// Verifies CUTLASS against host and device references
bool verify_with_reference_(

View File

@@ -559,7 +559,8 @@ void BlockScaledGemmOperationProfiler::update_workspace_(
std::array<int64_t, 3> const &preferred_cluster,
std::array<int64_t, 3> const &fallback_cluster,
cutlass::library::RasterOrder const &raster_order,
int swizzle_size
int swizzle_size,
bool is_dynamic_cluster_enabled
) {
gemm_workspace.arguments.problem_size.m() = problem_shape.m();
@@ -573,15 +574,17 @@ void BlockScaledGemmOperationProfiler::update_workspace_(
gemm_workspace.arguments.swizzle_size = swizzle_size;
gemm_workspace.arguments.raster_order = raster_order;
gemm_workspace.arguments.cluster_shape = {int(preferred_cluster[0]), int(preferred_cluster[1]), int(preferred_cluster[2])};
gemm_workspace.arguments.cluster_shape_fallback = {int(fallback_cluster[0]), int(fallback_cluster[1]), int(fallback_cluster[2])};
if (is_dynamic_cluster_enabled) {
gemm_workspace.arguments.cluster_shape = {int(preferred_cluster[0]), int(preferred_cluster[1]), int(preferred_cluster[2])};
gemm_workspace.arguments.cluster_shape_fallback = {int(fallback_cluster[0]), int(fallback_cluster[1]), int(fallback_cluster[2])};
gemm_workspace.configuration.cluster_shape = {int(preferred_cluster[0]), int(preferred_cluster[1]), int(preferred_cluster[2])};
gemm_workspace.configuration.cluster_shape_fallback = {int(fallback_cluster[0]), int(fallback_cluster[1]), int(fallback_cluster[2])};
}
gemm_workspace.configuration.problem_size.m() = problem_shape.m();
gemm_workspace.configuration.problem_size.n() = problem_shape.n();
gemm_workspace.configuration.problem_size.k() = problem_shape.k();
gemm_workspace.configuration.cluster_shape = {int(preferred_cluster[0]), int(preferred_cluster[1]), int(preferred_cluster[2])};
gemm_workspace.configuration.cluster_shape_fallback = {int(fallback_cluster[0]), int(fallback_cluster[1]), int(fallback_cluster[2])};
gemm_workspace.configuration.lda = leading_dim[0];
gemm_workspace.configuration.ldb = leading_dim[1];
@@ -598,7 +601,8 @@ void BlockScaledGemmOperationProfiler::update_result_(
cutlass::library::RasterOrder const &raster_order,
std::array<int64_t, 3> const &preferred_cluster,
std::array<int64_t, 3> const &fallback_cluster,
int swizzle_size
int swizzle_size,
bool is_dynamic_cluster_enabled
) {
result.bytes = problem_.bytes_with_problem_shape(operation_desc, problem_shape);
result.flops = problem_.flops_with_problem_shape(operation_desc, problem_shape);
@@ -609,12 +613,14 @@ void BlockScaledGemmOperationProfiler::update_result_(
set_argument(result, "raster_order", problem_space, library::to_string(raster_order));
set_argument(result, "swizzle_size", problem_space, swizzle_size);
set_argument(result, "cluster_m", problem_space, preferred_cluster[0]);
set_argument(result, "cluster_n", problem_space, preferred_cluster[1]);
set_argument(result, "cluster_k", problem_space, preferred_cluster[2]);
set_argument(result, "cluster_m_fallback", problem_space, fallback_cluster[0]);
set_argument(result, "cluster_n_fallback", problem_space, fallback_cluster[1]);
set_argument(result, "cluster_k_fallback", problem_space, fallback_cluster[2]);
if (is_dynamic_cluster_enabled) {
set_argument(result, "cluster_m", problem_space, preferred_cluster[0]);
set_argument(result, "cluster_n", problem_space, preferred_cluster[1]);
set_argument(result, "cluster_k", problem_space, preferred_cluster[2]);
set_argument(result, "cluster_m_fallback", problem_space, fallback_cluster[0]);
set_argument(result, "cluster_n_fallback", problem_space, fallback_cluster[1]);
set_argument(result, "cluster_k_fallback", problem_space, fallback_cluster[2]);
}
}
@@ -1377,9 +1383,8 @@ bool BlockScaledGemmOperationProfiler::profile(
library::BlockScaledGemmDescription const &operation_desc =
static_cast<library::BlockScaledGemmDescription const &>(operation->description());
auto min_cc = operation_desc.tile_description.minimum_compute_capability;
bool is_dynamic_cluster_enabled = (min_cc >= 100);
auto cluster_shape = operation_desc.tile_description.cluster_shape;
bool is_dynamic_cluster_enabled = cluster_shape.m() == 0 || cluster_shape.n() == 0 || cluster_shape.k() == 0;
// Helper function wrapping up performance test with flexible parameters.
auto initialize_and_profile = [&](
@@ -1419,7 +1424,7 @@ bool BlockScaledGemmOperationProfiler::profile(
gemm_workspace_.reduction_arguments.pointer_mode = library::ScalarPointerMode::kHost;
}
update_workspace_(gemm_workspace_, problem_shape, leading_dim, preferred_cluster, fallback_cluster, raster_order, swizzle_size);
update_workspace_(gemm_workspace_, problem_shape, leading_dim, preferred_cluster, fallback_cluster, raster_order, swizzle_size, is_dynamic_cluster_enabled);
const auto can_implement = operation->can_implement(&gemm_workspace_.configuration, &gemm_workspace_.arguments);
if (can_implement != Status::kSuccess) {
@@ -1447,7 +1452,7 @@ bool BlockScaledGemmOperationProfiler::profile(
}
PerformanceResult curr_result(result);
update_result_(curr_result, operation_desc, problem_space, problem_shape, raster_order, preferred_cluster, fallback_cluster, swizzle_size);
update_result_(curr_result, operation_desc, problem_space, problem_shape, raster_order, preferred_cluster, fallback_cluster, swizzle_size, is_dynamic_cluster_enabled);
curr_result.status = profile_cutlass_(
curr_result,
@@ -1490,16 +1495,12 @@ bool BlockScaledGemmOperationProfiler::profile(
PerformanceResult result_base = results_.back();
results_.pop_back();
bool dynamic_cluster = int64_t(operation_desc.tile_description.cluster_shape.m()) == 0 ||
int64_t(operation_desc.tile_description.cluster_shape.n()) == 0 ||
int64_t(operation_desc.tile_description.cluster_shape.k()) == 0;
std::vector<std::array<int64_t, 3>> preferred_clusters;
std::vector<std::array<int64_t, 3>> fallback_clusters;
// Only loop over built-in cluster shape lists for dynamic cluster kernels
// and for kernels that can leverage the dynamic cluster feature.
if (dynamic_cluster && is_dynamic_cluster_enabled) {
if (is_dynamic_cluster_enabled) {
preferred_clusters = this->problem_.preferred_clusters;
fallback_clusters = this->problem_.fallback_clusters;
}
@@ -1510,7 +1511,7 @@ bool BlockScaledGemmOperationProfiler::profile(
for (auto preferred_cluster : preferred_clusters) {
for (auto fallback_cluster : fallback_clusters) {
if (dynamic_cluster && !is_valid_dynamic_cluster_shape(preferred_cluster, fallback_cluster)) {
if (is_dynamic_cluster_enabled && !is_valid_dynamic_cluster_shape(preferred_cluster, fallback_cluster)) {
continue;
}
for (auto swizzle_size : this->problem_.swizzle_sizes) {

View File

@@ -661,7 +661,8 @@ void GemmOperationProfiler::update_workspace_(
std::array<int64_t, 3> const &preferred_cluster,
std::array<int64_t, 3> const &fallback_cluster,
cutlass::library::RasterOrder const &raster_order,
int swizzle_size
int swizzle_size,
bool is_dynamic_cluster_enabled
) {
gemm_workspace.arguments.problem_size.m() = problem_shape.m();
@@ -675,16 +676,17 @@ void GemmOperationProfiler::update_workspace_(
gemm_workspace.arguments.swizzle_size = swizzle_size;
gemm_workspace.arguments.raster_order = raster_order;
gemm_workspace.arguments.cluster_shape = {int(preferred_cluster[0]), int(preferred_cluster[1]), int(preferred_cluster[2])};
gemm_workspace.arguments.cluster_shape_fallback = {int(fallback_cluster[0]), int(fallback_cluster[1]), int(fallback_cluster[2])};
if (is_dynamic_cluster_enabled) {
gemm_workspace.arguments.cluster_shape = {int(preferred_cluster[0]), int(preferred_cluster[1]), int(preferred_cluster[2])};
gemm_workspace.arguments.cluster_shape_fallback = {int(fallback_cluster[0]), int(fallback_cluster[1]), int(fallback_cluster[2])};
gemm_workspace.configuration.cluster_shape = {int(preferred_cluster[0]), int(preferred_cluster[1]), int(preferred_cluster[2])};
gemm_workspace.configuration.cluster_shape_fallback = {int(fallback_cluster[0]), int(fallback_cluster[1]), int(fallback_cluster[2])};
}
gemm_workspace.configuration.problem_size.m() = problem_shape.m();
gemm_workspace.configuration.problem_size.n() = problem_shape.n();
gemm_workspace.configuration.problem_size.k() = problem_shape.k();
gemm_workspace.configuration.cluster_shape = {int(preferred_cluster[0]), int(preferred_cluster[1]), int(preferred_cluster[2])};
gemm_workspace.configuration.cluster_shape_fallback = {int(fallback_cluster[0]), int(fallback_cluster[1]), int(fallback_cluster[2])};
gemm_workspace.configuration.lda = leading_dim[0];
gemm_workspace.configuration.ldb = leading_dim[1];
gemm_workspace.configuration.ldc = leading_dim[2];
@@ -699,7 +701,8 @@ void GemmOperationProfiler::update_result_(
cutlass::library::RasterOrder const &raster_order,
std::array<int64_t, 3> const &preferred_cluster,
std::array<int64_t, 3> const &fallback_cluster,
int swizzle_size
int swizzle_size,
bool is_dynamic_cluster_enabled
) {
result.bytes = problem_.bytes_with_problem_shape(operation_desc, problem_shape);
result.flops = problem_.flops_with_problem_shape(operation_desc, problem_shape);
@@ -711,12 +714,14 @@ void GemmOperationProfiler::update_result_(
set_argument(result, "raster_order", problem_space, library::to_string(raster_order));
set_argument(result, "swizzle_size", problem_space, swizzle_size);
set_argument(result, "cluster_m", problem_space, preferred_cluster[0]);
set_argument(result, "cluster_n", problem_space, preferred_cluster[1]);
set_argument(result, "cluster_k", problem_space, preferred_cluster[2]);
set_argument(result, "cluster_m_fallback", problem_space, fallback_cluster[0]);
set_argument(result, "cluster_n_fallback", problem_space, fallback_cluster[1]);
set_argument(result, "cluster_k_fallback", problem_space, fallback_cluster[2]);
if (is_dynamic_cluster_enabled) {
set_argument(result, "cluster_m", problem_space, preferred_cluster[0]);
set_argument(result, "cluster_n", problem_space, preferred_cluster[1]);
set_argument(result, "cluster_k", problem_space, preferred_cluster[2]);
set_argument(result, "cluster_m_fallback", problem_space, fallback_cluster[0]);
set_argument(result, "cluster_n_fallback", problem_space, fallback_cluster[1]);
set_argument(result, "cluster_k_fallback", problem_space, fallback_cluster[2]);
}
}
@@ -1576,9 +1581,8 @@ bool GemmOperationProfiler::profile(
library::GemmDescription const &operation_desc =
static_cast<library::GemmDescription const &>(operation->description());
auto min_cc = operation_desc.tile_description.minimum_compute_capability;
bool is_dynamic_cluster_enabled = (min_cc >= 100);
auto cluster_shape = operation_desc.tile_description.cluster_shape;
bool is_dynamic_cluster_enabled = cluster_shape.m() == 0 || cluster_shape.n() == 0 || cluster_shape.k() == 0;
// Helper function wrapping up performance test with flexible parameters.
auto initialize_and_profile = [&](
@@ -1618,7 +1622,7 @@ bool GemmOperationProfiler::profile(
workspace.reduction_arguments.pointer_mode = library::ScalarPointerMode::kHost;
}
update_workspace_(workspace, problem_shape, leading_dim, preferred_cluster, fallback_cluster, raster_order, swizzle_size);
update_workspace_(workspace, problem_shape, leading_dim, preferred_cluster, fallback_cluster, raster_order, swizzle_size, is_dynamic_cluster_enabled);
const auto can_implement = operation->can_implement(&workspace.configuration, &workspace.arguments);
if (can_implement != Status::kSuccess) {
@@ -1672,7 +1676,7 @@ bool GemmOperationProfiler::profile(
}
PerformanceResult curr_result(result);
update_result_(curr_result, operation_desc, problem_space, problem_shape, raster_order, preferred_cluster, fallback_cluster, swizzle_size);
update_result_(curr_result, operation_desc, problem_space, problem_shape, raster_order, preferred_cluster, fallback_cluster, swizzle_size, is_dynamic_cluster_enabled);
curr_result.status = profile_cutlass_(
curr_result,
@@ -1712,17 +1716,13 @@ bool GemmOperationProfiler::profile(
std::vector<PerformanceResult> candidates;
PerformanceResult result_base = results_.back();
results_.pop_back();
bool dynamic_cluster = int64_t(operation_desc.tile_description.cluster_shape.m()) == 0 ||
int64_t(operation_desc.tile_description.cluster_shape.n()) == 0 ||
int64_t(operation_desc.tile_description.cluster_shape.k()) == 0;
std::vector<std::array<int64_t, 3>> preferred_clusters;
std::vector<std::array<int64_t, 3>> fallback_clusters;
// Only loop over built-in cluster shape lists for dynamic cluster kernels
// and for kernels that can leverage the dynamic cluster feature.
if (dynamic_cluster && is_dynamic_cluster_enabled) {
if (is_dynamic_cluster_enabled) {
preferred_clusters = this->problem_.preferred_clusters;
fallback_clusters = this->problem_.fallback_clusters;
}
@@ -1733,7 +1733,7 @@ bool GemmOperationProfiler::profile(
for (auto preferred_cluster : preferred_clusters) {
for (auto fallback_cluster : fallback_clusters) {
if (dynamic_cluster && !is_valid_dynamic_cluster_shape(preferred_cluster, fallback_cluster)) {
if (is_dynamic_cluster_enabled && !is_valid_dynamic_cluster_shape(preferred_cluster, fallback_cluster)) {
continue;
}
for (auto swizzle_size : this->problem_.swizzle_sizes) {

View File

@@ -538,23 +538,33 @@ void GroupedGemmOperationProfiler::GroupedGemmProblem::initialize_result(
library::lexical_cast(beta, operation_desc.gemm.element_epilogue));
}
void GroupedGemmOperationProfiler::update_result_(
void GroupedGemmOperationProfiler::update_workspace_and_result_(
GroupedGemmWorkspace &gemm_workspace,
PerformanceResult &result,
ProblemSpace const &problem_space,
cutlass::library::RasterOrder const &raster_order,
std::array<int64_t, 3> const &preferred_cluster,
std::array<int64_t, 3> const &fallback_cluster,
int swizzle_size
int swizzle_size,
bool is_dynamic_cluster_enabled
) {
gemm_workspace.arguments.swizzle_size = swizzle_size;
gemm_workspace.arguments.raster_order = raster_order;
set_argument(result, "raster_order", problem_space, library::to_string(raster_order));
set_argument(result, "swizzle_size", problem_space, swizzle_size);
set_argument(result, "cluster_m", problem_space, preferred_cluster[0]);
set_argument(result, "cluster_n", problem_space, preferred_cluster[1]);
set_argument(result, "cluster_k", problem_space, preferred_cluster[2]);
set_argument(result, "cluster_m_fallback", problem_space, fallback_cluster[0]);
set_argument(result, "cluster_n_fallback", problem_space, fallback_cluster[1]);
set_argument(result, "cluster_k_fallback", problem_space, fallback_cluster[2]);
if (is_dynamic_cluster_enabled) {
gemm_workspace.arguments.cluster_shape = {int(preferred_cluster[0]), int(preferred_cluster[1]), int(preferred_cluster[2])};
gemm_workspace.arguments.cluster_shape_fallback = {int(fallback_cluster[0]), int(fallback_cluster[1]), int(fallback_cluster[2])};
set_argument(result, "cluster_m", problem_space, preferred_cluster[0]);
set_argument(result, "cluster_n", problem_space, preferred_cluster[1]);
set_argument(result, "cluster_k", problem_space, preferred_cluster[2]);
set_argument(result, "cluster_m_fallback", problem_space, fallback_cluster[0]);
set_argument(result, "cluster_n_fallback", problem_space, fallback_cluster[1]);
set_argument(result, "cluster_k_fallback", problem_space, fallback_cluster[2]);
}
}
/////////////////////////////////////////////////////////////////////////////////////////////////
@@ -1605,9 +1615,8 @@ bool GroupedGemmOperationProfiler::profile_cutlass_for_fixed_shape_(
library::GroupedGemmDescription const &operation_desc =
static_cast<library::GroupedGemmDescription const &>(operation->description());
auto min_cc = operation_desc.tile_description.minimum_compute_capability;
bool is_dynamic_cluster_enabled = (min_cc >= 100);
auto cluster_shape = operation_desc.tile_description.cluster_shape;
bool is_dynamic_cluster_enabled = cluster_shape.m() == 0 || cluster_shape.n() == 0 || cluster_shape.k() == 0;
// Helper function to test validity of fallback cluster shapes and preferred cluster shapes.
auto is_valid_dynamic_cluster_shape = [](const std::array<int64_t, 3>& preferred_cluster, const std::array<int64_t, 3>& fallback_cluster) {
@@ -1636,19 +1645,15 @@ bool GroupedGemmOperationProfiler::profile_cutlass_for_fixed_shape_(
PerformanceResult result_base = results_.back();
results_.pop_back();
bool dynamic_cluster = int64_t(operation_desc.tile_description.cluster_shape.m()) == 0 ||
int64_t(operation_desc.tile_description.cluster_shape.n()) == 0 ||
int64_t(operation_desc.tile_description.cluster_shape.k()) == 0;
std::vector<std::array<int64_t, 3>> preferred_clusters;
std::vector<std::array<int64_t, 3>> fallback_clusters;
// Only loop over built-in cluster shape lists for dynamic cluster kernels
// and for kernels that can leverage the dynamic cluster feature.
if (dynamic_cluster && is_dynamic_cluster_enabled) {
if (is_dynamic_cluster_enabled) {
preferred_clusters = this->problem_.preferred_clusters;
fallback_clusters = this->problem_.fallback_clusters;
}
}
else {
preferred_clusters = {{int(problem_.cluster_m), int(problem_.cluster_n), int(problem_.cluster_k)}};
fallback_clusters = {{int(problem_.cluster_m_fallback), int(problem_.cluster_n_fallback), int(problem_.cluster_k_fallback)}};
@@ -1656,13 +1661,13 @@ bool GroupedGemmOperationProfiler::profile_cutlass_for_fixed_shape_(
for (auto preferred_cluster : preferred_clusters) {
for (auto fallback_cluster : fallback_clusters) {
if (dynamic_cluster && !is_valid_dynamic_cluster_shape(preferred_cluster, fallback_cluster)) {
if (is_dynamic_cluster_enabled && !is_valid_dynamic_cluster_shape(preferred_cluster, fallback_cluster)) {
continue;
}
for (auto swizzle_size : this->problem_.swizzle_sizes) {
for (auto raster_order : this->problem_.raster_orders) {
PerformanceResult curr_result(result_base);
update_result_(curr_result, problem_space, raster_order, preferred_cluster, fallback_cluster, swizzle_size);
update_workspace_and_result_(gemm_workspace_, curr_result, problem_space, raster_order, preferred_cluster, fallback_cluster, swizzle_size, is_dynamic_cluster_enabled);
curr_result.status = profile_cutlass_(
curr_result,
options,

View File

@@ -75,15 +75,7 @@ __global__ void dequantize_kernel(DequantizedElement* dq_buffer,
// Represent the full tensors to gmem elements.
// These are expected to have shape [MN, K, L]
cute::Tensor gmem_op_dq = cute::make_tensor(cute::make_gmem_ptr(dq_buffer), operand_layout);
auto init_quantized_iterator = [&]() {
if constexpr (cute::sizeof_bits_v<QuantizedElement> >= 8) {
return cute::make_gmem_ptr(q_buffer);
}
else {
return cute::subbyte_iterator<const QuantizedElement>(q_buffer);
}
};
cute::Tensor gmem_op_q = cute::make_tensor(init_quantized_iterator(), operand_layout);
cute::Tensor gmem_op_q = cute::make_tensor(cute::make_gmem_ptr<QuantizedElement const>(q_buffer), operand_layout);
// While the scales are expected to have shape [MN, G, L] but with a stride to allow broadcasting
// It is expected that K % G == 0
cute::Tensor gmem_scale_broadcasted = cute::make_tensor(make_gmem_ptr(scale_buffer), broadcasted_scale_layout);

View File

@@ -95,6 +95,7 @@ template<
class ElementCompute_,
class ElementC_,
class ElementOut_,
bool ResidualAdd_,
class TensorAlpha_,
class TensorBeta_,
class TensorBias_,
@@ -110,6 +111,8 @@ struct ConvEpilogueFusionParams {
using TensorBeta = TensorBeta_;
using TensorBias = TensorBias_;
using ActivationFunctor = ActivationFunctor_;
static constexpr bool ResidualAdd = ResidualAdd_; // Source added after activation
ElementScalar alpha = ElementScalar(1);
ElementScalar beta = ElementScalar(0);
@@ -228,12 +231,17 @@ private:
epi_fusion_params_.tensor_alpha[k] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data()) ?
epi_fusion_params_.tensor_beta[k] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(k, q, n, g));
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator);
if (not EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(k, q, n, g));
}
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[k]);
}
output = epi_activation(output);
if (EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(k, q, n, g));
}
tensor_d_(k, q, n, g) = output_converter(output);
}
}
@@ -279,12 +287,17 @@ private:
epi_fusion_params_.tensor_alpha[k] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data()) ?
epi_fusion_params_.tensor_beta[k] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(k, q, p, n, g));
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator);
if (not EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(k, q, p, n, g));
}
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[k]);
}
output = epi_activation(output);
if (EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(k, q, p, n, g));
}
tensor_d_(k, q, p, n, g) = output_converter(output);
}
}
@@ -337,12 +350,17 @@ private:
epi_fusion_params_.tensor_alpha[k] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data()) ?
epi_fusion_params_.tensor_beta[k] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(k, q, p, z, n, g));
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator);
if (not EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(k, q, p, z, n, g));
}
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[k]);
}
output = epi_activation(output);
if (EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(k, q, p, z, n, g));
}
tensor_d_(k, q, p, z, n, g) = output_converter(output);
}
}
@@ -389,12 +407,17 @@ private:
? epi_fusion_params_.tensor_alpha[c] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data())
? epi_fusion_params_.tensor_beta[c] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(c, w, n, g));
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator);
if (not EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(c, w, n, g));
}
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[c]);
}
output = epi_activation(output);
if (EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(c, w, n, g));
}
tensor_d_(c, w, n, g) = output_converter(output);
}
}
@@ -451,12 +474,17 @@ private:
? epi_fusion_params_.tensor_alpha[c] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data())
? epi_fusion_params_.tensor_beta[c] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(c, w, h, n, g));
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator);
if (not EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(c, w, h, n, g));
}
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[c]);
}
output = epi_activation(output);
if (EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(c, w, h, n, g));
}
tensor_d_(c, w, h, n, g) = output_converter(output);
}
@@ -527,12 +555,17 @@ private:
? epi_fusion_params_.tensor_alpha[c] : epi_fusion_params_.alpha;
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data())
? epi_fusion_params_.tensor_beta[c] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(c, w, h, d, n, g));
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator);
if (not EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(c, w, h, d, n, g));
}
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[c]);
}
output = epi_activation(output);
if (EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(c, w, h, d, n, g));
}
tensor_d_(c, w, h, d, n, g) = output_converter(output);
}
}
@@ -583,12 +616,17 @@ private:
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data()) ?
epi_fusion_params_.tensor_beta[c] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(c, s, k, g));
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator);
if (not EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(c, s, k, g));
}
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[c]);
}
output = epi_activation(output);
if (EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(c, s, k, g));
}
tensor_d_(c, s, k, g) = output_converter(output);
}
}
@@ -643,12 +681,17 @@ private:
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data()) ?
epi_fusion_params_.tensor_beta[c] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(c, s, r, k, g));
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator);
if (not EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(c, s, r, k, g));
}
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[c]);
}
output = epi_activation(output);
if (EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(c, s, r, k, g));
}
tensor_d_(c, s, r, k, g) = output_converter(output);
}
}
@@ -711,12 +754,17 @@ private:
ElementScalar beta = raw_pointer_cast(epi_fusion_params_.tensor_beta.data()) ?
epi_fusion_params_.tensor_beta[c] : epi_fusion_params_.beta;
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator) +
scale_converter(beta) * residual_converter(tensor_c_(c, s, r, t, k, g));
ElementCompute output = scale_converter(alpha) * acc_converter(accumulator);
if (not EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(c, s, r, t, k, g));
}
if (raw_pointer_cast(epi_fusion_params_.tensor_bias.data())) {
output += bias_converter(epi_fusion_params_.tensor_bias[c]);
}
output = epi_activation(output);
if (EpilogueFusionParams::ResidualAdd) {
output += scale_converter(beta) * residual_converter(tensor_c_(c, s, r, t, k, g));
}
tensor_d_(c, s, r, t, k, g) = output_converter(output);
}
}