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
deep_gemm/mega/__init__.py
+7 -5
View File
@@ -142,10 +142,11 @@ def _interleave_l1_weights(l1_weights: Tuple[torch.Tensor, torch.Tensor]) -> Tup
return _interleave_l1_weight_tensor(l1_weights[0]), _interleave_l1_weight_tensor(l1_weights[1])
def _transpose_sf_for_utccp(sf: torch.Tensor) -> torch.Tensor:
def _transpose_sf_for_utccp(sf: torch.Tensor, gran_k: int = 32) -> torch.Tensor:
num_groups, mn, packed_sf_k = sf.shape
assert sf.dtype == torch.int and mn % 128 == 0
result = (sf.reshape(num_groups, -1, 4, 32, packed_sf_k)
assert 128 % gran_k == 0
result = (sf.reshape(num_groups, -1, 128 // gran_k, gran_k, packed_sf_k)
.transpose(2, 3)
.reshape(num_groups, mn, packed_sf_k))
return torch.empty_like(sf).copy_(result)
@@ -163,14 +164,15 @@ def transform_weights_for_mega_moe_sm90(
def transform_weights_for_mega_moe(
l1_weights: Tuple[torch.Tensor, torch.Tensor],
l2_weights: Tuple[torch.Tensor, torch.Tensor]
l2_weights: Tuple[torch.Tensor, torch.Tensor],
weight_gran_k: int = 32,
) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
if _is_sm90():
return transform_weights_for_mega_moe_sm90(l1_weights, l2_weights)
# SM100: L1 interleave gate/up + UTCCP SF transpose, L2 UTCCP SF transpose
l1_interleaved = _interleave_l1_weights(l1_weights)
l1_weights = (l1_interleaved[0], _transpose_sf_for_utccp(l1_interleaved[1]))
l2_weights = (l2_weights[0], _transpose_sf_for_utccp(l2_weights[1]))
l1_weights = (l1_interleaved[0], _transpose_sf_for_utccp(l1_interleaved[1], weight_gran_k))
l2_weights = (l2_weights[0], _transpose_sf_for_utccp(l2_weights[1], weight_gran_k))
return l1_weights, l2_weights