biondizzle/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes (#304)

Browse Source 
* Merge with private repo

* Update README

* Update README

* Update README

* Add PyTorch requirements

* Fix sync scopes for MQA logits (#256)

* Update README
This commit is contained in:
Chenggang Zhao
2026-04-17 09:45:14 +08:00
committed by GitHub
parent d30fc36c8f
commit 7f2a703ed5
109 changed files with 12101 additions and 3219 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
csrc/jit_kernels/impls/epilogue.hpp
View File

@@ -6,7 +6,7 @@
namespace deep_gemm {
static std::string get_default_epilogue_type(const std::optional<std::string>& epilogue_type) {
return epilogue_type.value_or("EpilogueIdentity");
return epilogue_type.value_or("epilogue::transform::EpilogueIdentity");
}
} // namespace deep_gemm

72
csrc/jit_kernels/impls/runtime_utils.hpp
View File

@@ -20,6 +20,9 @@ static int get_non_contiguous_dim(const cute::UMMA::Major& major) {
}
static int get_compiled_dim(const int& dim, const char& name, const std::string& compiled_dims) {
if (heuristics_runtime->get_ignore_compile_dims())
return 0;
for (const char& c: compiled_dims) {
if (name == c)
return dim;
@@ -58,8 +61,19 @@ static std::string to_string(const at::ScalarType& dtype) {
}
}
static std::string to_string(const float& v) {
if (std::isfinite(v)) {
return fmt::format(R"({:a}f)", v);
} else if (std::isinf(v)) {
return v > 0 ? "cute::numeric_limits<float>::infinity()"
: "-cute::numeric_limits<float>::infinity()";
}
DG_HOST_UNREACHABLE("NaN input is not supported");
}
static CUtensorMapDataType aten_dtype_to_tensor_map_dtype(const at::ScalarType& dtype,
const bool& allow_tf32) {
const bool& allow_tf32,
const bool& fp4_unpacked_smem) {
if (allow_tf32 and dtype == torch::kFloat)
return CU_TENSOR_MAP_DATA_TYPE_TFLOAT32;
@@ -68,13 +82,16 @@ static CUtensorMapDataType aten_dtype_to_tensor_map_dtype(const at::ScalarType&
case torch::kFloat: return CU_TENSOR_MAP_DATA_TYPE_FLOAT32;
case torch::kBFloat16: return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16;
case torch::kFloat8_e4m3fn: return CU_TENSOR_MAP_DATA_TYPE_UINT8;
case kPackedFP4: return CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B;
#if CUDA_VERSION >= 12080
case kPackedFP4: return fp4_unpacked_smem ? CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN16B
: CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B;
#endif
default: DG_HOST_UNREACHABLE("Unsupported dtype");
}
}
static CUtensorMapSwizzle mode_into_tensor_map_swizzle(const int& mode, const int& base) {
#if CUDART_VERSION >= 12080
#if CUDA_VERSION >= 12080
if (base != 0) {
DG_HOST_ASSERT(base == 32 and mode == 128);
return CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B;
@@ -97,14 +114,20 @@ static CUtensorMap make_tma_2d_desc(const torch::Tensor& t,
int smem_inner_dim, int smem_outer_dim,
const int& gmem_outer_stride,
const int& swizzle_mode, const int& swizzle_base = 0,
const bool& allow_tf32 = false) {
const auto& elem_size = static_cast<int>(t.element_size());
const bool& allow_tf32 = false,
const bool& fp4_unpacked_smem = true) {
const auto elem_size = static_cast<int>(t.element_size());
if (swizzle_mode != 0)
smem_inner_dim = swizzle_mode / elem_size;
// Inner dim must be a multiple of 64B for .b4x16_p64
if (t.scalar_type() == kPackedFP4)
DG_HOST_ASSERT(gmem_inner_dim % 128 == 0);
if (t.scalar_type() == kPackedFP4) {
// Inner dim must be a multiple of 64B for .b4x16_p64
DG_HOST_ASSERT(not fp4_unpacked_smem or gmem_inner_dim % 128 == 0);
// Fix FP4 packed smem
if (not fp4_unpacked_smem and swizzle_mode != 0)
smem_inner_dim = swizzle_mode * 2;
}
CUtensorMap tensor_map;
const cuuint64_t gmem_dims[2] = {static_cast<cuuint64_t>(gmem_inner_dim), static_cast<cuuint64_t>(gmem_outer_dim)};
@@ -112,12 +135,13 @@ static CUtensorMap make_tma_2d_desc(const torch::Tensor& t,
const cuuint64_t gmem_strides[1] = {static_cast<cuuint64_t>(gmem_outer_stride * elem_size), };
const cuuint32_t elem_strides[2] = {1, 1};
if (get_env<int>("DG_JIT_DEBUG")) {
printf("Making TMA desc: global memory: %d %d, shared memory: %d %d, outer stride: %d, swizzle: %d (base: %d), elem size: %d\n",
printf("Making TMA desc: global memory: %d %d, shared memory: %d %d, outer stride: %d, swizzle: %d (base: %d), elem size: %d, pointer: %llu\n",
gmem_inner_dim, gmem_outer_dim, smem_inner_dim, smem_outer_dim,
gmem_outer_stride, swizzle_mode, swizzle_base, elem_size);
gmem_outer_stride, swizzle_mode, swizzle_base, elem_size,
reinterpret_cast<unsigned long long>(t.data_ptr()));
}
DG_CUDA_DRIVER_CHECK(lazy_cuTensorMapEncodeTiled(
&tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32),
&tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32, fp4_unpacked_smem),
2, t.data_ptr(), gmem_dims, gmem_strides, smem_dims, elem_strides,
CU_TENSOR_MAP_INTERLEAVE_NONE, mode_into_tensor_map_swizzle(swizzle_mode, swizzle_base),
CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));
@@ -129,14 +153,20 @@ static CUtensorMap make_tma_3d_desc(const torch::Tensor& t,
int smem_dim_0, int smem_dim_1, int smem_dim_2,
const int& gmem_stride_0, const int& gmem_stride_1,
const int& swizzle_mode, const int& swizzle_base = 0,
const bool& allow_tf32 = false) {
const auto& elem_size = static_cast<int>(t.element_size());
const bool& allow_tf32 = false,
const bool& fp4_unpacked_smem = true) {
const auto elem_size = static_cast<int>(t.element_size());
if (swizzle_mode != 0)
smem_dim_0 = swizzle_mode / elem_size;
// Inner dim must be a multiple of 64B for .b4x16_p64
if (t.scalar_type() == kPackedFP4)
DG_HOST_ASSERT(gmem_dim_0 % 128 == 0);
if (t.scalar_type() == kPackedFP4) {
// Inner dim must be a multiple of 64B for .b4x16_p64
DG_HOST_ASSERT(not fp4_unpacked_smem or gmem_dim_0 % 128 == 0);
// Fix fp4 packed smem
if (not fp4_unpacked_smem and swizzle_mode != 0)
smem_dim_0 = swizzle_mode * 2;
}
CUtensorMap tensor_map;
const cuuint64_t gmem_dims[3] = {static_cast<cuuint64_t>(gmem_dim_0), static_cast<cuuint64_t>(gmem_dim_1), static_cast<cuuint64_t>(gmem_dim_2),};
@@ -149,7 +179,7 @@ static CUtensorMap make_tma_3d_desc(const torch::Tensor& t,
gmem_stride_0, gmem_stride_1, swizzle_mode, elem_size);
}
DG_CUDA_DRIVER_CHECK(lazy_cuTensorMapEncodeTiled(
&tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32),
&tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32, fp4_unpacked_smem),
3, t.data_ptr(), gmem_dims, gmem_strides, smem_dims, elem_strides,
CU_TENSOR_MAP_INTERLEAVE_NONE, mode_into_tensor_map_swizzle(swizzle_mode, swizzle_base),
CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE));
@@ -166,8 +196,8 @@ static CUtensorMap make_tma_a_desc(const cute::UMMA::Major& major,
const bool& allow_tf32 = false) {
if (num_groups > 1)
DG_HOST_ASSERT(major == cute::UMMA::Major::K);
const auto& [gmem_inner_dim, gmem_outer_dim] = get_inner_outer_dims(major, shape_k, shape_m * num_groups);
const auto& [smem_inner_dim, smem_outer_dim] = get_inner_outer_dims(major, block_k, block_m);
const auto [gmem_inner_dim, gmem_outer_dim] = get_inner_outer_dims(major, shape_k, shape_m * num_groups);
const auto [smem_inner_dim, smem_outer_dim] = get_inner_outer_dims(major, block_k, block_m);
return make_tma_2d_desc(t,
gmem_inner_dim, gmem_outer_dim,
smem_inner_dim, smem_outer_dim,
@@ -184,8 +214,8 @@ static CUtensorMap make_tma_b_desc(const cute::UMMA::Major& major,
const int& num_groups,
const int& swizzle_mode, const int& swizzle_base = 0,
const bool& allow_tf32 = false) {
const auto& [gmem_inner_dim, gmem_outer_dim] = get_inner_outer_dims(major, shape_k, shape_n);
const auto& [smem_inner_dim, smem_outer_dim] = get_inner_outer_dims(major, block_k, block_n);
const auto [gmem_inner_dim, gmem_outer_dim] = get_inner_outer_dims(major, shape_k, shape_n);
const auto [smem_inner_dim, smem_outer_dim] = get_inner_outer_dims(major, block_k, block_n);
// `num_groups` is always applied into the outer dimensions
return make_tma_2d_desc(t,

426
csrc/jit_kernels/impls/sm100_bf16_gemm.hpp
View File

@@ -16,9 +16,7 @@ namespace deep_gemm {
class SM100BF16GemmRuntime final: public LaunchRuntime<SM100BF16GemmRuntime> {
public:
struct Args {
int m, n, k, num_groups;
const std::string& compiled_dims;
GemmDesc gemm_desc;
GemmConfig gemm_config;
LaunchArgs launch_args;
@@ -45,28 +43,32 @@ static void __instantiate_kernel() {{
{}, {},
{}, {},
{},
{},
{}, {}, {},
{}
>);
}};
)",
to_string(args.gemm_config.major_a), to_string(args.gemm_config.major_b),
get_compiled_dim(args.m, 'm', args.compiled_dims), get_compiled_dim(args.n, 'n', args.compiled_dims), get_compiled_dim(args.k, 'k', args.compiled_dims),
args.gemm_config.block_m, args.gemm_config.block_n, args.gemm_config.block_k,
args.num_groups,
args.gemm_config.smem_config.swizzle_a_mode, args.gemm_config.smem_config.swizzle_b_mode, args.gemm_config.smem_config.swizzle_cd_mode,
args.gemm_config.num_stages,
args.gemm_config.thread_config.num_non_epilogue_threads, args.gemm_config.thread_config.num_epilogue_threads,
args.gemm_config.multicast_config.num_multicast, args.gemm_config.multicast_config.is_multicast_on_a,
args.gemm_config.num_sms,
to_string(args.gemm_config.gemm_type), args.gemm_config.with_accumulation, to_string(args.gemm_config.cd_dtype),
args.gemm_config.tc_util);
to_string(args.gemm_desc.major_a), to_string(args.gemm_desc.major_b),
get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims),
args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k,
args.gemm_desc.num_groups,
args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode, args.gemm_config.storage_config.swizzle_cd_mode,
args.gemm_config.pipeline_config.num_stages,
args.gemm_config.launch_config.num_non_epilogue_threads, args.gemm_config.launch_config.num_epilogue_threads,
args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1,
args.gemm_config.launch_config.num_sms,
args.gemm_config.layout.swap_ab,
to_string(args.gemm_desc.gemm_type), args.gemm_desc.with_accumulation, to_string(args.gemm_desc.cd_dtype),
args.gemm_desc.tc_util);
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
// TODO: optimize `args` copy
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.grouped_layout, args.m, args.n, args.k,
args.grouped_layout, args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
args.tensor_map_a, args.tensor_map_b,
args.tensor_map_cd));
}
@@ -79,45 +81,49 @@ static void sm100_bf16_gemm(const torch::Tensor& a,
const int& m, const int& n, const int& k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::Normal, KernelType::KernelNoSF,
m, n, k, 1, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Normal,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = k, .num_groups = 1,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM100ArchSpec>(desc);
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.storage_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = 1,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_gemm", code);
const auto code = SM100BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm100_bf16_gemm", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
@@ -130,53 +136,61 @@ static void sm100_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a,
const std::string& compiled_dims,
const bool& use_psum_layout,
const std::optional<int>& expected_m_for_psum_layout) {
const auto& gemm_type = use_psum_layout ? GemmType::MGroupedContiguousWithPsumLayout : GemmType::MGroupedContiguous;
const auto gemm_type = use_psum_layout ?
GemmType::MGroupedContiguousWithPsumLayout : GemmType::MGroupedContiguous;
// Only psum layout can use expected m
if (expected_m_for_psum_layout)
DG_HOST_ASSERT(use_psum_layout);
// NOTES: If actual M is dynamic, estimate config via `num_groups` and `expected_m`.
// Otherwise, treat the contiguous layout as a whole.
const auto& m_for_config = expected_m_for_psum_layout.has_value() ? expected_m_for_psum_layout.value() : m;
const auto& num_groups_for_config = expected_m_for_psum_layout.has_value() ? num_groups : 1;
const auto desc = GemmDesc {
.gemm_type = gemm_type,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = expected_m_for_psum_layout.value_or(m),
.expected_n = n, .expected_k = k,
.expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1
};
const auto config = get_best_config<SM100ArchSpec>(desc);
const auto& config = get_best_config<SM100ArchSpec>(
gemm_type, KernelType::KernelNoSF,
// NOTES: `num_groups` is 1, since the contiguous layout is seen as a whole
m_for_config, n, k, num_groups_for_config, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), false,
device_runtime->get_num_sms());
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.storage_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
const SM100BF16GemmRuntime::Args args = {
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = grouped_layout.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_m_grouped_gemm_contiguous", code);
const auto code = SM100BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm100_bf16_m_grouped_gemm_contiguous", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
@@ -188,45 +202,50 @@ static void sm100_m_grouped_bf16_gemm_masked(const torch::Tensor& a,
const int& expected_m,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::MGroupedMasked, KernelType::KernelNoSF,
expected_m, n, k, num_groups, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), false,
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::MGroupedMasked,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM100ArchSpec>(desc);
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), num_groups,
config.smem_config.swizzle_cd_mode);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), num_groups,
config.storage_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
const SM100BF16GemmRuntime::Args args = {
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = masked_m.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_m_grouped_gemm_masked", code);
const auto code = SM100BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm100_bf16_m_grouped_gemm_masked", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
@@ -241,54 +260,59 @@ static void sm100_bf16_k_grouped_gemm(const torch::Tensor& a,
DG_HOST_ASSERT(major_a == cute::UMMA::Major::MN and major_b == cute::UMMA::Major::MN);
int sum_k = 0;
for (const auto& k: ks) {
for (const auto k: ks) {
sum_k += k;
DG_HOST_ASSERT(k % 128 == 0);
}
const auto& num_groups = static_cast<int>(ks.size());
const auto num_groups = static_cast<int>(ks.size());
// Get config using max K for better performance
const auto& max_k = *std::max_element(ks.begin(), ks.end());
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::KGroupedContiguous, KernelType::KernelNoSF,
m, n, max_k, num_groups, cute::UMMA::Major::MN, cute::UMMA::Major::MN,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto max_k = *std::max_element(ks.begin(), ks.end());
const auto desc = GemmDesc {
.gemm_type = GemmType::KGroupedContiguous,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = sum_k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM100ArchSpec>(desc);
// Create tensor descriptors
const auto& tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(0)), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(cute::UMMA::Major::MN, b, n, sum_k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(0)), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(1)), num_groups,
config.smem_config.swizzle_cd_mode);
const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(0)), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::MN, b, n, sum_k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(0)), 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(1)), num_groups,
config.storage_config.swizzle_cd_mode);
// Launch kernel
const SM100BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = sum_k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = ks_tensor.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_k_grouped_gemm", code);
const auto code = SM100BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm100_bf16_k_grouped_gemm", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
@@ -297,46 +321,46 @@ static void sm100_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a,
const torch::Tensor& tensor_d,
const int& b, const int& h, const int& r, const int& d,
const std::string& compiled_dims = "nk") {
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::Batched, KernelType::KernelNoSF,
b, d, r, h, cute::UMMA::Major::K, cute::UMMA::Major::K,
tensor_a.scalar_type(), tensor_b.scalar_type(),
tensor_d.scalar_type(), false,
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Batched,
.kernel_type = KernelType::KernelNoSF,
.m = b, .n = d, .k = r, .num_groups = h,
.a_dtype = tensor_a.scalar_type(), .b_dtype = tensor_b.scalar_type(),
.cd_dtype = tensor_d.scalar_type(),
.major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::K,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM100ArchSpec>(desc);
const int& load_block_m = SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& tensor_map_a = make_tma_3d_desc(tensor_a, r, b, h,
config.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.smem_config.swizzle_a_mode);
const int& load_block_n = SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
config.block_k, load_block_n, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.smem_config.swizzle_b_mode);
const int& store_block_m = SM100ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM100ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_cd = make_tma_3d_desc(tensor_d, d, b, h,
store_block_n, store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.smem_config.swizzle_cd_mode);
const auto tensor_map_a = make_tma_3d_desc(tensor_a, r, b, h,
config.layout.block_k, config.storage_config.load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
config.layout.block_k, config.storage_config.load_block_n, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_3d_desc(tensor_d, d, b, h,
config.storage_config.store_block_n, config.storage_config.store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.storage_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
.m = b, .n = d, .k = r,
.num_groups = h,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_bhr_hdr_bhd", code);
const auto code = SM100BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm100_bf16_bhr_hdr_bhd", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
@@ -345,46 +369,46 @@ static void sm100_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a,
const torch::Tensor& tensor_d,
const int& b, const int& h, const int& r, const int& d,
const std::string& compiled_dims = "nk") {
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::Batched, KernelType::KernelNoSF,
b, r, d, h, cute::UMMA::Major::K, cute::UMMA::Major::MN,
tensor_a.scalar_type(), tensor_b.scalar_type(),
tensor_d.scalar_type(), false,
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Batched,
.kernel_type = KernelType::KernelNoSF,
.m = b, .n = r, .k = d, .num_groups = h,
.a_dtype = tensor_a.scalar_type(), .b_dtype = tensor_b.scalar_type(),
.cd_dtype = tensor_d.scalar_type(),
.major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::MN,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM100ArchSpec>(desc);
const int& load_block_m = SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& tensor_map_a = make_tma_3d_desc(tensor_a, d, b, h,
config.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.smem_config.swizzle_a_mode);
const int& load_block_n = SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
load_block_n, config.block_k, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.smem_config.swizzle_b_mode);
const int& store_block_m = SM100ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM100ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_cd = make_tma_3d_desc(tensor_d, r, b, h,
store_block_n, store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.smem_config.swizzle_cd_mode);
const auto tensor_map_a = make_tma_3d_desc(tensor_a, d, b, h,
config.layout.block_k, config.storage_config.load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
config.storage_config.load_block_n, config.layout.block_k, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_3d_desc(tensor_d, r, b, h,
config.storage_config.store_block_n, config.storage_config.store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.storage_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
.m = b, .n = r, .k = d,
.num_groups = h,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_bhd_hdr_bhr", code);
const auto code = SM100BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm100_bf16_bhd_hdr_bhr", code);
SM100BF16GemmRuntime::launch(runtime, args);
}

22
csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp
View File

@@ -85,11 +85,11 @@ static void sm100_bmn_bnk_mn_gemm(const torch::Tensor &a,
// NOTES: we select 4 as start, as it is tested to be faster than values > 4
int num_stages = 4, smem_size = 0;
while (true) {
const int& smem_cd = block_m * swizzle_cd_mode * 2;
const int& smem_a_per_stage = block_m * block_k * sizeof(cutlass::bfloat16_t);
const int& smem_b_per_stage = block_n * block_k * sizeof(cutlass::bfloat16_t);
const int& smem_barrier = SM100ArchSpec::get_barrier_smem_size(num_stages);
const int& smem_tmem_ptr = SM100ArchSpec::get_tmem_ptr_smem_size();
const int smem_cd = block_m * swizzle_cd_mode * 2;
const int smem_a_per_stage = block_m * block_k * sizeof(cutlass::bfloat16_t);
const int smem_b_per_stage = block_n * block_k * sizeof(cutlass::bfloat16_t);
const int smem_barrier = num_stages * 8 * 3 + 2 * 8 * 2 + 8;
const int smem_tmem_ptr = 4;
smem_size = 0;
smem_size += smem_cd;
@@ -112,11 +112,11 @@ static void sm100_bmn_bnk_mn_gemm(const torch::Tensor &a,
num_stages, smem_size, swizzle_ab_mode, swizzle_cd_mode);
}
const auto& tensor_map_a = make_tma_2d_desc(a, k, s * m, block_k, block_m, k, swizzle_ab_mode);
const auto& tensor_map_b = make_tma_2d_desc(b, k, s * n, block_k, block_n, k, swizzle_ab_mode);
const auto& tensor_map_d = make_tma_2d_desc(d, n, m, block_n, block_m, n, swizzle_cd_mode);
const auto tensor_map_a = make_tma_2d_desc(a, k, s * m, block_k, block_m, k, swizzle_ab_mode);
const auto tensor_map_b = make_tma_2d_desc(b, k, s * n, block_k, block_n, k, swizzle_ab_mode);
const auto tensor_map_d = make_tma_2d_desc(d, n, m, block_n, block_m, n, swizzle_cd_mode);
const SM100BmkBnkMnRuntime::Args& args = {
const SM100BmkBnkMnRuntime::Args args = {
.s = s, .m = m, .n = n, .k = k,
.block_m = block_m, .block_n = block_n, .block_k = block_k,
.split_factor = split_factor,
@@ -129,8 +129,8 @@ static void sm100_bmn_bnk_mn_gemm(const torch::Tensor &a,
.tensor_map_b = tensor_map_b,
.tensor_map_d = tensor_map_d
};
const auto& code = SM100BmkBnkMnRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bmn_bnk_mn_gemm", code);
const auto code = SM100BmkBnkMnRuntime::generate(args);
const auto runtime = compiler->build("sm100_bmn_bnk_mn_gemm", code);
SM100BmkBnkMnRuntime::launch(runtime, args);
}

459
csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp Normal file
View File

@@ -0,0 +1,459 @@
#pragma once
#include <torch/python.h>
#include "../../jit/compiler.hpp"
#include "../../jit/device_runtime.hpp"
#include "../../jit/kernel_runtime.hpp"
#include "../../utils/exception.hpp"
#include "../../utils/format.hpp"
#include "../../utils/math.hpp"
#include "../heuristics/sm100.hpp"
#include "epilogue.hpp"
#include "runtime_utils.hpp"
namespace deep_gemm {
class SM100FP8FP4Gemm1D1DRuntime final: public LaunchRuntime<SM100FP8FP4Gemm1D1DRuntime> {
public:
struct Args {
GemmDesc gemm_desc;
GemmConfig gemm_config;
LaunchArgs launch_args;
// TODO: move into descriptor
const std::optional<std::string> epilogue_type;
// TODO: move into descriptor
int gran_k_a, gran_k_b;
void* grouped_layout;
CUtensorMap tensor_map_a;
CUtensorMap tensor_map_b;
CUtensorMap tensor_map_sfa;
CUtensorMap tensor_map_sfb;
CUtensorMap tensor_map_cd;
};
static std::string generate_impl(const Args& args) {
// TODO: rename files
return fmt::format(R"(
#include <deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm100_fp8_fp4_gemm_1d1d_impl<
{}, {},
{}, {},
{}, {}, {},
{}, {}, {},
{},
{}, {}, {},
{},
{}, {},
{}, {},
{},
{},
{}, {},
{}, {}, {},
{}
>);
}};
)",
to_string(args.gemm_desc.major_a), to_string(args.gemm_desc.major_b),
args.gran_k_a, args.gran_k_b,
get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims),
args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k,
args.gemm_desc.num_groups,
args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode, args.gemm_config.storage_config.swizzle_cd_mode,
args.gemm_config.pipeline_config.num_stages,
args.gemm_config.launch_config.num_non_epilogue_threads, args.gemm_config.launch_config.num_epilogue_threads,
args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1,
args.gemm_config.launch_config.num_sms,
args.gemm_config.layout.swap_ab,
to_string(args.gemm_desc.gemm_type), args.gemm_desc.with_accumulation,
to_string(args.gemm_desc.a_dtype), to_string(args.gemm_desc.b_dtype), to_string(args.gemm_desc.cd_dtype),
get_default_epilogue_type(args.epilogue_type));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
// TODO: optimize `args` copy
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.grouped_layout, args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
args.tensor_map_a, args.tensor_map_b,
args.tensor_map_sfa, args.tensor_map_sfb,
args.tensor_map_cd));
}
};
static void sm100_fp8_fp4_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const std::optional<torch::Tensor>& c,
const torch::Tensor& d,
const int& m, const int& n, const int& k,
const int& gran_k_a, const int& gran_k_b,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims,
const std::optional<std::string>& epilogue_type = std::nullopt) {
const auto desc = GemmDesc {
.gemm_type = GemmType::Normal,
.kernel_type = KernelType::Kernel1D1D,
.m = m, .n = n, .k = k, .num_groups = 1,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(),
.compiled_dims = compiled_dims
};
const auto config = get_best_config<SM100ArchSpec>(desc);
const auto cd = c.value_or(d);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, static_cast<int>(d.size(-1)),
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, gran_k_a, 1, 0);
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.layout.block_n, gran_k_b, 1, 0);
// Launch
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = epilogue_type,
.gran_k_a = gran_k_a,
.gran_k_b = gran_k_b,
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd
};
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
const auto runtime = compiler->build("sm100_fp8_fp4_gemm_1d1d", code);
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
}
static void sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const torch::Tensor& d,
const torch::Tensor& grouped_layout,
const int& num_groups, const int& m, const int& n, const int& k,
const int& gran_k_a, const int& gran_k_b,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims,
const bool& use_psum_layout,
const std::optional<int>& expected_m_for_psum_layout) {
const auto gemm_type = use_psum_layout ?
GemmType::MGroupedContiguousWithPsumLayout : GemmType::MGroupedContiguous;
// Only psum layout can use expected m
if (expected_m_for_psum_layout)
DG_HOST_ASSERT(use_psum_layout);
// NOTES: If actual M is dynamic, estimate config via `num_groups` and `expected_m`.
// Otherwise, treat the contiguous layout as a whole.
const auto desc = GemmDesc {
.gemm_type = gemm_type,
.kernel_type = KernelType::Kernel1D1D,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(),
.compiled_dims = compiled_dims,
.expected_m = expected_m_for_psum_layout.value_or(m),
.expected_n = n, .expected_k = k,
.expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1
};
const auto config = get_best_config<SM100ArchSpec>(desc);
// Create tensor descriptors
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, gran_k_a, 1, 0);
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.layout.block_n, gran_k_b, num_groups, 0);
// Launch kernel
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = std::nullopt,
.gran_k_a = gran_k_a,
.gran_k_b = gran_k_b,
.grouped_layout = grouped_layout.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd
};
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
const auto runtime = compiler->build("sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d", code);
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
}
static void sm100_m_grouped_fp8_fp4_gemm_masked_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const torch::Tensor& d,
const torch::Tensor& masked_m,
const int& num_groups, const int& m, const int& n, const int& k,
const int& expected_m,
const int& gran_k_a, const int& gran_k_b,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto desc = GemmDesc {
.gemm_type = GemmType::MGroupedMasked,
.kernel_type = KernelType::Kernel1D1D,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(),
.compiled_dims = compiled_dims,
.expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM100ArchSpec>(desc);
// Create tensor descriptors
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), num_groups,
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, gran_k_a, num_groups, 0);
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.layout.block_n, gran_k_b, num_groups, 0);
// Launch kernel
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = std::nullopt,
.gran_k_a = gran_k_a,
.gran_k_b = gran_k_b,
.grouped_layout = masked_m.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd
};
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
const auto runtime = compiler->build("sm100_m_grouped_fp8_fp4_gemm_masked_1d1d", code);
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
}
static void sm100_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const std::optional<torch::Tensor>& c,
const torch::Tensor& d,
const int& m, const int& n,
const std::vector<int>& ks, const torch::Tensor& ks_tensor,
const int& gran_k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
DG_HOST_ASSERT(major_a == cute::UMMA::Major::MN and major_b == cute::UMMA::Major::MN);
DG_HOST_ASSERT(gran_k == 32 or gran_k == 128);
const int gran_k_a = gran_k;
const int gran_k_b = gran_k;
int sum_k = 0, sum_sf_k = 0;
for (const auto k: ks) {
sum_k += k, sum_sf_k += ceil_div(k, gran_k * 4);
DG_HOST_ASSERT(k % gran_k == 0);
}
const auto num_groups = static_cast<int>(ks.size());
// Get config using max K for better performance
const auto max_k = *std::max_element(ks.begin(), ks.end());
const auto desc = GemmDesc {
.gemm_type = GemmType::KGroupedContiguous,
.kernel_type = KernelType::Kernel1D1D,
.m = m, .n = n, .k = sum_k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(),
.compiled_dims = compiled_dims,
.expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM100ArchSpec>(desc);
// Create tensor descriptors
const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(0)), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::MN, b, n, sum_k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(0)), 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(1)), num_groups,
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, sum_sf_k * gran_k_a * 4,
config.layout.block_m, gran_k_a, 1, 0);
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, sum_sf_k * gran_k_b * 4,
config.layout.block_n, gran_k_b, 1, 0);
// Launch kernel
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = std::nullopt,
.gran_k_a = gran_k_a,
.gran_k_b = gran_k_b,
.grouped_layout = ks_tensor.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd
};
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
const auto runtime = compiler->build("sm100_k_grouped_fp8_gemm_1d1d", code);
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
}
static void sm100_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const std::optional<torch::Tensor>& c,
const torch::Tensor& d,
const int& batch_size, const int& m, const int& n, const int& k,
const int& gran_k_a, const int& gran_k_b,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto desc = GemmDesc {
.gemm_type = GemmType::Batched,
.kernel_type = KernelType::Kernel1D1D,
.m = m, .n = n, .k = k, .num_groups = batch_size,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(),
.compiled_dims = compiled_dims
};
const auto config = get_best_config<SM100ArchSpec>(desc);
const int load_block_m = config.storage_config.load_block_m;
const auto [inner_dim_a, outer_dim_a] = get_inner_outer_dims(major_a, k, m);
const auto [inner_block_a, outer_block_a] = get_inner_outer_dims(major_a, config.layout.block_k, load_block_m);
const auto tensor_map_a = make_tma_3d_desc(a, inner_dim_a, outer_dim_a, batch_size,
inner_block_a, outer_block_a, 1,
a.stride(major_a == cute::UMMA::Major::K ? 1 : 2),
a.stride(0),
config.storage_config.swizzle_a_mode);
const int load_block_n = config.storage_config.load_block_n;
const auto [inner_dim_b, outer_dim_b] = get_inner_outer_dims(major_b, k, n);
const auto [inner_block_b, outer_block_b] = get_inner_outer_dims(major_b, config.layout.block_k, load_block_n);
const auto tensor_map_b = make_tma_3d_desc(b, inner_dim_b, outer_dim_b, batch_size,
inner_block_b, outer_block_b, 1,
b.stride(major_b == cute::UMMA::Major::K ? 1 : 2),
b.stride(0),
config.storage_config.swizzle_b_mode);
const int store_block_m = config.storage_config.store_block_m;
const int store_block_n = config.storage_config.store_block_n;
const auto tensor_map_cd = make_tma_3d_desc(d, n, m, batch_size,
store_block_n, store_block_m, 1,
d.stride(1), d.stride(0),
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, gran_k_a, batch_size, 0);
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.layout.block_n, gran_k_b, batch_size, 0);
// Launch
const SM100FP8FP4Gemm1D1DRuntime::Args args = {
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = std::nullopt,
.gran_k_a = gran_k_a,
.gran_k_b = gran_k_b,
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd
};
const auto code = SM100FP8FP4Gemm1D1DRuntime::generate(args);
const auto runtime = compiler->build("sm100_fp8_gemm_1d1d", code);
SM100FP8FP4Gemm1D1DRuntime::launch(runtime, args);
}
} // namespace deep_gemm

