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

feat(megamoe): add nvfp4 group16 capability gate

Browse Source 
Allow SM100 FP4 scale layout transforms to accept group16 and thread weight granularity through the MegaMoE Python wrapper, API checks, and synthetic benchmark entrypoint.

Keep fused SM100 MegaMoE compute behind an explicit group16 capability gate until the SFB/TMEM/MMA scale path is updated and validated.

Tested: PYTHONPYCACHEPREFIX=/private/tmp/deepgemm_pycache python3 -m py_compile deep_gemm/mega/__init__.py tests/test_mega_moe.py tests/generators.py

Tested: git diff --check

Not-tested: CUDA build and SM100/B300 runtime validation are not available locally.
This commit is contained in:
LuminolT
2026-07-08 18:29:09 +08:00
parent 0ff91f1285
commit 79fcfd6abf
9 changed files with 99 additions and 31 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/apis/layout.hpp
+6 -4
View File
@@ -44,8 +44,10 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf,
if (sf.scalar_type() == torch::kFloat and gran_mn == 128 and gran_k == 128 and (arch_major == 9 or disable_ue8m0_cast))
return check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups, false, true, torch::kFloat);
// (FP32, x, gran_k) on SM100: transform to (INT, 1, gran_k), TMA-aligned and MN-major
if (sf.scalar_type() == torch::kFloat and (gran_k == 32 or gran_k == 128) and arch_major == 10) {
// (FP32, x, gran_k) on SM100: transform to (INT, 1, gran_k), TMA-aligned and MN-major.
// GLM-5.2 NVFP4 checkpoints use weight granularity 16, while the original
// SM100 MegaMoE path only exercised 32.
if (sf.scalar_type() == torch::kFloat and (gran_k == 16 or gran_k == 32 or gran_k == 128) and arch_major == 10) {
DG_HOST_ASSERT(not disable_ue8m0_cast);
const auto broadcasted = gran_mn == 1 ? sf :
sf.index_select(-2, torch::arange(mn, at::TensorOptions().device(sf.device())).floor_divide_(gran_mn));
@@ -53,7 +55,7 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf,
}
// (INT, 1, gran_k) on SM100: transform to TMA-aligned and MN-major
if (sf.scalar_type() == torch::kInt and gran_mn == 1 and (gran_k == 32 or gran_k == 128) and arch_major == 10)
if (sf.scalar_type() == torch::kInt and gran_mn == 1 and (gran_k == 16 or gran_k == 32 or gran_k == 128) and arch_major == 10)
return check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups, true, false, torch::kInt);
DG_HOST_UNREACHABLE("Unknown SF transformation");
@@ -94,7 +96,7 @@ static torch::Tensor transform_k_grouped_sf_into_required_layout(const torch::Te
DG_HOST_ASSERT(std::get<0>(recipe) == 1 and std::get<1>(recipe) == 1);
const int gran_k = std::get<2>(recipe);
DG_HOST_ASSERT(gran_k == 32 or gran_k == 128);
DG_HOST_ASSERT(gran_k == 16 or gran_k == 32 or gran_k == 128);
const auto arch_major = device_runtime->get_arch_major();
csrc/apis/mega.hpp
+11 -4
View File
@@ -150,7 +150,7 @@ static void fp8_fp4_mega_moe(
// Config checks
const auto num_tokens = static_cast<int>(y.size(0));
const auto [rm, rn, rk] = recipe;
DG_HOST_ASSERT(rm == 1 and rn == 1 and rk == 32);
DG_HOST_ASSERT(rm == 1 and rn == 1 and (rk == 16 or rk == 32));
DG_HOST_ASSERT(activation == "swiglu");
// Activation checks
@@ -173,11 +173,17 @@ static void fp8_fp4_mega_moe(
DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous());
// Check weight SF layout for UE8M0 packing, MN-major, and TMA alignment
constexpr int kGranMN = 1, kGranK = 32;
check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, kGranK,
constexpr int kGranMN = 1;
const int weight_gran_k = rk;
check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, weight_gran_k,
num_experts_per_rank, true, false, torch::kInt);
check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, kGranK,
check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, weight_gran_k,
num_experts_per_rank, true, false, torch::kInt);
if (weight_gran_k == 16) {
DG_HOST_UNREACHABLE(
"SM100 FP8xFP4 MegaMoE weight granularity 16 requires kernel support for "
"NVFP4 group16; the current fused compute path still uses mxf4.block_scale.block32");
}
// Check stats counter
if (cumulative_local_expert_recv_stats.has_value()) {
@@ -213,6 +219,7 @@ static void fp8_fp4_mega_moe(
num_experts_per_rank,
num_tokens, num_topk,
hidden, intermediate_hidden,
weight_gran_k,
activation_clamp, fast_math);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp
+10 -5
View File
@@ -22,6 +22,7 @@ public:
int num_max_tokens_per_rank;
int hidden, intermediate_hidden;
int num_experts, num_topk;
int weight_gran_k;
int num_ranks;
float activation_clamp;
bool fast_math;
@@ -60,6 +61,7 @@ static void __instantiate_kernel() {{
{}, {},
{}, {},
{},
{},
{}, {}, {},
{},
{}, {},
@@ -75,6 +77,7 @@ static void __instantiate_kernel() {{
)", args.num_max_tokens_per_rank,
args.hidden, args.intermediate_hidden,
args.num_experts, args.num_topk,
args.weight_gran_k,
args.config.num_experts_per_wave,
args.config.block_m, args.config.block_n, args.config.block_k,
args.config.store_block_m,
@@ -120,6 +123,7 @@ static void sm100_fp8_fp4_mega_moe(
const int& num_experts_per_rank,
const int& num_tokens, const int& num_topk,
const int& hidden, const int& intermediate_hidden,
const int& weight_gran_k,
const float& activation_clamp,
const bool& fast_math
) {
@@ -133,7 +137,7 @@ static void sm100_fp8_fp4_mega_moe(
num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, num_padded_sf_pool_tokens);
// Make tensormap
constexpr int kGranK = 32;
constexpr int kActivationGranK = 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,
@@ -141,7 +145,7 @@ static void sm100_fp8_fp4_mega_moe(
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,
config.sf_block_m, kActivationGranK,
1, 0);
const auto tensor_map_l1_weights = make_tma_2d_desc(l1_weights,
hidden, num_experts_per_rank * intermediate_hidden * 2,
@@ -150,7 +154,7 @@ static void sm100_fp8_fp4_mega_moe(
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,
config.block_n, weight_gran_k,
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),
@@ -167,7 +171,7 @@ static void sm100_fp8_fp4_mega_moe(
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,
config.sf_block_m, kActivationGranK,
1, 0);
const auto tensor_map_l2_weights = make_tma_2d_desc(l2_weights,
intermediate_hidden, num_experts_per_rank * hidden,
@@ -176,7 +180,7 @@ static void sm100_fp8_fp4_mega_moe(
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,
config.block_n, weight_gran_k,
num_experts_per_rank, 0);
// Stats can be optional
@@ -190,6 +194,7 @@ static void sm100_fp8_fp4_mega_moe(
.num_max_tokens_per_rank = num_max_tokens_per_rank,
.hidden = hidden, .intermediate_hidden = intermediate_hidden,
.num_experts = num_experts, .num_topk = num_topk,
.weight_gran_k = weight_gran_k,
.num_ranks = num_ranks,
.activation_clamp = activation_clamp,
.fast_math = fast_math,
csrc/jit_kernels/impls/smxx_layout.hpp
+1 -1
View File
@@ -228,7 +228,7 @@ static torch::Tensor get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(cons
const torch::Tensor& ks_tensor,
const std::vector<int>& ks,
const int gran_k) {
DG_HOST_ASSERT(gran_k == 32 or gran_k == 128);
DG_HOST_ASSERT(gran_k == 16 or 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());