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

4
tests/generators.py
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:

20
tests/test_mega_moe.py
View File

@@ -82,7 +82,8 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
assert intermediate_hidden % 128 == 0
assert l1_weights.shape[2] % 128 == 0 and l2_weights.shape[2] % 128 == 0
# Cast inputs to FP8 with per-32 UE8M0 SF
# Cast inputs to FP8 with per-32 UE8M0 SF. GLM NVFP4 group16 applies to
# FP4 weights only; activation dispatch remains per-32 in this test.
x = per_token_cast_to_fp8(x, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True)
# Cast grouped BF16 weights to FP4 with MN-major SF
@@ -90,15 +91,17 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
def cast_grouped_weights_to_fp4(bf16_weights: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
num_groups, n, k = bf16_weights.shape
w = torch.empty((num_groups, n, k // 2), device='cuda', dtype=torch.int8)
w_sf = torch.empty((num_groups, n, k // 32), device='cuda', dtype=torch.float)
w_sf = torch.empty((num_groups, n, k // args.weight_gran_k), device='cuda', dtype=torch.float)
for i in range(num_groups):
w[i], w_sf[i] = per_token_cast_to_fp4(bf16_weights[i], use_ue8m0=True, gran_k=32)
w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, 32), num_groups)
w[i], w_sf[i] = per_token_cast_to_fp4(
bf16_weights[i], use_ue8m0=True, gran_k=args.weight_gran_k)
w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, args.weight_gran_k), num_groups)
return w, w_sf
l1_weights = cast_grouped_weights_to_fp4(l1_weights)
l2_weights = cast_grouped_weights_to_fp4(l2_weights)
transformed_l1_weights, transformed_l2_weights = deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights)
transformed_l1_weights, transformed_l2_weights = deep_gemm.transform_weights_for_mega_moe(
l1_weights, l2_weights, weight_gran_k=args.weight_gran_k)
# Run fused mega MoE
# NOTES: copy x into buffer before each call because debug mode zeros the entire buffer
@@ -115,6 +118,7 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
transformed_l1_weights, transformed_l2_weights,
buffer,
cumulative_local_expert_recv_stats=cumulative_local_expert_recv_stats_fused,
recipe=(1, 1, args.weight_gran_k),
activation_clamp=args.activation_clamp,
fast_math=bool(args.fast_math)
)
@@ -166,7 +170,7 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
l1_y = torch.empty((n, intermediate_hidden * 2), dtype=torch.bfloat16, device='cuda')
deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(
recv_x, l1_weights, l1_y, handle.psum_num_recv_tokens_per_expert,
use_psum_layout=True, recipe=(1, 1, 32))
use_psum_layout=True, recipe=(1, 1, args.weight_gran_k))
# noinspection PyCallingNonCallable
l1_y = tilelang_ops.swiglu_apply_weight_to_fp8(
x=l1_y,
@@ -183,7 +187,7 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
l2_y = torch.empty((n, hidden), dtype=torch.bfloat16, device='cuda')
deep_gemm.m_grouped_fp8_fp4_gemm_nt_contiguous(
l1_y, l2_weights, l2_y, handle.psum_num_recv_tokens_per_expert,
use_psum_layout=True, recipe=(1, 1, 32))
use_psum_layout=True, recipe=(1, 1, args.weight_gran_k))
return ep_buffer.combine(l2_y, handle=handle)[0], cumulative_local_expert_recv_stats_baseline
# Check correctness (must be bitwise identical)
@@ -275,6 +279,8 @@ if __name__ == '__main__':
parser.add_argument('--num-topk', type=int, default=6, help='Number of expert selections')
parser.add_argument('--masked-ratio', type=float, default=0.0, help='Mask some expert selections')
parser.add_argument('--fast-math', type=int, default=1, help='Enable fast math (0 or 1, default: 1)')
parser.add_argument('--weight-gran-k', type=int, default=32, choices=(16, 32),
help='FP4 weight scale granularity along K')
# Test settings
parser.add_argument('--num-correctness-tests', type=int, default=None, help='Pressure test')