210
csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp Normal file
View File

@@ -0,0 +1,210 @@
#pragma once
#include <torch/python.h>
#include "../../jit/compiler.hpp"
#include "../../jit/kernel_runtime.hpp"
#include "../../utils/exception.hpp"
#include "../../utils/format.hpp"
#include "runtime_utils.hpp"
#include <deep_gemm/layout/mega_moe.cuh>
#include <deep_gemm/layout/sym_buffer.cuh>
#include "../heuristics/mega_moe.hpp"
namespace deep_gemm {
class SM100FP8FP4MegaMoERuntime final : public LaunchRuntime<SM100FP8FP4MegaMoERuntime> {
public:
struct Args {
// Templated arguments
int num_max_tokens_per_rank;
int hidden, intermediate_hidden;
int num_experts, num_topk;
int num_ranks;
float activation_clamp;
bool fast_math;
MegaMoEConfig config;
// Runtime arguments
void* y;
int num_tokens;
layout::SymBuffer<> sym_buffer_ptrs;
// Tensormap
CUtensorMap tensor_map_l1_acts;
CUtensorMap tensor_map_l1_acts_sf;
CUtensorMap tensor_map_l1_weights;
CUtensorMap tensor_map_l1_weights_sf;
CUtensorMap tensor_map_l1_output;
CUtensorMap tensor_map_l2_acts;
CUtensorMap tensor_map_l2_acts_sf;
CUtensorMap tensor_map_l2_weights;
CUtensorMap tensor_map_l2_weights_sf;
// Launch configs
LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#include <deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm100_fp8_fp4_mega_moe_impl<
{},
{}, {},
{}, {},
{},
{}, {}, {},
{},
{}, {},
{},
{},
{},
{}, {}, {},
{}, {},
{},
{}
>);
}};
)", args.num_max_tokens_per_rank,
args.hidden, args.intermediate_hidden,
args.num_experts, args.num_topk,
args.config.num_experts_per_wave,
args.config.block_m, args.config.block_n, args.config.block_k,
args.config.store_block_m,
args.config.sf_block_m, args.config.sf_block_n,
args.config.num_max_pool_tokens,
args.config.num_padded_sf_pool_tokens,
args.config.num_stages,
args.config.num_dispatch_threads, args.config.num_non_epilogue_threads, args.config.num_epilogue_threads,
args.launch_args.grid_dim.first, args.num_ranks,
to_string(args.activation_clamp),
args.fast_math ? "true" : "false");
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
// TODO: optimize `args` copy
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.y,
args.num_tokens,
args.sym_buffer_ptrs,
args.tensor_map_l1_acts,
args.tensor_map_l1_acts_sf,
args.tensor_map_l1_weights,
args.tensor_map_l1_weights_sf,
args.tensor_map_l1_output,
args.tensor_map_l2_acts,
args.tensor_map_l2_acts_sf,
args.tensor_map_l2_weights,
args.tensor_map_l2_weights_sf
));
}
};
static void sm100_fp8_fp4_mega_moe(
const torch::Tensor& y,
const torch::Tensor& l1_acts, const torch::Tensor& l1_acts_sf,
const torch::Tensor& l2_acts, const torch::Tensor& l2_acts_sf,
const torch::Tensor& l1_weights, const torch::Tensor& l2_weights,
const torch::Tensor& l1_weights_sf, const torch::Tensor& l2_weights_sf,
const std::vector<int64_t>& sym_buffer_ptrs,
const int& rank_idx, const int& num_max_tokens_per_rank,
const int& num_experts_per_rank,
const int& num_tokens, const int& num_topk,
const int& hidden, const int& intermediate_hidden,
const float& activation_clamp,
const bool& fast_math
) {
const auto num_ranks = static_cast<int>(sym_buffer_ptrs.size());
const auto num_experts = num_experts_per_rank * num_ranks;
// Heuristics
const auto config = get_mega_moe_config(
num_ranks, num_experts, num_experts_per_rank,
num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden);
// Make tensormap
constexpr int kGranK = 32;
const auto tensor_map_l1_acts = make_tma_2d_desc(l1_acts,
hidden, config.num_max_pool_tokens,
config.block_k, config.load_block_m,
static_cast<int>(l1_acts.stride(-2)),
config.swizzle_acts_mode);
const auto tensor_map_l1_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_acts_sf,
config.num_padded_sf_pool_tokens, hidden,
config.sf_block_m, kGranK,
1, 0);
const auto tensor_map_l1_weights = make_tma_2d_desc(l1_weights,
hidden, num_experts_per_rank * intermediate_hidden * 2,
config.block_k, config.load_block_n,
static_cast<int>(l1_weights.stride(-2)),
config.swizzle_weights_mode);
const auto tensor_map_l1_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l1_weights_sf,
intermediate_hidden * 2, hidden,
config.block_n, kGranK,
num_experts_per_rank, 0);
// NOTES: L1 output and L2 activations are essentially the same tensor.
// Post-SwiGLU output has half the N width (`BLOCK_N / 2` per input tile),
// so the swizzle mode is also halved (128 -> 64).
const auto tensor_map_l1_output = make_tma_2d_desc(l2_acts,
intermediate_hidden, config.num_max_pool_tokens,
config.block_n / 2, config.store_block_m,
static_cast<int>(l2_acts.stride(-2)),
config.swizzle_acts_mode / 2);
const auto tensor_map_l2_acts = make_tma_2d_desc(l2_acts,
intermediate_hidden, config.num_max_pool_tokens,
config.block_k, config.load_block_m,
static_cast<int>(l2_acts.stride(-2)),
config.swizzle_acts_mode);
const auto tensor_map_l2_acts_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_acts_sf,
config.num_padded_sf_pool_tokens, intermediate_hidden,
config.sf_block_m, kGranK,
1, 0);
const auto tensor_map_l2_weights = make_tma_2d_desc(l2_weights,
intermediate_hidden, num_experts_per_rank * hidden,
config.block_k, config.load_block_n,
static_cast<int>(l2_weights.stride(-2)),
config.swizzle_weights_mode);
const auto tensor_map_l2_weights_sf = make_tma_sf_desc(cute::UMMA::Major::MN, l2_weights_sf,
hidden, intermediate_hidden,
config.block_n, kGranK,
num_experts_per_rank, 0);
// Launch
const auto num_sms = device_runtime->get_num_sms();
const SM100FP8FP4MegaMoERuntime::Args args = {
.num_max_tokens_per_rank = num_max_tokens_per_rank,
.hidden = hidden, .intermediate_hidden = intermediate_hidden,
.num_experts = num_experts, .num_topk = num_topk,
.num_ranks = num_ranks,
.activation_clamp = activation_clamp,
.fast_math = fast_math,
.config = config,
.y = y.data_ptr(),
.num_tokens = num_tokens,
.sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx),
.tensor_map_l1_acts = tensor_map_l1_acts,
.tensor_map_l1_acts_sf = tensor_map_l1_acts_sf,
.tensor_map_l1_weights = tensor_map_l1_weights,
.tensor_map_l1_weights_sf = tensor_map_l1_weights_sf,
.tensor_map_l1_output = tensor_map_l1_output,
.tensor_map_l2_acts = tensor_map_l2_acts,
.tensor_map_l2_acts_sf = tensor_map_l2_acts_sf,
.tensor_map_l2_weights = tensor_map_l2_weights,
.tensor_map_l2_weights_sf = tensor_map_l2_weights_sf,
.launch_args = LaunchArgs(num_sms,
config.num_dispatch_threads + config.num_non_epilogue_threads + config.num_epilogue_threads,
config.smem_size, 2)
};
const auto code = SM100FP8FP4MegaMoERuntime::generate(args);
const auto runtime = compiler->build("sm100_fp8_fp4_mega_moe", code);
SM100FP8FP4MegaMoERuntime::launch(runtime, args);
}
} // namespace deep_gemm

