feat(megamoe): add nvfp4 group16 capability gate

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
+2 -2
View File
@@ -202,7 +202,7 @@ def enumerate_k_grouped_contiguous(dtype: torch.dtype):
def enumerate_sf_layout():
gran_k_list = (128, ) if get_arch_major() == 9 else (32, 128)
gran_k_list = (128, ) if get_arch_major() == 9 else (16, 32, 128)
for use_ue8m0 in (False, True):
for with_transpose in (True, False):
for mn in (4096, 4097, 8192):
@@ -214,7 +214,7 @@ def enumerate_sf_layout():
def enumerate_k_grouped_sf_layout():
gran_k_list = (128, ) if get_arch_major() == 9 else (32, 128)
gran_k_list = (128, ) if get_arch_major() == 9 else (16, 32, 128)
for mn in (4096, 7168):
for num_groups, avg_k in ((16, 2048), (8, 4096), (72, 384), (128, 256)):
for gran_k in gran_k_list: