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

feat(megamoe): expose fp4 weight preparation helper

Browse Source 
Add a top-level MegaMoE helper that handles source/runtime FP4 weight granularity before applying the existing MegaMoE weight layout transform.

Use the helper from the synthetic MegaMoE benchmark so SGLang can later follow the same contract for GLM-5.2 NVFP4 group16 checkpoints.

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

Tested: git diff --check

Not-tested: CUDA build, SM100/B300 runtime, and GLM-5.2 accuracy validation are not available locally.
This commit is contained in:
LuminolT
2026-07-08 18:36:40 +08:00
parent 8ad348fb11
commit 2c7543130b
4 changed files with 33 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
tests/test_mega_moe.py
View File

@@ -7,7 +7,7 @@ import torch.distributed as dist
from typing import Tuple
import deep_gemm
from deep_gemm.utils import per_token_cast_to_fp4, per_token_cast_to_fp8, requant_fp4_to_gran_k
from deep_gemm.utils import per_token_cast_to_fp4, per_token_cast_to_fp8
from deep_gemm.utils.dist import dist_print, init_dist, uneven_all_gather
from deep_gemm.testing import bench_kineto
@@ -97,15 +97,18 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
for i in range(num_groups):
w[i], w_sf[i] = per_token_cast_to_fp4(
bf16_weights[i], use_ue8m0=True, gran_k=args.weight_gran_k)
if args.requant_group16_to_group32:
w, w_sf = requant_fp4_to_gran_k(w, w_sf, src_gran_k=16, dst_gran_k=32)
w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, runtime_weight_gran_k), num_groups)
if not args.requant_group16_to_group32:
w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, runtime_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, weight_gran_k=runtime_weight_gran_k)
if args.requant_group16_to_group32:
transformed_l1_weights, transformed_l2_weights = deep_gemm.prepare_fp4_weights_for_mega_moe(
l1_weights, l2_weights, source_weight_gran_k=16, runtime_weight_gran_k=32)
else:
transformed_l1_weights, transformed_l2_weights = deep_gemm.transform_weights_for_mega_moe(
l1_weights, l2_weights, weight_gran_k=runtime_weight_gran_k)
# Run fused mega MoE
# NOTES: copy x into buffer before each call because debug mode zeros the entire buffer