36
csrc/jit_kernels/impls/sm100_tf32_hc_prenorm_gemm.hpp
View File

@@ -79,21 +79,21 @@ static void sm100_tf32_hc_prenorm_gemm(const torch::Tensor& a,
DG_HOST_ASSERT(n <= 128 and n % 8 == 0);
DG_HOST_ASSERT(k % block_k == 0);
const auto& swizzle_cd_mode = get_swizzle_mode(block_n, sizeof(float));
const auto& tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k,
block_m, block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, a.element_size()), 0,
true);
const auto& tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b, n, k,
block_n, block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, b.element_size()), 0,
true);
const auto& tensor_map_d = num_splits == 1 ? make_tma_cd_desc(d, m, n,
block_m, block_n,
static_cast<int>(d.stride(-2)), 1,
swizzle_cd_mode)
const auto swizzle_cd_mode = get_swizzle_mode(block_n, sizeof(float));
const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k,
block_m, block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, a.element_size()), 0,
true);
const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b, n, k,
block_n, block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, b.element_size()), 0,
true);
const auto tensor_map_d = num_splits == 1 ? make_tma_cd_desc(d, m, n,
block_m, block_n,
static_cast<int>(d.stride(-2)), 1,
swizzle_cd_mode)
: make_tma_3d_desc(d, n, m, num_splits,
block_n, block_m, 1,
static_cast<int>(d.stride(-2)),
@@ -135,14 +135,14 @@ static void sm100_tf32_hc_prenorm_gemm(const torch::Tensor& a,
.num_stages = num_stages,
.num_mma_threads = num_mma_threads,
.num_cast_and_reduce_threads = num_cast_and_reduce_threads,
.launch_args = LaunchArgs(num_splits * ceil_div(m, block_m), num_mma_threads + num_cast_and_reduce_threads, smem_size, 1),
.launch_args = LaunchArgs(num_splits * ceil_div(m, block_m), num_mma_threads + num_cast_and_reduce_threads, smem_size),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_d = tensor_map_d,
.sqr_sum = sqr_sum.data_ptr<float>()
};
const auto& code = SM100BF16HCPrenormGemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_tf32_hc_prenorm_gemm", code);
const auto code = SM100BF16HCPrenormGemmRuntime::generate(args);
const auto runtime = compiler->build("sm100_tf32_hc_prenorm_gemm", code);
SM100BF16HCPrenormGemmRuntime::launch(runtime, args);
}

418
csrc/jit_kernels/impls/sm90_bf16_gemm.hpp
View File

