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

21
deep_gemm/mega/__init__.py
View File

@@ -1,6 +1,6 @@
import torch
from typing import Tuple, Optional
from ..utils.math import align
from ..utils.math import align, requant_fp4_to_gran_k
# noinspection PyBroadException
try:
@@ -176,6 +176,25 @@ def transform_weights_for_mega_moe(
return l1_weights, l2_weights
def prepare_fp4_weights_for_mega_moe(
l1_weights: Tuple[torch.Tensor, torch.Tensor],
l2_weights: Tuple[torch.Tensor, torch.Tensor],
source_weight_gran_k: int = 32,
runtime_weight_gran_k: int = 32,
) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
if source_weight_gran_k != runtime_weight_gran_k:
if source_weight_gran_k != 16 or runtime_weight_gran_k != 32:
raise RuntimeError(
f'Unsupported MegaMoE FP4 weight granularity conversion: '
f'{source_weight_gran_k} -> {runtime_weight_gran_k}')
l1_weights = requant_fp4_to_gran_k(
l1_weights[0], l1_weights[1], source_weight_gran_k, runtime_weight_gran_k)
l2_weights = requant_fp4_to_gran_k(
l2_weights[0], l2_weights[1], source_weight_gran_k, runtime_weight_gran_k)
return transform_weights_for_mega_moe(
l1_weights, l2_weights, weight_gran_k=runtime_weight_gran_k)
def fp8_mega_moe(y: torch.Tensor,
l1_weights: Tuple[torch.Tensor, torch.Tensor],
l2_weights: Tuple[torch.Tensor, torch.Tensor],