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/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");