@@ -14,9 +14,7 @@ namespace deep_gemm {
class SM90BF16GemmRuntime final: public LaunchRuntime<SM90BF16GemmRuntime> {
public:
struct Args {
int m, n, k, num_groups;
const std::string& compiled_dims;
GemmDesc gemm_desc;
GemmConfig gemm_config;
LaunchArgs launch_args;
@@ -49,24 +47,29 @@ static void __instantiate_kernel() {{
}};
)",
// TODO: add CD dtype
to_string(args.gemm_config.major_a), to_string(args.gemm_config.major_b),
get_compiled_dim(args.m, 'm', args.compiled_dims), get_compiled_dim(args.n, 'n', args.compiled_dims), get_compiled_dim(args.k, 'k', args.compiled_dims),
args.num_groups,
args.gemm_config.block_m, args.gemm_config.block_n, args.gemm_config.block_k,
args.gemm_config.smem_config.swizzle_a_mode, args.gemm_config.smem_config.swizzle_b_mode, args.gemm_config.smem_config.swizzle_cd_mode,
args.gemm_config.num_stages,
args.gemm_config.thread_config.num_tma_threads, args.gemm_config.thread_config.num_math_threads,
args.gemm_config.multicast_config.num_multicast, args.gemm_config.multicast_config.is_multicast_on_a,
args.gemm_config.num_sms,
to_string(args.gemm_config.gemm_type), args.gemm_config.with_accumulation,
to_string(args.gemm_config.cd_dtype));
to_string(args.gemm_desc.major_a), to_string(args.gemm_desc.major_b),
get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims),
args.gemm_desc.num_groups,
args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k,
args.gemm_config.storage_config.swizzle_a_mode,
args.gemm_config.storage_config.swizzle_b_mode,
args.gemm_config.storage_config.swizzle_cd_mode,
args.gemm_config.pipeline_config.num_stages,
args.gemm_config.launch_config.num_tma_threads, args.gemm_config.launch_config.num_math_threads,
// TODO: refactor with cluster M/N
args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1,
args.gemm_config.launch_config.num_sms,
to_string(args.gemm_desc.gemm_type), args.gemm_desc.with_accumulation,
to_string(args.gemm_desc.cd_dtype));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
// TODO: optimize `args` copy
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.grouped_layout,
args.m, args.n, args.k,
args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
args.tensor_map_a, args.tensor_map_b,
args.tensor_map_cd));
}
@@ -79,46 +82,50 @@ static void sm90_bf16_gemm(const torch::Tensor& a,
const int& m, const int& n, const int& k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::Normal, KernelType::KernelNoSF,
m, n, k, 1, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Normal,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = k, .num_groups = 1,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
config.storage_config.swizzle_cd_mode);
// Launch
const SM90BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = 1,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_gemm", code);
const auto code = SM90BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm90_bf16_gemm", code);
SM90BF16GemmRuntime::launch(runtime, args);
}
@@ -128,51 +135,67 @@ static void sm90_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a,
const torch::Tensor& m_indices,
const int& num_groups, const int& m, const int& n, const int& k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const std::string& compiled_dims,
const bool& use_psum_layout,
const std::optional<int>& expected_m_for_psum_layout) {
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K);
DG_HOST_ASSERT(k % 64 == 0);
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::MGroupedContiguous, KernelType::KernelNoSF,
m, n, k, 1, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), false,
device_runtime->get_num_sms());
const auto gemm_type = use_psum_layout ?
GemmType::MGroupedContiguousWithPsumLayout : GemmType::MGroupedContiguous;
// Only psum layout can use expected m
if (expected_m_for_psum_layout)
DG_HOST_ASSERT(use_psum_layout);
const auto desc = GemmDesc {
.gemm_type = gemm_type,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = expected_m_for_psum_layout.value_or(m),
.expected_n = n, .expected_k = k,
.expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
config.storage_config.swizzle_cd_mode);
// Launch
const SM90BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = m_indices.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_m_grouped_bf16_gemm_contiguous", code);
const auto code = SM90BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm90_m_grouped_bf16_gemm_contiguous", code);
SM90BF16GemmRuntime::launch(runtime, args);
}
@@ -188,46 +211,51 @@ static void sm90_bf16_m_grouped_gemm_masked(const torch::Tensor& a,
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
DG_HOST_ASSERT(k % 64 == 0);
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::MGroupedMasked, KernelType::KernelNoSF,
expected_m, n, k, num_groups, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), false,
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::MGroupedMasked,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = expected_m, .expected_n = 0, .expected_k = 0, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), num_groups,
config.smem_config.swizzle_cd_mode);
config.storage_config.swizzle_cd_mode);
// Launch
const SM90BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = masked_m.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_m_grouped_gemm_masked", code);
const auto code = SM90BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm90_bf16_m_grouped_gemm_masked", code);
SM90BF16GemmRuntime::launch(runtime, args);
}
@@ -242,54 +270,59 @@ static void sm90_bf16_k_grouped_gemm(const torch::Tensor& a,
DG_HOST_ASSERT(major_a == cute::UMMA::Major::MN and major_b == cute::UMMA::Major::MN);
int sum_k = 0;
for (const auto& k: ks) {
for (const auto k: ks) {
sum_k += k;
DG_HOST_ASSERT(k % 128 == 0);
}
const auto& num_groups = static_cast<int>(ks.size());
const auto num_groups = static_cast<int>(ks.size());
// Get config using max K for better performance
const auto& max_k = *std::max_element(ks.begin(), ks.end());
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::KGroupedContiguous, KernelType::KernelNoSF,
m, n, max_k, num_groups, cute::UMMA::Major::MN, cute::UMMA::Major::MN,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto max_k = *std::max_element(ks.begin(), ks.end());
const auto desc = GemmDesc {
.gemm_type = GemmType::KGroupedContiguous,
.kernel_type = KernelType::KernelNoSF,
.m = m, .n = n, .k = sum_k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Create tensor descriptors
const auto& tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(0)), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(cute::UMMA::Major::MN, b, n, sum_k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(0)), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(1)), num_groups,
config.smem_config.swizzle_cd_mode);
const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(0)), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::MN, b, n, sum_k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(0)), 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(1)), num_groups,
config.storage_config.swizzle_cd_mode);
// Launch kernel
const SM90BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = sum_k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = ks_tensor.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_k_grouped_gemm", code);
const auto code = SM90BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm90_bf16_k_grouped_gemm", code);
SM90BF16GemmRuntime::launch(runtime, args);
}
@@ -298,45 +331,50 @@ static void sm90_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a,
const torch::Tensor& tensor_d,
const int& b, const int& h, const int& r, const int& d,
const std::string& compiled_dims = "nk") {
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::Batched, KernelType::KernelNoSF,
b, d, r, h, cute::UMMA::Major::K, cute::UMMA::Major::K,
tensor_a.scalar_type(), tensor_b.scalar_type(),
tensor_d.scalar_type(), false,
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Batched,
.kernel_type = KernelType::KernelNoSF,
.m = b, .n = d, .k = r, .num_groups = h,
.a_dtype = tensor_a.scalar_type(), .b_dtype = tensor_b.scalar_type(),
.cd_dtype = tensor_d.scalar_type(),
.major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::K,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM90ArchSpec>(desc);
const int& load_block_m = SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& tensor_map_a = make_tma_3d_desc(tensor_a, r, b, h,
config.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.smem_config.swizzle_a_mode);
const int& load_block_n = SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
config.block_k, load_block_n, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.smem_config.swizzle_b_mode);
const int& store_block_m = SM90ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM90ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_cd = make_tma_3d_desc(tensor_d, d, b, h,
const int load_block_m = config.storage_config.load_block_m;
const auto tensor_map_a = make_tma_3d_desc(tensor_a, r, b, h,
config.layout.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.storage_config.swizzle_a_mode);
const int load_block_n = config.storage_config.load_block_n;
const auto tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
config.layout.block_k, load_block_n, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.storage_config.swizzle_b_mode);
const int store_block_m = config.storage_config.store_block_m;
const int store_block_n = config.storage_config.store_block_n;
const auto tensor_map_cd = make_tma_3d_desc(tensor_d, d, b, h,
store_block_n, store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.smem_config.swizzle_cd_mode);
config.storage_config.swizzle_cd_mode);
// Launch
const SM90BF16GemmRuntime::Args& args = {
.m = b, .n = d, .k = r,
.num_groups = h,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_bhr_hdr_bhd", code);
const auto code = SM90BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm90_bf16_bhr_hdr_bhd", code);
SM90BF16GemmRuntime::launch(runtime, args);
}
@@ -345,45 +383,49 @@ static void sm90_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a,
const torch::Tensor& tensor_d,
const int& b, const int& h, const int& r, const int& d,
const std::string& compiled_dims = "nk") {
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::Batched, KernelType::KernelNoSF,
b, r, d, h, cute::UMMA::Major::K, cute::UMMA::Major::MN,
tensor_a.scalar_type(), tensor_b.scalar_type(),
tensor_d.scalar_type(), false,
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Batched,
.kernel_type = KernelType::KernelNoSF,
.m = b, .n = r, .k = d, .num_groups = h,
.a_dtype = tensor_a.scalar_type(), .b_dtype = tensor_b.scalar_type(),
.cd_dtype = tensor_d.scalar_type(),
.major_a = cute::UMMA::Major::K, .major_b = cute::UMMA::Major::MN,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM90ArchSpec>(desc);
const int& load_block_m = SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& tensor_map_a = make_tma_3d_desc(tensor_a, d, b, h,
config.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.smem_config.swizzle_a_mode);
const int& load_block_n = SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
load_block_n, config.block_k, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.smem_config.swizzle_b_mode);
const int& store_block_m = SM90ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM90ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_cd = make_tma_3d_desc(tensor_d, r, b, h,
const int load_block_m = config.storage_config.load_block_m;
const auto tensor_map_a = make_tma_3d_desc(tensor_a, d, b, h,
config.layout.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.storage_config.swizzle_a_mode);
const int load_block_n = config.storage_config.load_block_n;
const auto tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
load_block_n, config.layout.block_k, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.storage_config.swizzle_b_mode);
const int store_block_m = config.storage_config.store_block_m;
const int store_block_n = config.storage_config.store_block_n;
const auto tensor_map_cd = make_tma_3d_desc(tensor_d, r, b, h,
store_block_n, store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.smem_config.swizzle_cd_mode);
config.storage_config.swizzle_cd_mode);
// Launch
const SM90BF16GemmRuntime::Args& args = {
.m = b, .n = r, .k = d,
.num_groups = h,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_bhd_hdr_bhr", code);
const auto code = SM90BF16GemmRuntime::generate(args);
const auto runtime = compiler->build("sm90_bf16_bhd_hdr_bhr", code);
SM90BF16GemmRuntime::launch(runtime, args);
}

14
csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp
View File

@@ -84,9 +84,9 @@ static void sm90_bmn_bnk_mn_gemm(const torch::Tensor &a,
// Select best number of stages
int num_stages = 4, smem_size = 0;
while (true) {
const int& smem_a_per_stage = block_m * block_k * sizeof(cutlass::bfloat16_t);
const int& smem_b_per_stage = block_n * block_k * sizeof(cutlass::bfloat16_t);
const int& smem_barrier = SM90ArchSpec::get_barrier_smem_size(num_stages);
const int smem_a_per_stage = block_m * block_k * sizeof(cutlass::bfloat16_t);
const int smem_b_per_stage = block_n * block_k * sizeof(cutlass::bfloat16_t);
const int smem_barrier = num_stages * 8 * 2;
smem_size = 0;
smem_size += (smem_a_per_stage + smem_b_per_stage) * num_stages;
@@ -108,8 +108,8 @@ static void sm90_bmn_bnk_mn_gemm(const torch::Tensor &a,
num_stages, smem_size, swizzle_ab_mode);
}
const auto& tensor_map_a = make_tma_2d_desc(a, k, s * m, block_k, block_m, k, swizzle_ab_mode);
const auto& tensor_map_b = make_tma_2d_desc(b, k, s * n, block_k, block_n, k, swizzle_ab_mode);
const auto tensor_map_a = make_tma_2d_desc(a, k, s * m, block_k, block_m, k, swizzle_ab_mode);
const auto tensor_map_b = make_tma_2d_desc(b, k, s * n, block_k, block_n, k, swizzle_ab_mode);
const SM90BmkBnkMnRuntime::Args& args = {
.s = s, .m = m, .n = n, .k = k,
@@ -123,8 +123,8 @@ static void sm90_bmn_bnk_mn_gemm(const torch::Tensor &a,
.tensor_map_b = tensor_map_b,
.d = d.data_ptr<float>()
};
const auto& code = SM90BmkBnkMnRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bmn_bnk_mn_gemm", code);
const auto code = SM90BmkBnkMnRuntime::generate(args);
const auto runtime = compiler->build("sm90_bmn_bnk_mn_gemm", code);
SM90BmkBnkMnRuntime::launch(runtime, args);
}

183
csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp
View File

@@ -15,9 +15,7 @@ namespace deep_gemm {
class SM90FP8Gemm1D1DRuntime final: public LaunchRuntime<SM90FP8Gemm1D1DRuntime> {
public:
struct Args {
int m, n, k, num_groups;
const std::string& compiled_dims;
GemmDesc gemm_desc;
GemmConfig gemm_config;
LaunchArgs launch_args;
@@ -52,15 +50,17 @@ static void __instantiate_kernel() {{
>);
}};
)",
get_compiled_dim(args.m, 'm', args.compiled_dims), get_compiled_dim(args.n, 'n', args.compiled_dims), get_compiled_dim(args.k, 'k', args.compiled_dims),
args.num_groups,
args.gemm_config.block_m, args.gemm_config.block_n, args.gemm_config.block_k,
args.gemm_config.smem_config.swizzle_a_mode, args.gemm_config.smem_config.swizzle_b_mode,
args.gemm_config.num_stages,
args.gemm_config.thread_config.num_tma_threads, args.gemm_config.thread_config.num_math_threads,
args.gemm_config.multicast_config.num_multicast, args.gemm_config.multicast_config.is_multicast_on_a,
args.gemm_config.num_sms, to_string(args.gemm_config.gemm_type),
to_string(args.gemm_config.cd_dtype));
get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims),
args.gemm_desc.num_groups,
args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k,
args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode,
args.gemm_config.pipeline_config.num_stages,
args.gemm_config.launch_config.num_tma_threads, args.gemm_config.launch_config.num_math_threads,
args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1,
args.gemm_config.launch_config.num_sms, to_string(args.gemm_desc.gemm_type),
to_string(args.gemm_desc.cd_dtype));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
@@ -68,7 +68,7 @@ static void __instantiate_kernel() {{
args.gmem_a_ptr, args.gmem_b_ptr,
args.grouped_layout,
args.tensor_map_buffer,
args.m, args.n, args.k,
args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
args.tensor_map_a_base, args.tensor_map_b_base,
args.tensor_map_sfa, args.tensor_map_sfb,
args.tensor_map_cd));
@@ -85,44 +85,48 @@ static void sm90_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
DG_HOST_ASSERT(c.has_value() and d.scalar_type() == torch::kFloat);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::Normal, KernelType::Kernel1D1D,
m, n, k, 1, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Normal,
.kernel_type = KernelType::Kernel1D1D,
.m = m, .n = n, .k = k, .num_groups = 1,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
DG_HOST_ASSERT(config.smem_config.swizzle_a_mode == config.block_k);
DG_HOST_ASSERT(config.smem_config.swizzle_b_mode == config.block_k);
DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k);
DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k);
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k, k, 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k, k, 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.block_m, config.block_k, 1, 0);
const auto& tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.block_n, config.block_k, 1, 0);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m, true),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
0);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k, k, 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k, k, 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, config.layout.block_k, 1, 0);
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.layout.block_n, config.layout.block_k, 1, 0);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
0);
// Launch
const SM90FP8Gemm1D1DRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = 1,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.gmem_a_ptr = nullptr,
.gmem_b_ptr = nullptr,
.grouped_layout = nullptr,
@@ -133,8 +137,8 @@ static void sm90_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90FP8Gemm1D1DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_gemm_1d1d", code);
const auto code = SM90FP8Gemm1D1DRuntime::generate(args);
const auto runtime = compiler->build("sm90_fp8_gemm_1d1d", code);
SM90FP8Gemm1D1DRuntime::launch(runtime, args);
}
@@ -151,54 +155,61 @@ static void sm90_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Te
DG_HOST_ASSERT(c.has_value() and d.scalar_type() == torch::kFloat);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
// Get config using max K for better performance
const auto& num_groups = static_cast<int>(ks.size());
const auto& max_k = *std::max_element(ks.begin(), ks.end());
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::KGroupedContiguous, KernelType::Kernel1D1D,
m, n, max_k, num_groups, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
// Requires no TMA splits
DG_HOST_ASSERT(config.smem_config.swizzle_a_mode == config.block_k);
DG_HOST_ASSERT(config.smem_config.swizzle_b_mode == config.block_k);
int first_k = 0, sum_k = 0, sum_sf_k = 0;
// TODO: refactor with the mk alignment function
const auto num_groups = static_cast<int>(ks.size());
int first_k = 0, sum_k = 0, sum_sf_k = 0, max_k = 0;
for (int i = 0; i < num_groups; ++ i) {
if (first_k == 0 and ks[i] != 0)
first_k = ks[i];
sum_k += ks[i], sum_sf_k += ceil_div(ks[i], 128);
max_k = std::max(max_k, ks[i]);
DG_HOST_ASSERT(ks[i] % 128 == 0);
}
const auto& tensor_map_a_base = make_tma_a_desc(major_a, a, m, first_k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k, first_k, 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b_base = make_tma_b_desc(major_b, b, n, first_k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k, first_k, 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, sum_sf_k * 128,
config.block_m, config.block_k, 1, 0);
const auto& tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, sum_sf_k * 128,
config.block_n, config.block_k, 1, 0);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m, true),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), num_groups,
config.smem_config.swizzle_cd_mode);
// Get config using max K for better performance
const auto desc = GemmDesc {
.gemm_type = GemmType::KGroupedContiguous,
.kernel_type = KernelType::Kernel1D1D,
.m = m, .n = n, .k = sum_k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = m, .expected_n = n, .expected_k = max_k, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k);
DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k);
const auto tensor_map_a_base = make_tma_a_desc(major_a, a, m, first_k,
config.storage_config.load_block_m,
config.layout.block_k, first_k, 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b_base = make_tma_b_desc(major_b, b, n, first_k,
config.storage_config.load_block_n,
config.layout.block_k, first_k, 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, sum_sf_k * 128,
config.layout.block_m, config.layout.block_k, 1, 0);
const auto tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, sum_sf_k * 128,
config.layout.block_n, config.layout.block_k, 1, 0);
const auto tensor_map_cd = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), num_groups,
config.storage_config.swizzle_cd_mode);
// Launch
const SM90FP8Gemm1D1DRuntime::Args& args = {
.m = m, .n = n, .k = sum_k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.gmem_a_ptr = a.data_ptr(),
.gmem_b_ptr = b.data_ptr(),
.grouped_layout = ks_tensor.data_ptr(),
@@ -209,8 +220,8 @@ static void sm90_k_grouped_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Te
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90FP8Gemm1D1DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_gemm_1d1d", code);
const auto code = SM90FP8Gemm1D1DRuntime::generate(args);
const auto runtime = compiler->build("sm90_fp8_gemm_1d1d", code);
SM90FP8Gemm1D1DRuntime::launch(runtime, args);
}

348
csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp
View File

@@ -17,14 +17,13 @@ namespace deep_gemm {
class SM90FP8Gemm1D2DRuntime final: public LaunchRuntime<SM90FP8Gemm1D2DRuntime> {
public:
struct Args {
cute::UMMA::Major major_sfb;
int m, n, k, num_groups;
const std::string& compiled_dims;
const std::optional<std::string>& epilogue_type;
GemmDesc gemm_desc;
GemmConfig gemm_config;
LaunchArgs launch_args;
// TODO: move this into `gemm_desc`
const std::optional<std::string>& epilogue_type;
cute::UMMA::Major major_sfb;
void *sfb, *grouped_layout;
CUtensorMap tensor_map_a;
CUtensorMap tensor_map_b;
@@ -45,7 +44,7 @@ static void __instantiate_kernel() {{
{},
{}, {}, {},
{}, {}, {},
{}, {},
{},
{}, {},
{}, {},
{}, {},
@@ -55,14 +54,16 @@ static void __instantiate_kernel() {{
)",
// TODO: add CD dtype
to_string(args.major_sfb),
get_compiled_dim(args.m, 'm', args.compiled_dims), get_compiled_dim(args.n, 'n', args.compiled_dims), get_compiled_dim(args.k, 'k', args.compiled_dims),
args.num_groups,
args.gemm_config.block_m, args.gemm_config.block_n, args.gemm_config.block_k,
args.gemm_config.smem_config.swizzle_a_mode, args.gemm_config.smem_config.swizzle_b_mode, args.gemm_config.smem_config.swizzle_cd_mode,
args.gemm_config.num_stages, args.gemm_config.num_last_stages,
args.gemm_config.thread_config.num_tma_threads, args.gemm_config.thread_config.num_math_threads,
args.gemm_config.multicast_config.num_multicast, args.gemm_config.multicast_config.is_multicast_on_a,
args.gemm_config.num_sms, to_string(args.gemm_config.gemm_type),
get_compiled_dim(args.gemm_desc.m, 'm', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.n, 'n', args.gemm_desc.compiled_dims),
get_compiled_dim(args.gemm_desc.k, 'k', args.gemm_desc.compiled_dims),
args.gemm_desc.num_groups,
args.gemm_config.layout.block_m, args.gemm_config.layout.block_n, args.gemm_config.layout.block_k,
args.gemm_config.storage_config.swizzle_a_mode, args.gemm_config.storage_config.swizzle_b_mode, args.gemm_config.storage_config.swizzle_cd_mode,
args.gemm_config.pipeline_config.num_stages,
args.gemm_config.launch_config.num_tma_threads, args.gemm_config.launch_config.num_math_threads,
args.gemm_config.layout.get_cluster_size(), args.gemm_config.layout.cluster_n > 1,
args.gemm_config.launch_config.num_sms, to_string(args.gemm_desc.gemm_type),
get_default_epilogue_type(args.epilogue_type));
}
@@ -70,7 +71,7 @@ static void __instantiate_kernel() {{
// TODO: optimize `args` copy
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.sfb, args.grouped_layout,
args.m, args.n, args.k,
args.gemm_desc.m, args.gemm_desc.n, args.gemm_desc.k,
args.tensor_map_a, args.tensor_map_b,
args.tensor_map_d, args.tensor_map_sfa));
}
@@ -87,45 +88,49 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
DG_HOST_ASSERT(not c.has_value() and d.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::Normal, KernelType::Kernel1D2D,
m, n, k, 1, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Normal,
.kernel_type = KernelType::Kernel1D2D,
.m = m, .n = n, .k = k, .num_groups = 1,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
DG_HOST_ASSERT(config.smem_config.swizzle_a_mode == config.block_k);
DG_HOST_ASSERT(config.smem_config.swizzle_b_mode == config.block_k);
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, static_cast<int>(d.size(-1)),
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.block_m, config.block_k, 1, 0);
DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k);
DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.storage_config.swizzle_b_mode);
const auto tensor_map_d = make_tma_cd_desc(d, m, static_cast<int>(d.size(-1)),
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, config.layout.block_k, 1, 0);
// Launch
const SM90FP8Gemm1D2DRuntime::Args& args = {
.major_sfb = major_sfb,
.m = m, .n = n, .k = k,
.num_groups = 1,
.compiled_dims = compiled_dims,
.epilogue_type = epilogue_type,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = epilogue_type,
.major_sfb = major_sfb,
.sfb = sfb.data_ptr(),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
@@ -133,8 +138,8 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
.tensor_map_d = tensor_map_d,
.tensor_map_sfa = tensor_map_sfa,
};
const auto& code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_gemm_1d2d", code);
const auto code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto runtime = compiler->build("sm90_fp8_gemm_1d2d", code);
SM90FP8Gemm1D2DRuntime::launch(runtime, args);
}
@@ -144,49 +149,65 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons
const torch::Tensor& m_indices,
const int& num_groups, const int& m, const int& n, const int& k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, const cute::UMMA::Major& major_sfb,
const std::string& compiled_dims) {
const std::string& compiled_dims,
const bool& use_psum_layout,
const std::optional<int>& expected_m_for_psum_layout) {
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::MGroupedContiguous, KernelType::Kernel1D2D,
m, n, k, 1, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), false,
device_runtime->get_num_sms());
const auto gemm_type = use_psum_layout ?
GemmType::MGroupedContiguousWithPsumLayout : GemmType::MGroupedContiguous;
// Only psum layout can use expected m
if (expected_m_for_psum_layout)
DG_HOST_ASSERT(use_psum_layout);
const auto desc = GemmDesc {
.gemm_type = gemm_type,
.kernel_type = KernelType::Kernel1D2D,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = expected_m_for_psum_layout.value_or(m),
.expected_n = n, .expected_k = k,
.expected_num_groups = expected_m_for_psum_layout.has_value() ? num_groups : 1
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
DG_HOST_ASSERT(config.smem_config.swizzle_a_mode == config.block_k);
DG_HOST_ASSERT(config.smem_config.swizzle_b_mode == config.block_k);
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.block_m, config.block_k, 1, 0);
DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k);
DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_d = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), 1,
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, config.layout.block_k, 1, 0);
// Launch
const SM90FP8Gemm1D2DRuntime::Args& args = {
.major_sfb = major_sfb,
.m = m, .n = n, .k = k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.epilogue_type = std::nullopt,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = std::nullopt,
.major_sfb = major_sfb,
.sfb = sfb.data_ptr(),
.grouped_layout = m_indices.data_ptr(),
.tensor_map_a = tensor_map_a,
@@ -194,8 +215,8 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons
.tensor_map_d = tensor_map_d,
.tensor_map_sfa = tensor_map_sfa,
};
const auto& code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_m_grouped_fp8_gemm_contiguous_1d2d", code);
const auto code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto runtime = compiler->build("sm90_m_grouped_fp8_gemm_contiguous_1d2d", code);
SM90FP8Gemm1D2DRuntime::launch(runtime, args);
}
@@ -210,45 +231,50 @@ static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const to
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::MGroupedMasked, KernelType::Kernel1D2D,
expected_m, n, k, num_groups, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), false,
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::MGroupedMasked,
.kernel_type = KernelType::Kernel1D2D,
.m = m, .n = n, .k = k, .num_groups = num_groups,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = false,
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims,
.expected_m = expected_m, .expected_n = n, .expected_k = k, .expected_num_groups = num_groups
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
DG_HOST_ASSERT(config.smem_config.swizzle_a_mode == config.block_k);
DG_HOST_ASSERT(config.smem_config.swizzle_b_mode == config.block_k);
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), num_groups,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.block_m, config.block_k, num_groups, 0);
DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k);
DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k);
const auto tensor_map_a = make_tma_a_desc(major_a, a, m, k,
config.storage_config.load_block_m,
config.layout.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.storage_config.swizzle_a_mode);
const auto tensor_map_b = make_tma_b_desc(major_b, b, n, k,
config.storage_config.load_block_n,
config.layout.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.storage_config.swizzle_b_mode);
const auto tensor_map_d = make_tma_cd_desc(d, m, n,
config.storage_config.store_block_m,
config.storage_config.store_block_n,
static_cast<int>(d.stride(-2)), num_groups,
config.storage_config.swizzle_cd_mode);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, config.layout.block_k, num_groups, 0);
// Launch
const SM90FP8Gemm1D2DRuntime::Args& args = {
.major_sfb = major_sfb,
.m = m, .n = n, .k = k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.epilogue_type = std::nullopt,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = std::nullopt,
.major_sfb = major_sfb,
.sfb = sfb.data_ptr(),
.grouped_layout = masked_m.data_ptr(),
.tensor_map_a = tensor_map_a,
@@ -256,8 +282,8 @@ static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const to
.tensor_map_d = tensor_map_d,
.tensor_map_sfa = tensor_map_sfa,
};
const auto& code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_m_grouped_gemm_masked_1d2d", code);
const auto code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto runtime = compiler->build("sm90_fp8_m_grouped_gemm_masked_1d2d", code);
SM90FP8Gemm1D2DRuntime::launch(runtime, args);
}
@@ -271,51 +297,55 @@ static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::Batched, KernelType::Kernel1D2D,
m, n, k, batch_size, major_a, major_b,
a.scalar_type(), b.scalar_type(),
d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto desc = GemmDesc {
.gemm_type = GemmType::Batched,
.kernel_type = KernelType::Kernel1D2D,
.m = m, .n = n, .k = k, .num_groups = batch_size,
.a_dtype = a.scalar_type(), .b_dtype = b.scalar_type(),
.cd_dtype = d.scalar_type(),
.major_a = major_a, .major_b = major_b,
.with_accumulation = c.has_value(),
.num_sms = device_runtime->get_num_sms(),
.tc_util = device_runtime->get_tc_util(), .compiled_dims = compiled_dims
};
const auto config = get_best_config<SM90ArchSpec>(desc);
// Requires no TMA splits
DG_HOST_ASSERT(config.smem_config.swizzle_a_mode == config.block_k);
DG_HOST_ASSERT(config.smem_config.swizzle_b_mode == config.block_k);
const int& load_block_m = SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& tensor_map_a = make_tma_3d_desc(a, k, m, batch_size,
config.block_k, load_block_m, 1,
a.stride(1),
a.stride(0),
config.smem_config.swizzle_a_mode);
DG_HOST_ASSERT(config.storage_config.swizzle_a_mode == config.layout.block_k);
DG_HOST_ASSERT(config.storage_config.swizzle_b_mode == config.layout.block_k);
const int load_block_m = config.storage_config.load_block_m;
const auto tensor_map_a = make_tma_3d_desc(a, k, m, batch_size,
config.layout.block_k, load_block_m, 1,
a.stride(1),
a.stride(0),
config.storage_config.swizzle_a_mode);
const int& load_block_n = SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& tensor_map_b = make_tma_3d_desc(b, k, n, batch_size,
config.block_k, load_block_n, 1,
b.stride(1),
b.stride(0),
config.smem_config.swizzle_b_mode);
const int load_block_n = config.storage_config.load_block_n;
const auto tensor_map_b = make_tma_3d_desc(b, k, n, batch_size,
config.layout.block_k, load_block_n, 1,
b.stride(1),
b.stride(0),
config.storage_config.swizzle_b_mode);
const int& store_block_m = SM90ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM90ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_d = make_tma_3d_desc(d, n, m, batch_size,
store_block_n, store_block_m, 1,
d.stride(1), d.stride(0),
config.smem_config.swizzle_cd_mode);
const int store_block_m = config.storage_config.store_block_m;
const int store_block_n = config.storage_config.store_block_n;
const auto tensor_map_d = make_tma_3d_desc(d, n, m, batch_size,
store_block_n, store_block_m, 1,
d.stride(1), d.stride(0),
config.storage_config.swizzle_cd_mode);
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.block_m, config.block_k, batch_size, 0);
const auto tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
config.layout.block_m, config.layout.block_k, batch_size, 0);
// Launch
const SM90FP8Gemm1D2DRuntime::Args& args = {
.major_sfb = major_sfb,
.m = m, .n = n, .k = k,
.num_groups = batch_size,
.compiled_dims = compiled_dims,
.epilogue_type = std::nullopt,
.gemm_desc = desc,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.launch_args = LaunchArgs(config.launch_config.num_sms, config.launch_config.num_threads,
config.pipeline_config.smem_size,
config.layout.get_cluster_size()),
.epilogue_type = std::nullopt,
.major_sfb = major_sfb,
.sfb = sfb.data_ptr(),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
@@ -323,8 +353,8 @@ static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
.tensor_map_d = tensor_map_d,
.tensor_map_sfa = tensor_map_sfa,
};
const auto& code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_gemm_1d2d", code);
const auto code = SM90FP8Gemm1D2DRuntime::generate(args);
const auto runtime = compiler->build("sm90_fp8_gemm_1d2d", code);
SM90FP8Gemm1D2DRuntime::launch(runtime, args);
}

36
csrc/jit_kernels/impls/sm90_tf32_hc_prenorm_gemm.hpp
View File

@@ -81,21 +81,21 @@ static void sm90_tf32_hc_prenorm_gemm(const torch::Tensor& a,
DG_HOST_ASSERT(n <= 32 and n % 8 == 0);
DG_HOST_ASSERT(k % block_k == 0);
const auto& swizzle_cd_mode = get_swizzle_mode(block_n, sizeof(float));
const auto& tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k,
block_m, block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, a.element_size()), 0,
true);
const auto& tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b, n, k,
block_n, block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, b.element_size()), 0,
true);
const auto& tensor_map_d = num_splits == 1 ? make_tma_cd_desc(d, m, n,
block_m, block_n,
static_cast<int>(d.stride(-2)), 1,
swizzle_cd_mode)
const auto swizzle_cd_mode = get_swizzle_mode(block_n, sizeof(float));
const auto tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k,
block_m, block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, a.element_size()), 0,
true);
const auto tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b, n, k,
block_n, block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(cute::UMMA::Major::K))), 1,
get_swizzle_mode(block_k, b.element_size()), 0,
true);
const auto tensor_map_d = num_splits == 1 ? make_tma_cd_desc(d, m, n,
block_m, block_n,
static_cast<int>(d.stride(-2)), 1,
swizzle_cd_mode)
: make_tma_3d_desc(d, n, m, num_splits,
block_n, block_m, 1,
static_cast<int>(d.stride(-2)),
@@ -138,14 +138,14 @@ static void sm90_tf32_hc_prenorm_gemm(const torch::Tensor& a,
.num_stages = num_stages,
.num_math_threads = num_math_threads,
.num_tma_threads = num_tma_threads,
.launch_args = LaunchArgs(num_splits * ceil_div(m, block_m), num_threads, smem_size, 1),
.launch_args = LaunchArgs(num_splits * ceil_div(m, block_m), num_threads, smem_size),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_d = tensor_map_d,
.sqr_sum = sqr_sum.data_ptr<float>()
};
const auto& code = SM90BF16HCPrenormGemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_tf32_hc_prenorm_gemm", code);
const auto code = SM90BF16HCPrenormGemmRuntime::generate(args);
const auto runtime = compiler->build("sm90_tf32_hc_prenorm_gemm", code);
SM90BF16HCPrenormGemmRuntime::launch(runtime, args);
}

14
csrc/jit_kernels/impls/smxx_clean_logits.hpp
View File

@@ -17,7 +17,8 @@ public:
int* cu_seq_len_k_start;
int* cu_seq_len_k_end;
float* logits;
void* logits;
at::ScalarType logits_dtype;
int block_kv;
int num_warps;
@@ -33,10 +34,10 @@ using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&smxx_clean_logits<
{}, {}, {}
{}, {}, {}, {}
>);
}};
)", args.next_n, args.block_kv, args.num_warps);
)", args.next_n, args.block_kv, args.num_warps, to_string(args.logits_dtype));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
@@ -65,14 +66,15 @@ static void smxx_clean_logits(const torch::Tensor& logits,
.stride_logits = stride_logits,
.cu_seq_len_k_start = cu_seq_len_k_start.has_value() ? cu_seq_len_k_start.value().data_ptr<int>() : nullptr,
.cu_seq_len_k_end = cu_seq_len_k_end.data_ptr<int>(),
.logits = logits.data_ptr<float>(),
.logits = logits.data_ptr(),
.logits_dtype = logits.scalar_type(),
.block_kv = block_kv,
.num_warps = num_warps,
.launch_args = LaunchArgs(device_runtime->get_num_sms(),
num_warps * 32, smem_size)
};
const auto& code = SMXXCleanLogitsRuntime::generate(args);
const auto& runtime = compiler->build("smxx_clean_logits", code);
const auto code = SMXXCleanLogitsRuntime::generate(args);
const auto runtime = compiler->build("smxx_clean_logits", code);
SMXXCleanLogitsRuntime::launch(runtime, args);
}

73
csrc/jit_kernels/impls/smxx_cublaslt.hpp
View File

@@ -46,7 +46,7 @@ static void call_cublaslt_api(const cublasOperation_t& trans_a,
DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_SCALE_TYPE, &scale_type, sizeof(scale_type)));
#if DG_CUBLASLT_ADVANCED_FEATURES_COMPATIBLE
const int& math_sms = device_runtime->get_num_sms();
const int math_sms = device_runtime->get_num_sms();
DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET, &math_sms, sizeof(math_sms)));
#endif
@@ -57,10 +57,10 @@ static void call_cublaslt_api(const cublasOperation_t& trans_a,
#endif
// Get cuBLASLt handle, workspace, and stream
const auto& handle = device_runtime->get_cublaslt_handle();
const auto& workspace = device_runtime->get_cublaslt_workspace();
const auto& workspace_bytes = workspace.nbytes();
const auto& stream = at::cuda::getCurrentCUDAStream();
const auto handle = device_runtime->get_cublaslt_handle();
const auto workspace = device_runtime->get_cublaslt_workspace();
const auto workspace_bytes = workspace.nbytes();
const auto stream = at::cuda::getCurrentCUDAStream();
// Algorithm selection
cublasLtMatmulPreference_t pref;
@@ -77,7 +77,7 @@ static void call_cublaslt_api(const cublasOperation_t& trans_a,
DG_HOST_ASSERT(num_heuristic_results == 1 and "Unable to find any algorithm for the GEMM");
// Call: D = alpha * (A @ B) + beta * C
const float& alpha = 1.0, beta = accumulate ? 1.0 : 0.0;
const float alpha = 1.0, beta = accumulate ? 1.0 : 0.0;
DG_CUBLASLT_CHECK(cublasLtMatmul(handle, // Light handle
desc, // Operation description
&alpha, // Alpha
@@ -99,47 +99,36 @@ static void call_cublaslt_api(const cublasOperation_t& trans_a,
}
static void cublaslt_gemm(const torch::Tensor& lhs, const torch::Tensor& rhs,
const std::optional<torch::Tensor>& acc,
const torch::Tensor& out,
const int& m, const int& n, const int& k,
const cute::UMMA::Major& a_major, const cute::UMMA::Major& b_major) {
const auto& trans_a = b_major == cute::UMMA::Major::K ? CUBLAS_OP_T : CUBLAS_OP_N;
const auto& trans_b = a_major == cute::UMMA::Major::K ? CUBLAS_OP_N : CUBLAS_OP_T;
// Duplicate the accumulator if necessary
// TODO: remove this
if (acc.has_value()) {
if (acc->data_ptr() == out.data_ptr()) {
DG_HOST_ASSERT(acc->sizes() == out.sizes() and acc->strides() == out.strides());
} else {
out.copy_(acc.value());
}
}
const cute::UMMA::Major& a_major, const cute::UMMA::Major& b_major,
const bool& accumulate) {
const auto trans_a = b_major == cute::UMMA::Major::K ? CUBLAS_OP_T : CUBLAS_OP_N;
const auto trans_b = a_major == cute::UMMA::Major::K ? CUBLAS_OP_N : CUBLAS_OP_T;
// Matrix layouts
const auto& cuda_type_a = at::cuda::ScalarTypeToCudaDataType(rhs.scalar_type());
const auto& cuda_type_b = at::cuda::ScalarTypeToCudaDataType(lhs.scalar_type());
const auto& cuda_type_d = at::cuda::ScalarTypeToCudaDataType(out.scalar_type());
const auto& layout_a = b_major == cute::UMMA::Major::K ? get_cublaslt_layout(cuda_type_a, k, n, rhs.stride(0))
: get_cublaslt_layout(cuda_type_a, n, k, rhs.stride(1));
const auto& layout_b = a_major == cute::UMMA::Major::K ? get_cublaslt_layout(cuda_type_b, k, m, lhs.stride(0))
: get_cublaslt_layout(cuda_type_b, m, k, lhs.stride(1));
const auto& layout_d = get_cublaslt_layout(cuda_type_d, n, m, out.stride(0));
const auto cuda_type_a = at::cuda::ScalarTypeToCudaDataType(rhs.scalar_type());
const auto cuda_type_b = at::cuda::ScalarTypeToCudaDataType(lhs.scalar_type());
const auto cuda_type_d = at::cuda::ScalarTypeToCudaDataType(out.scalar_type());
const auto layout_a = b_major == cute::UMMA::Major::K ? get_cublaslt_layout(cuda_type_a, k, n, rhs.stride(0))
: get_cublaslt_layout(cuda_type_a, n, k, rhs.stride(1));
const auto layout_b = a_major == cute::UMMA::Major::K ? get_cublaslt_layout(cuda_type_b, k, m, lhs.stride(0))
: get_cublaslt_layout(cuda_type_b, m, k, lhs.stride(1));
const auto layout_d = get_cublaslt_layout(cuda_type_d, n, m, out.stride(0));
call_cublaslt_api(trans_a, trans_b, layout_a, layout_b, layout_d, lhs, rhs, out, acc.has_value());
call_cublaslt_api(trans_a, trans_b, layout_a, layout_b, layout_d, lhs, rhs, out, accumulate);
}
static void cublaslt_bhr_hdr_bhd(const torch::Tensor& lhs, const torch::Tensor& rhs, const torch::Tensor& out,
const int& b, const int& h, const int& r, const int& d) {
const auto& m = d, n = b, k = r;
const auto& trans_a = CUBLAS_OP_T;
const auto& trans_b = CUBLAS_OP_N;
const auto m = d, n = b, k = r;
const auto trans_a = CUBLAS_OP_T;
const auto trans_b = CUBLAS_OP_N;
// Matrix layouts
const auto& layout_a = get_cublaslt_layout(CUDA_R_16BF, k, m, rhs.stride(1), h, rhs.stride(0));
const auto& layout_b = get_cublaslt_layout(CUDA_R_16BF, k, n, lhs.stride(0), h, lhs.stride(1));
const auto& layout_d = get_cublaslt_layout(CUDA_R_16BF, m, n, out.stride(0), h, out.stride(1));
const auto layout_a = get_cublaslt_layout(CUDA_R_16BF, k, m, rhs.stride(1), h, rhs.stride(0));
const auto layout_b = get_cublaslt_layout(CUDA_R_16BF, k, n, lhs.stride(0), h, lhs.stride(1));
const auto layout_d = get_cublaslt_layout(CUDA_R_16BF, m, n, out.stride(0), h, out.stride(1));
call_cublaslt_api(trans_a, trans_b, layout_a, layout_b, layout_d, lhs, rhs, out, false);
}
@@ -147,14 +136,14 @@ static void cublaslt_bhr_hdr_bhd(const torch::Tensor& lhs, const torch::Tensor&
static void cublaslt_bhd_hdr_bhr(const torch::Tensor& lhs, const torch::Tensor& rhs, const torch::Tensor& out,
const int& b, const int& h, const int& r, const int& d) {
const auto& m = r, n = b, k = d;
const auto& trans_a = CUBLAS_OP_N;
const auto& trans_b = CUBLAS_OP_N;
const auto m = r, n = b, k = d;
const auto trans_a = CUBLAS_OP_N;
const auto trans_b = CUBLAS_OP_N;
// Matrix layouts
const auto& layout_a = get_cublaslt_layout(CUDA_R_16BF, m, k, rhs.stride(1), h, rhs.stride(0));
const auto& layout_b = get_cublaslt_layout(CUDA_R_16BF, k, n, lhs.stride(0), h, lhs.stride(1));
const auto& layout_d = get_cublaslt_layout(CUDA_R_16BF, m, n, out.stride(0), h, out.stride(1));
const auto layout_a = get_cublaslt_layout(CUDA_R_16BF, m, k, rhs.stride(1), h, rhs.stride(0));
const auto layout_b = get_cublaslt_layout(CUDA_R_16BF, k, n, lhs.stride(0), h, lhs.stride(1));
const auto layout_d = get_cublaslt_layout(CUDA_R_16BF, m, n, out.stride(0), h, out.stride(1));
call_cublaslt_api(trans_a, trans_b, layout_a, layout_b, layout_d, lhs, rhs, out, false);
}

328
csrc/jit_kernels/impls/smxx_fp8_fp4_mqa_logits.hpp Normal file
View File

@@ -0,0 +1,328 @@
#pragma once
#include "../../jit/compiler.hpp"
#include "../../jit/device_runtime.hpp"
#include "../../jit/kernel_runtime.hpp"
#include "../heuristics/sm90.hpp"
#include "../heuristics/sm100.hpp"
#include "runtime_utils.hpp"
namespace deep_gemm {
class SMXXFP8MQALogitsRuntime final: public LaunchRuntime<SMXXFP8MQALogitsRuntime> {
public:
struct Args {
int seq_len;
int seq_len_kv;
int max_seqlen_k;
int stride_logits;
int num_heads, head_dim;
bool is_compressed_logits;
int num_q_stages;
int num_kv_stages;
int block_q;
int block_kv;
int* cu_seq_len_k_start;
int* cu_seq_len_k_end;
void* logits;
CUtensorMap tensor_map_q;
CUtensorMap tensor_map_kv;
CUtensorMap tensor_map_kv_scales;
CUtensorMap tensor_map_weights;
at::ScalarType logits_dtype;
int num_specialized_threads;
int num_math_threads;
LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
// TODO: optimize performance by tuning args
// Block sizes are fixed in this kernel
DG_HOST_ASSERT(128 % args.num_heads == 0);
const auto arch = device_runtime->get_arch(true);
return fmt::format(R"(
#include <deep_gemm/impls/sm{}_fp8_mqa_logits.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm{}_fp8_mqa_logits<
{}, {},
{},
{}, {},
{}, {},
{},
{}, {},
{}
>);
}};
)", arch, arch,
args.num_heads, args.head_dim,
args.is_compressed_logits,
args.block_q, args.block_kv,
args.num_q_stages, args.num_kv_stages,
args.launch_args.grid_dim.first,
args.num_specialized_threads, args.num_math_threads,
to_string(args.logits_dtype));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.seq_len, args.seq_len_kv,
args.max_seqlen_k, args.stride_logits,
args.cu_seq_len_k_start, args.cu_seq_len_k_end,
args.logits,
args.tensor_map_q, args.tensor_map_kv,
args.tensor_map_kv_scales, args.tensor_map_weights
));
}
};
static void smxx_fp8_mqa_logits(const torch::Tensor& q,
const torch::Tensor& kv, const torch::Tensor& kv_scales,
const torch::Tensor& weights,
const torch::Tensor& cu_seq_len_k_start,
const torch::Tensor& cu_seq_len_k_end,
const torch::Tensor& logits,
const at::ScalarType& logits_dtype,
const int& seq_len, const int& seq_len_kv,
const int& max_seqlen_k, const int& stride_logits,
const int& num_heads, const int& head_dim,
const int& block_q, const int& block_kv) {
constexpr int num_specialized_threads = 128;
constexpr int num_q_stages = 3, num_kv_stages = 3;
const int num_math_threads = (device_runtime->get_arch_major() == 10 ? 256 : 512);
// Use compressed logits format when max_seqlen_k is specified
const bool is_compressed_logits = (max_seqlen_k > 0);
// Construct TMAs
DG_HOST_ASSERT(head_dim == 32 or head_dim == 64 or head_dim == 128);
const auto tensor_map_q = make_tma_2d_desc(q, head_dim, seq_len * num_heads,
head_dim, block_q * num_heads, head_dim, head_dim);
const auto tensor_map_kv = make_tma_2d_desc(kv, head_dim, seq_len_kv,
head_dim, block_kv, head_dim, head_dim);
// According to the driver API, the minimal alignment is 256 bytes
// So it is safe for us to do a 16-byte OOB
const auto tensor_map_kv_scales = make_tma_2d_desc(kv_scales,
get_tma_aligned_size(seq_len_kv, static_cast<int>(kv_scales.element_size())),
1, block_kv, 1, 0, 0);
const auto tensor_map_weights = make_tma_2d_desc(weights, num_heads, seq_len,
num_heads, block_q, num_heads, 0);
// Calculate shared memory size
int smem_size = 0;
const int smem_q_size_per_stage = block_q * num_heads * head_dim * static_cast<int>(q.element_size());
const int smem_weight_size_per_stage = block_q * num_heads * static_cast<int>(weights.element_size());
const int smem_kv_size_per_stage = block_kv * head_dim * static_cast<int>(kv.element_size());
const int kv_scale_size_per_stage = block_kv * static_cast<int>(kv_scales.element_size());
smem_size += num_q_stages * smem_q_size_per_stage;
smem_size += num_kv_stages * smem_kv_size_per_stage;
smem_size += num_q_stages * smem_weight_size_per_stage;
smem_size += num_kv_stages * kv_scale_size_per_stage;
smem_size += (num_q_stages * 2 + num_kv_stages * 2 + (num_math_threads / 128) * 2) * 8;
smem_size += 4;
DG_HOST_ASSERT(smem_size <= SM90ArchSpec::smem_capacity);
DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity);
// Launch
const SMXXFP8MQALogitsRuntime::Args args = {
.seq_len = seq_len,
.seq_len_kv = seq_len_kv,
.max_seqlen_k = max_seqlen_k,
.stride_logits = stride_logits,
.num_heads = num_heads, .head_dim = head_dim,
.is_compressed_logits = is_compressed_logits,
.num_q_stages = num_q_stages,
.num_kv_stages = num_kv_stages,
.block_q = block_q,
.block_kv = block_kv,
.cu_seq_len_k_start = cu_seq_len_k_start.data_ptr<int>(),
.cu_seq_len_k_end = cu_seq_len_k_end.data_ptr<int>(),
.logits = logits.data_ptr(),
.tensor_map_q = tensor_map_q,
.tensor_map_kv = tensor_map_kv,
.tensor_map_kv_scales = tensor_map_kv_scales,
.tensor_map_weights = tensor_map_weights,
.logits_dtype = logits_dtype,
.num_specialized_threads = num_specialized_threads,
.num_math_threads = num_math_threads,
.launch_args = LaunchArgs(device_runtime->get_num_sms(),
num_specialized_threads + num_math_threads,
smem_size)
};
const auto code = SMXXFP8MQALogitsRuntime::generate(args);
const auto runtime = compiler->build("smxx_fp8_mqa_logits", code);
SMXXFP8MQALogitsRuntime::launch(runtime, args);
}
class SM100FP4MQALogitsRuntime final: public LaunchRuntime<SM100FP4MQALogitsRuntime> {
public:
struct Args {
int seq_len;
int seq_len_kv;
int max_seqlen_k;
int stride_logits;
int num_heads, head_dim;
bool is_compressed_logits;
int num_q_stages;
int num_kv_stages;
int block_q;
int block_kv;
int* cu_seq_len_k_start;
int* cu_seq_len_k_end;
void* logits;
CUtensorMap tensor_map_q;
CUtensorMap tensor_map_sf_q;
CUtensorMap tensor_map_kv;
CUtensorMap tensor_map_sf_kv;
CUtensorMap tensor_map_weights;
at::ScalarType logits_dtype;
int num_specialized_threads;
int num_math_threads;
LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
// TODO: optimize performance by tuning args
// Block sizes are fixed in this kernel
DG_HOST_ASSERT(128 % args.num_heads == 0);
const auto arch = device_runtime->get_arch(true);
return fmt::format(R"(
#include <deep_gemm/impls/sm100_fp4_mqa_logits.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm100_fp4_mqa_logits<
{}, {},
{},
{}, {},
{}, {},
{},
{}, {},
{}
>);
}};
)", args.num_heads, args.head_dim,
args.is_compressed_logits,
args.block_q, args.block_kv,
args.num_q_stages, args.num_kv_stages,
args.launch_args.grid_dim.first,
args.num_specialized_threads, args.num_math_threads,
to_string(args.logits_dtype));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.seq_len, args.seq_len_kv,
args.max_seqlen_k, args.stride_logits,
args.cu_seq_len_k_start, args.cu_seq_len_k_end,
args.logits,
args.tensor_map_q, args.tensor_map_sf_q,
args.tensor_map_kv, args.tensor_map_sf_kv,
args.tensor_map_weights
));
}
};
static void sm100_fp4_mqa_logits(const torch::Tensor& q, const torch::Tensor& sf_q,
const torch::Tensor& kv, const torch::Tensor& sf_kv,
const torch::Tensor& weights,
const torch::Tensor& cu_seq_len_k_start,
const torch::Tensor& cu_seq_len_k_end,
const torch::Tensor& logits,
const at::ScalarType& logits_dtype,
const int& seq_len, const int& seq_len_kv,
const int& max_seqlen_k, const int& stride_logits,
const int& num_heads, const int& head_dim,
const int& block_q, const int& block_kv) {
constexpr int num_specialized_threads = 128;
const int num_math_threads = 2 * 128;
constexpr int num_q_stages = 3, num_kv_stages = 6, num_tmem_stages = 3;
// Use compressed logits format when max_seqlen_k is specified
const bool is_compressed_logits = (max_seqlen_k > 0);
// Construct TMAs
// `head_dim` must be 128 for 64B swizzling
DG_HOST_ASSERT(head_dim == 128);
const auto tensor_map_q = make_tma_2d_desc(q, head_dim, seq_len * num_heads,
head_dim, block_q * num_heads,
static_cast<int>(q.stride(1)),
head_dim / 2, 0, false, false);
const auto tensor_map_sf_q = make_tma_2d_desc(sf_q, num_heads, seq_len,
num_heads, block_q,
static_cast<int>(sf_q.stride(0)), 0);
const auto tensor_map_weights = make_tma_2d_desc(weights, num_heads, seq_len,
num_heads, block_q,
static_cast<int>(weights.stride(0)), 0);
const auto tensor_map_kv = make_tma_2d_desc(kv, head_dim, seq_len_kv,
head_dim, block_kv,
static_cast<int>(kv.stride(0)),
head_dim / 2, 0, false, false);
// According to the driver API, the minimal alignment is 256 bytes
// So it is safe for us to do a 16-byte OOB
const auto tensor_map_sf_kv = make_tma_2d_desc(sf_kv,
get_tma_aligned_size(seq_len_kv, static_cast<int>(sf_kv.element_size())), 1,
block_kv, 1, 0, 0);
// Calculate shared memory size
const int smem_q_size_per_stage = block_q * num_heads * head_dim / 2;
const int smem_sf_q_size_per_stage = align(block_q * num_heads, 128) * sizeof(int);
const int smem_kv_size_per_stage = block_kv * head_dim / 2;
const int smem_sf_kv_size_per_stage = align(block_kv, 128) * sizeof(int);
const int smem_weight_size_per_stage = block_q * num_heads * sizeof(float);
const int smem_barriers = (num_q_stages + num_kv_stages + num_tmem_stages) * 2 * 8;
const int smem_tmem_ptr = 4;
const int smem_size = num_q_stages * (smem_q_size_per_stage + smem_sf_q_size_per_stage + smem_weight_size_per_stage) +
num_kv_stages * (smem_kv_size_per_stage + smem_sf_kv_size_per_stage) +
smem_barriers + smem_tmem_ptr;
DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity);
// Launch
const SM100FP4MQALogitsRuntime::Args args = {
.seq_len = seq_len,
.seq_len_kv = seq_len_kv,
.max_seqlen_k = max_seqlen_k,
.stride_logits = stride_logits,
.num_heads = num_heads, .head_dim = head_dim,
.is_compressed_logits = is_compressed_logits,
.num_q_stages = num_q_stages,
.num_kv_stages = num_kv_stages,
.block_q = block_q,
.block_kv = block_kv,
.cu_seq_len_k_start = cu_seq_len_k_start.data_ptr<int>(),
.cu_seq_len_k_end = cu_seq_len_k_end.data_ptr<int>(),
.logits = logits.data_ptr(),
.tensor_map_q = tensor_map_q,
.tensor_map_sf_q = tensor_map_sf_q,
.tensor_map_kv = tensor_map_kv,
.tensor_map_sf_kv = tensor_map_sf_kv,
.tensor_map_weights = tensor_map_weights,
.logits_dtype = logits_dtype,
.num_specialized_threads = num_specialized_threads,
.num_math_threads = num_math_threads,
.launch_args = LaunchArgs(device_runtime->get_num_sms(),
num_specialized_threads + num_math_threads,
smem_size)
};
const auto code = SM100FP4MQALogitsRuntime::generate(args);
const auto runtime = compiler->build("sm100_fp4_mqa_logits", code);
SM100FP4MQALogitsRuntime::launch(runtime, args);
}
} // namespace deep_gemm

445
csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp Normal file
View File

@@ -0,0 +1,445 @@
#pragma once
#include "../../jit/compiler.hpp"
#include "../../jit/device_runtime.hpp"
#include "../../jit/kernel_runtime.hpp"
#include "../heuristics/sm90.hpp"
#include "runtime_utils.hpp"
namespace deep_gemm {
class SMXXPagedMQALogitsMetadataRuntime final: public LaunchRuntime<SMXXPagedMQALogitsMetadataRuntime> {
public:
struct Args {
int aligned_batch_size;
int split_kv;
int num_sms;
int batch_size;
int next_n;
bool is_context_lens_2d;
int* context_lens;
int* schedule_metadata;
LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#include <deep_gemm/scheduler/paged_mqa_logits.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sched::smxx_paged_mqa_logits_metadata<
{}, {}, {}
>);
}};
)", args.aligned_batch_size, args.split_kv, args.num_sms);
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.batch_size,
args.next_n,
args.is_context_lens_2d,
args.context_lens,
args.schedule_metadata
));
}
};
static void smxx_paged_mqa_logits_metadata(const torch::Tensor& context_lens,
const torch::Tensor& schedule_metadata,
const int& batch_size, const int& next_n,
const int& block_kv, const int& num_sms,
const bool& is_context_lens_2d) {
constexpr int split_kv = 256;
constexpr int num_threads = 32;
const int aligned_batch_size = align(batch_size, 32);
DG_HOST_ASSERT(split_kv % block_kv == 0);
// Calculate shared memory size
const int smem_size = aligned_batch_size * static_cast<int>(sizeof(int));
DG_HOST_ASSERT(smem_size <= SM90ArchSpec::smem_capacity);
DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity);
// Launch
const SMXXPagedMQALogitsMetadataRuntime::Args& args = {
.aligned_batch_size = aligned_batch_size,
.split_kv = split_kv,
.num_sms = num_sms,
.batch_size = batch_size,
.next_n = next_n,
.is_context_lens_2d = is_context_lens_2d,
.context_lens = context_lens.data_ptr<int>(),
.schedule_metadata = schedule_metadata.data_ptr<int>(),
.launch_args = LaunchArgs(1, num_threads, smem_size)
};
const auto code = SMXXPagedMQALogitsMetadataRuntime::generate(args);
const auto runtime = compiler->build("smxx_paged_mqa_logits_metadata", code);
SMXXPagedMQALogitsMetadataRuntime::launch(runtime, args);
}
class SMXXFP8PagedMQALogitsRuntime final: public LaunchRuntime<SMXXFP8PagedMQALogitsRuntime> {
public:
struct Args {
int batch_size;
int next_n;
int num_heads;
int head_dim;
int block_kv;
bool is_context_lens_2d;
int block_table_stride;
int logits_stride;
int num_q_stages;
int num_kv_stages;
int split_kv;
int* context_lens;
void* logits;
int* block_table;
int* schedule_meta;
CUtensorMap tensor_map_q;
CUtensorMap tensor_map_kv;
CUtensorMap tensor_map_kv_scales;
CUtensorMap tensor_map_weights;
at::ScalarType logits_dtype;
int num_specialized_threads;
int num_math_threads;
LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
// TODO: optimize performance by tuning args
// Block sizes are fixed in this kernel
DG_HOST_ASSERT(128 % args.num_heads == 0);
const auto arch = device_runtime->get_arch(true);
return fmt::format(R"(
#include <deep_gemm/impls/sm{}_fp8_paged_mqa_logits.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm{}_fp8_paged_mqa_logits<
{}, {},
{}, {},
{},
{}, {},
{},
{}, {},
{}
>);
}};
)", arch, arch,
args.next_n, args.num_heads,
args.head_dim, args.block_kv,
args.is_context_lens_2d,
args.num_q_stages, args.num_kv_stages,
args.split_kv,
args.num_specialized_threads, args.num_math_threads,
to_string(args.logits_dtype));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.batch_size,
args.logits_stride, args.block_table_stride,
args.context_lens, args.logits,
args.block_table, args.schedule_meta,
args.tensor_map_q, args.tensor_map_kv,
args.tensor_map_kv_scales, args.tensor_map_weights
));
}
};
static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q,
const torch::Tensor& kv_cache,
const torch::Tensor& kv_cache_scales,
const torch::Tensor& weights,
const torch::Tensor& context_lens,
const torch::Tensor& logits,
const torch::Tensor& block_table,
const torch::Tensor& schedule_meta,
const at::ScalarType& logits_dtype,
const int& batch_size, const int& next_n,
const int& num_heads, const int& head_dim,
const int& num_kv_blocks, const int& block_kv,
const bool& is_context_lens_2d,
const int& logits_stride,
const int& block_table_stride,
const int& num_sms,
const int& split_kv) {
const int num_specialized_threads = 128;
const int mma_m = (device_runtime->get_arch_major() == 10 ? 128 : 64);
const int num_math_warp_groups = split_kv / mma_m;
const int num_math_threads = num_math_warp_groups * 128;
const int num_q_stages = 3, num_kv_stages = (device_runtime->get_arch_major() == 10 ? 4 : 3);
DG_HOST_ASSERT(split_kv % mma_m == 0 and logits_stride % split_kv == 0);
// Construct TMAs
const int next_n_atom = (next_n % 2 == 0) ? 2 : 1;
const auto tensor_map_q = make_tma_2d_desc(q, head_dim, batch_size * next_n * num_heads,
head_dim, next_n_atom * num_heads,
static_cast<int>(q.stride(2)),
head_dim);
const auto tensor_map_kv = make_tma_3d_desc(kv_cache, head_dim, block_kv, num_kv_blocks,
head_dim, block_kv, 1,
static_cast<int>(kv_cache.stride(1)),
static_cast<int>(kv_cache.stride(0)),
head_dim);
const auto tensor_map_kv_scales = make_tma_2d_desc(kv_cache_scales, block_kv, num_kv_blocks,
block_kv, 1,
static_cast<int>(kv_cache_scales.stride(0)), 0);
const auto tensor_map_weights = make_tma_2d_desc(weights, num_heads, batch_size * next_n,
num_heads, next_n_atom,
static_cast<int>(weights.stride(0)), 0);
// Calculate shared memory size
int smem_size = 0;
if (device_runtime->get_arch_major() == 9) {
const int swizzle_alignment = head_dim * 8;
const int smem_q_size_per_stage = next_n * num_heads * head_dim * static_cast<int>(q.element_size());
const int aligned_smem_weight_size_per_stage = align(next_n * num_heads * static_cast<int>(weights.element_size()), swizzle_alignment);
const int smem_q_pipe_size = num_q_stages * (smem_q_size_per_stage + aligned_smem_weight_size_per_stage) + align(num_q_stages * 8 * 2, swizzle_alignment);
const int smem_kv_size_per_stage = block_kv * head_dim * static_cast<int>(kv_cache.element_size());
const int aligned_smem_kv_scale_size_per_stage = align(block_kv * static_cast<int>(kv_cache_scales.element_size()), swizzle_alignment);
const int smem_kv_pipe_size = num_kv_stages * (smem_kv_size_per_stage + aligned_smem_kv_scale_size_per_stage) + align(num_kv_stages * 8 * 2, swizzle_alignment);
// Allocate some shared memory for UMMA barriers and tensor memory pointer, although it is not used in SM90
const int smem_umma_barriers = num_math_warp_groups * 2 * 8;
const int smem_tmem_ptr = 4;
smem_size = smem_q_pipe_size + num_math_warp_groups * smem_kv_pipe_size + smem_umma_barriers + smem_tmem_ptr;
DG_HOST_ASSERT(smem_size <= SM90ArchSpec::smem_capacity);
DG_HOST_ASSERT(next_n == 1 or next_n == 2);
} else {
const int smem_q_size_per_stage = next_n_atom * num_heads * head_dim * static_cast<int>(q.element_size());
const int smem_kv_size_per_stage = split_kv * head_dim * static_cast<int>(kv_cache.element_size());
const int smem_kv_scale_size_per_stage = split_kv * static_cast<int>(kv_cache_scales.element_size());
const int smem_weight_size_per_stage = next_n_atom * num_heads * static_cast<int>(weights.element_size());
const int smem_barriers = (num_q_stages + num_kv_stages) * 2 * 8;
const int smem_umma_barriers = num_math_warp_groups * 2 * 8;
const int smem_tmem_ptr = 4;
smem_size = num_q_stages * (smem_q_size_per_stage + smem_weight_size_per_stage) +
num_kv_stages * (smem_kv_size_per_stage + smem_kv_scale_size_per_stage) +
smem_barriers + smem_umma_barriers + smem_tmem_ptr;
DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity);
}
// Launch
const SMXXFP8PagedMQALogitsRuntime::Args args = {
.batch_size = batch_size,
.next_n = next_n,
.num_heads = num_heads,
.head_dim = head_dim,
.block_kv = block_kv,
.is_context_lens_2d = is_context_lens_2d,
.block_table_stride = block_table_stride,
.logits_stride = logits_stride,
.num_q_stages = num_q_stages,
.num_kv_stages = num_kv_stages,
.split_kv = split_kv,
.context_lens = context_lens.data_ptr<int>(),
.logits = logits.data_ptr(),
.block_table = block_table.data_ptr<int>(),
.schedule_meta = schedule_meta.data_ptr<int>(),
.tensor_map_q = tensor_map_q,
.tensor_map_kv = tensor_map_kv,
.tensor_map_kv_scales = tensor_map_kv_scales,
.tensor_map_weights = tensor_map_weights,
.logits_dtype = logits_dtype,
.num_specialized_threads = num_specialized_threads,
.num_math_threads = num_math_threads,
.launch_args = LaunchArgs(num_sms,
num_specialized_threads + num_math_threads,
smem_size)
};
const auto code = SMXXFP8PagedMQALogitsRuntime::generate(args);
const auto runtime = compiler->build("smxx_fp8_paged_mqa_logits", code);
SMXXFP8PagedMQALogitsRuntime::launch(runtime, args);
}
class SM100FP4PagedMQALogitsRuntime final: public LaunchRuntime<SM100FP4PagedMQALogitsRuntime> {
public:
struct Args {
int batch_size;
int next_n;
int num_heads;
int head_dim;
int block_kv;
bool is_context_lens_2d;
int block_table_stride;
int logits_stride;
int num_q_stages;
int num_kv_stages;
int split_kv;
int* context_lens;
void* logits;
int* block_table;
int* schedule_meta;
CUtensorMap tensor_map_q;
CUtensorMap tensor_map_sf_q;
CUtensorMap tensor_map_kv;
CUtensorMap tensor_map_sf_kv;
CUtensorMap tensor_map_weights;
at::ScalarType logits_dtype;
int num_specialized_threads;
int num_math_threads;
LaunchArgs launch_args;
};
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#include <deep_gemm/impls/sm100_fp4_paged_mqa_logits.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm100_fp4_paged_mqa_logits<
{}, {},
{}, {},
{},
{}, {},
{},
{}, {},
{}
>);
}};
)", args.next_n, args.num_heads,
args.head_dim, args.block_kv,
args.is_context_lens_2d,
args.num_q_stages, args.num_kv_stages,
args.split_kv,
args.num_specialized_threads, args.num_math_threads,
to_string(args.logits_dtype));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.batch_size,
args.logits_stride, args.block_table_stride,
args.context_lens, args.logits,
args.block_table, args.schedule_meta,
args.tensor_map_q, args.tensor_map_sf_q,
args.tensor_map_kv, args.tensor_map_sf_kv,
args.tensor_map_weights
));
}
};
static void sm100_fp4_paged_mqa_logits(const torch::Tensor& q,
const torch::Tensor& sf_q,
const torch::Tensor& kv_cache,
const torch::Tensor& kv_cache_sf,
const torch::Tensor& weights,
const torch::Tensor& context_lens,
const torch::Tensor& logits,
const torch::Tensor& block_table,
const torch::Tensor& schedule_meta,
const at::ScalarType& logits_dtype,
const int& batch_size, const int& next_n,
const int& num_heads, const int& head_dim,
const int& num_kv_blocks, const int& block_kv,
const bool& is_context_lens_2d,
const int& logits_stride,
const int& block_table_stride,
const int& num_sms,
const int& split_kv) {
const int num_specialized_threads = 128;
const int num_math_threads = 2 * 128;
DG_HOST_ASSERT(split_kv == 256 and logits_stride % split_kv == 0);
// TODO: tuning num_stages
const int num_q_stages = 3, num_kv_stages = 6, num_tmem_stages = 3;
const int next_n_atom = (next_n % 2 == 0) ? 2 : 1;
// `head_dim` must be 128 for 64B swizzling
DG_HOST_ASSERT(head_dim == 128);
// Using 2D TMA as tensor q is asserted contiguous
const auto tensor_map_q = make_tma_2d_desc(q, head_dim, batch_size * next_n * num_heads,
head_dim, next_n_atom * num_heads,
static_cast<int>(q.stride(2)),
head_dim / 2, 0, false, false);
// NOTES: `sf_q` is a 3D tensor, while `weights` is a 2D tensor
const auto tensor_map_sf_q = make_tma_2d_desc(sf_q, num_heads, batch_size * next_n,
num_heads, next_n_atom,
static_cast<int>(sf_q.stride(1)), 0);
const auto tensor_map_weights = make_tma_2d_desc(weights, num_heads, batch_size * next_n,
num_heads, next_n_atom,
static_cast<int>(weights.stride(0)), 0);
const auto tensor_map_kv = make_tma_3d_desc(kv_cache, head_dim, block_kv, num_kv_blocks,
head_dim, block_kv, 1,
static_cast<int>(kv_cache.stride(1)),
static_cast<int>(kv_cache.stride(0)),
head_dim / 2, 0, false, false);
const auto tensor_map_sf_kv = make_tma_2d_desc(kv_cache_sf, block_kv, num_kv_blocks,
block_kv, 1,
static_cast<int>(kv_cache_sf.stride(0)), 0);
// Calculate shared memory size
const int smem_q_size_per_stage = next_n_atom * num_heads * head_dim / 2;
const int smem_sf_q_size_per_stage = align(next_n_atom * num_heads, 128) * sizeof(int);
const int smem_kv_size_per_stage = split_kv * head_dim / 2;
const int smem_sf_kv_size_per_stage = align(split_kv, 128) * sizeof(int);
const int smem_weight_size_per_stage = next_n_atom * num_heads * sizeof(float);
const int smem_barriers = (num_q_stages + num_kv_stages + num_tmem_stages) * 2 * 8;
const int smem_tmem_ptr = 4;
const int smem_size = num_q_stages * (smem_q_size_per_stage + smem_sf_q_size_per_stage + smem_weight_size_per_stage) +
num_kv_stages * (smem_kv_size_per_stage + smem_sf_kv_size_per_stage) +
smem_barriers + smem_tmem_ptr;
DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity);
// Launch
const SM100FP4PagedMQALogitsRuntime::Args args = {
.batch_size = batch_size,
.next_n = next_n,
.num_heads = num_heads,
.head_dim = head_dim,
.block_kv = block_kv,
.is_context_lens_2d = is_context_lens_2d,
.block_table_stride = block_table_stride,
.logits_stride = logits_stride,
.num_q_stages = num_q_stages,
.num_kv_stages = num_kv_stages,
.split_kv = split_kv,
.context_lens = context_lens.data_ptr<int>(),
.logits = logits.data_ptr(),
.block_table = block_table.data_ptr<int>(),
.schedule_meta = schedule_meta.data_ptr<int>(),
.tensor_map_q = tensor_map_q,
.tensor_map_sf_q = tensor_map_sf_q,
.tensor_map_kv = tensor_map_kv,
.tensor_map_sf_kv = tensor_map_sf_kv,
.tensor_map_weights = tensor_map_weights,
.logits_dtype = logits_dtype,
.num_specialized_threads = num_specialized_threads,
.num_math_threads = num_math_threads,
.launch_args = LaunchArgs(num_sms,
num_specialized_threads + num_math_threads,
smem_size)
};
const auto code = SM100FP4PagedMQALogitsRuntime::generate(args);
const auto runtime = compiler->build("sm100_fp4_paged_mqa_logits", code);
SM100FP4PagedMQALogitsRuntime::launch(runtime, args);
}
} // namespace deep_gemm

75
csrc/jit_kernels/impls/smxx_layout.hpp
View File

@@ -72,7 +72,7 @@ static void __instantiate_kernel() {{
class PackFP32IntoUE8M0Runtime final: public LaunchRuntime<PackFP32IntoUE8M0Runtime> {
public:
struct Args {
int num_groups, mn, sf_k, packed_sf_k;
int num_groups, mn, sf_k, packed_sf_k, gran_k;
int block_mn, block_packed_sf_k;
void *sf, *out, *ks;
@@ -95,32 +95,32 @@ static void __instantiate_kernel() {{
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.sf, args.out, args.ks, args.mn, args.sf_k, args.packed_sf_k));
args.sf, args.out, args.ks, args.mn, args.sf_k, args.packed_sf_k, args.gran_k));
}
};
static std::tuple<int, int, int, int, int, torch::Tensor> preprocess_sf(const torch::Tensor& sf) {
// NOTES: for the extreme performance, you may rewrite/fuse this function in CUDA
const auto& dim = sf.dim();
const auto dim = sf.dim();
DG_HOST_ASSERT(dim == 2 or dim == 3);
DG_HOST_ASSERT(sf.scalar_type() == torch::kFloat);
const auto& batched_sf = dim == 2 ? sf.unsqueeze(0) : sf;
const auto batched_sf = dim == 2 ? sf.unsqueeze(0) : sf;
const auto& [num_groups, mn, sf_k] = get_shape<3>(batched_sf);
const auto& tma_aligned_mn = get_tma_aligned_size(mn, static_cast<int>(sf.element_size()));
const auto [num_groups, mn, sf_k] = get_shape<3>(batched_sf);
const auto tma_aligned_mn = get_tma_aligned_size(mn, static_cast<int>(sf.element_size()));
return {dim, num_groups, mn, sf_k, tma_aligned_mn, batched_sf};
}
static torch::Tensor get_mn_major_tma_aligned_tensor(const torch::Tensor& sf) {
const auto& [dim, num_groups, mn, sf_k, tma_aligned_mn, batched_sf] = preprocess_sf(sf);
const auto [dim, num_groups, mn, sf_k, tma_aligned_mn, batched_sf] = preprocess_sf(sf);
// The last kernel already gives a column-major TMA aligned layout
if ((batched_sf.stride(0) == tma_aligned_mn * sf_k or dim == 2) and batched_sf.stride(1) == 1 and batched_sf.stride(2) == tma_aligned_mn)
return (dim == 2) ? batched_sf.squeeze(0) : batched_sf;
const auto& out = torch::empty_strided({num_groups, mn, sf_k},
{tma_aligned_mn * sf_k, 1, tma_aligned_mn},
batched_sf.options());
const auto out = torch::empty_strided({num_groups, mn, sf_k},
{tma_aligned_mn * sf_k, 1, tma_aligned_mn},
batched_sf.options());
if (not batched_sf.is_contiguous()) {
// Fallback to PyTorch's slow copy if not contiguous
@@ -129,7 +129,7 @@ static torch::Tensor get_mn_major_tma_aligned_tensor(const torch::Tensor& sf) {
} else {
constexpr int block_mn = 64;
constexpr int num_threads = 512;
const auto& smem_size = block_mn * (sf_k + (1 - (sf_k % 2))) * static_cast<int>(sizeof(float));
const auto smem_size = block_mn * (sf_k + (1 - (sf_k % 2))) * static_cast<int>(sizeof(float));
const TransposeFP32Runtime::Args& args = {
.mn = mn,
.sf_k = sf_k,
@@ -139,25 +139,25 @@ static torch::Tensor get_mn_major_tma_aligned_tensor(const torch::Tensor& sf) {
.launch_args = LaunchArgs({ceil_div(mn, block_mn), num_groups}, num_threads, smem_size)
};
const auto& code = TransposeFP32Runtime::generate(args);
const auto& runtime = compiler->build("transpose_fp32", code);
const auto code = TransposeFP32Runtime::generate(args);
const auto runtime = compiler->build("transpose_fp32", code);
TransposeFP32Runtime::launch(runtime, args);
}
return (dim == 2) ? out.squeeze(0) : out;
}
static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor_torch(const torch::Tensor& sf) {
const auto& sf_reshaped = (sf.dim() == 2) ? sf.unsqueeze(0) : sf;
const auto sf_reshaped = (sf.dim() == 2) ? sf.unsqueeze(0) : sf;
// First, convert into UE8M0 `uint8_t`
const auto& ue8m0_tensor = sf_reshaped.view(torch::kInt32).bitwise_right_shift(23).to(torch::kUInt8);
const auto ue8m0_tensor = sf_reshaped.view(torch::kInt32).bitwise_right_shift(23).to(torch::kUInt8);
// Second, make padded packed tensors
const auto& [num_groups, mn, k] = get_shape<3>(sf_reshaped);
const auto& aligned_mn = get_tma_aligned_size(mn, 4);
const auto& aligned_k = align(k, 4);
const auto [num_groups, mn, k] = get_shape<3>(sf_reshaped);
const auto aligned_mn = get_tma_aligned_size(mn, 4);
const auto aligned_k = align(k, 4);
const auto& options = torch::TensorOptions().device(sf.device()).dtype(torch::kUInt8);
const auto options = torch::TensorOptions().device(sf.device()).dtype(torch::kUInt8);
auto padded = torch::zeros({num_groups, aligned_mn, aligned_k}, options);
// ReSharper disable once CppExpressionWithoutSideEffects
padded.slice(1, 0, mn).slice(2, 0, k).copy_(ue8m0_tensor);
@@ -172,11 +172,11 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor_torch(const to
}
static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::Tensor& sf) {
const auto& [dim, num_groups, mn, sf_k, tma_aligned_mn, batched_sf] = preprocess_sf(sf);
const auto& packed_sf_k = ceil_div(sf_k, 4);
const auto& out = torch::empty_strided({num_groups, mn, packed_sf_k},
{packed_sf_k * tma_aligned_mn, 1, tma_aligned_mn},
at::TensorOptions().device(batched_sf.device()).dtype(torch::kInt));
const auto [dim, num_groups, mn, sf_k, tma_aligned_mn, batched_sf] = preprocess_sf(sf);
const auto packed_sf_k = ceil_div(sf_k, 4);
const auto out = torch::empty_strided({num_groups, mn, packed_sf_k},
{packed_sf_k * tma_aligned_mn, 1, tma_aligned_mn},
at::TensorOptions().device(batched_sf.device()).dtype(torch::kInt));
// Launch the kernel
if (batched_sf.is_contiguous()) {
if ((mn * sf_k) % 4 != 0 and num_groups > 1)
@@ -193,8 +193,8 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::T
.launch_args = LaunchArgs({ceil_div(mn, block_mn), num_groups}, num_threads, block_mn * sf_k * 4)
};
const auto& code = TransposeAndPackFP32IntoUE8M0Runtime::generate(args);
const auto& runtime = compiler->build("transpose_and_pack_fp32_into_ue8m0", code);
const auto code = TransposeAndPackFP32IntoUE8M0Runtime::generate(args);
const auto runtime = compiler->build("transpose_and_pack_fp32_into_ue8m0", code);
TransposeAndPackFP32IntoUE8M0Runtime::launch(runtime, args);
} else {
if (mn % 4 != 0 or num_groups > 1)
@@ -217,8 +217,8 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::T
.launch_args = LaunchArgs({ceil_div(mn, block_mn), ceil_div(packed_sf_k, block_packed_sf_k)}, num_threads)
};
const auto& code = PackFP32IntoUE8M0Runtime::generate(args);
const auto& runtime = compiler->build("pack_fp32_into_ue8m0", code);
const auto code = PackFP32IntoUE8M0Runtime::generate(args);
const auto runtime = compiler->build("pack_fp32_into_ue8m0", code);
PackFP32IntoUE8M0Runtime::launch(runtime, args);
}
return (dim == 2) ? out.squeeze(0) : out;
@@ -226,18 +226,20 @@ static torch::Tensor get_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::T
static torch::Tensor get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(const torch::Tensor& sf,
const torch::Tensor& ks_tensor,
const std::vector<int>& ks) {
const auto& [sf_k, mn] = get_shape<2>(sf);
const auto& num_groups = static_cast<int>(ks.size());
const std::vector<int>& ks,
const int gran_k) {
DG_HOST_ASSERT(gran_k == 32 or gran_k == 128);
const auto [sf_k, mn] = get_shape<2>(sf);
const auto num_groups = static_cast<int>(ks.size());
int ref_sf_k = 0, packed_sf_k = 0;
for (const auto& k: ks)
ref_sf_k += ceil_div(k, 128), packed_sf_k += ceil_div(k, 512);
for (const auto k: ks)
ref_sf_k += ceil_div(k, gran_k), packed_sf_k += ceil_div(k, gran_k * 4);
DG_HOST_ASSERT(sf.is_contiguous());
DG_HOST_ASSERT(ref_sf_k == sf_k);
DG_HOST_ASSERT(num_groups <= 128 and mn % 4 == 0);
const auto& out = torch::empty({packed_sf_k, mn}, at::TensorOptions().device(sf.device()).dtype(torch::kInt));
const auto out = torch::empty({packed_sf_k, mn}, at::TensorOptions().device(sf.device()).dtype(torch::kInt));
constexpr int block_mn = 128;
constexpr int block_packed_sf_k = 16;
@@ -247,6 +249,7 @@ static torch::Tensor get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(cons
.mn = mn,
.sf_k = sf_k,
.packed_sf_k = packed_sf_k,
.gran_k = gran_k,
.block_mn = block_mn,
.block_packed_sf_k = block_packed_sf_k,
.sf = sf.data_ptr(),
@@ -255,8 +258,8 @@ static torch::Tensor get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(cons
.launch_args = LaunchArgs({ceil_div(mn, block_mn), ceil_div(packed_sf_k, block_packed_sf_k)}, num_threads)
};
const auto& code = PackFP32IntoUE8M0Runtime::generate(args);
const auto& runtime = compiler->build("pack_fp32_into_ue8m0", code);
const auto code = PackFP32IntoUE8M0Runtime::generate(args);
const auto runtime = compiler->build("pack_fp32_into_ue8m0", code);
PackFP32IntoUE8M0Runtime::launch(runtime, args);
return out;
}