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

fix(megamoe): normalize fp4 weight preparation contract

Browse Source 
Separate source and destination FP4 scale packing in requant_fp4_to_gran_k so group16-to-group32 conversion always recomputes UE8M0 runtime scales by default.

Make prepare_fp4_weights_for_mega_moe accept raw grouped FP4 weights and scales, then perform optional requantization, DeepGEMM scale layout transform, and MegaMoE UTCCP weight transform internally.

Update the MegaMoE synthetic benchmark so baseline grouped GEMM uses runtime-layout weights while fused MegaMoE uses transformed weights from the same raw source tensors.

Tested: PYTHONPYCACHEPREFIX=/private/tmp/deepgemm_pycache python3 -m py_compile deep_gemm/__init__.py deep_gemm/mega/__init__.py deep_gemm/utils/math.py tests/test_layout.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:48:04 +08:00
parent 2c7543130b
commit 007c645f87
5 changed files with 86 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
deep_gemm/mega/__init__.py
View File

@@ -1,6 +1,6 @@
import torch
from typing import Tuple, Optional
from ..utils.math import align, requant_fp4_to_gran_k
from ..utils.math import align, requant_fp4_to_gran_k, unpack_ue8m0_from_int
# noinspection PyBroadException
try:
@@ -176,21 +176,45 @@ 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]]:
def _prepare_raw_fp4_weight_for_mega_moe(
weights: Tuple[torch.Tensor, torch.Tensor],
source_weight_gran_k: int,
runtime_weight_gran_k: int,
source_scale_packed_ue8m0: bool,
) -> Tuple[torch.Tensor, torch.Tensor]:
weight, weight_sf = weights
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)
weight, weight_sf = requant_fp4_to_gran_k(
weight, weight_sf,
source_weight_gran_k, runtime_weight_gran_k,
src_scale_packed_ue8m0=source_scale_packed_ue8m0)
source_scale_packed_ue8m0 = False
if source_scale_packed_ue8m0:
weight_sf = unpack_ue8m0_from_int(weight_sf)
num_groups, mn, packed_k = weight.shape
weight_sf = _C.transform_sf_into_required_layout(
weight_sf, mn, packed_k * 2, (1, runtime_weight_gran_k), num_groups)
return weight, weight_sf
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,
source_scale_packed_ue8m0: bool = False,
) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]:
l1_weights = _prepare_raw_fp4_weight_for_mega_moe(
l1_weights, source_weight_gran_k, runtime_weight_gran_k,
source_scale_packed_ue8m0)
l2_weights = _prepare_raw_fp4_weight_for_mega_moe(
l2_weights, source_weight_gran_k, runtime_weight_gran_k,
source_scale_packed_ue8m0)
return transform_weights_for_mega_moe(
l1_weights, l2_weights, weight_gran_k=runtime_weight_gran_k)

9
deep_gemm/utils/math.py
View File

@@ -144,7 +144,8 @@ def cast_back_from_fp4(packed: torch.Tensor, sf: torch.Tensor, gran_k: int = 128
def requant_fp4_to_gran_k(packed: torch.Tensor, sf: torch.Tensor, src_gran_k: int,
dst_gran_k: int, use_packed_ue8m0: bool = False) -> Tuple[torch.Tensor, torch.Tensor]:
dst_gran_k: int, src_scale_packed_ue8m0: bool = False,
dst_scale_packed_ue8m0: bool = False) -> Tuple[torch.Tensor, torch.Tensor]:
assert packed.dtype == torch.int8
assert packed.dim() >= 2
assert sf.shape[:-1] == packed.shape[:-1]
@@ -153,10 +154,10 @@ def requant_fp4_to_gran_k(packed: torch.Tensor, sf: torch.Tensor, src_gran_k: in
packed_2d = packed.reshape(-1, packed.size(-1))
sf_2d = sf.reshape(-1, sf.size(-1))
restored = cast_back_from_fp4(packed_2d, sf_2d, src_gran_k, use_packed_ue8m0)
restored = cast_back_from_fp4(packed_2d, sf_2d, src_gran_k, src_scale_packed_ue8m0)
repacked, rescaled = per_token_cast_to_fp4(
restored, use_ue8m0=use_packed_ue8m0, gran_k=dst_gran_k,
use_packed_ue8m0=use_packed_ue8m0)
restored, use_ue8m0=True, gran_k=dst_gran_k,
use_packed_ue8m0=dst_scale_packed_ue8m0)
return (
repacked.reshape(original_packed_shape),

18
megamoe-research-reports/glm52_nvfp4_group16_notes.md
View File

@@ -20,7 +20,8 @@ carried an implicit group32 assumption.
existing block32 kernels can consume.
- `prepare_fp4_weights_for_mega_moe(...)` wraps the source/runtime granularity
decision for callers such as SGLang: source group16 can be requantized to
runtime group32 and then passed through the normal MegaMoE weight transform.
runtime group32, converted into DeepGEMM scale layout, and then passed through
the normal MegaMoE weight transform.
- The fused SM100 MegaMoE compute API now performs an explicit capability check
for `recipe=(1, 1, 16)` instead of failing earlier with
`Unknown SF transformation`.
@@ -39,6 +40,21 @@ group16 correctly requires auditing at least:
Until that kernel work is complete and validated on B300/SM100, group16 should
be treated as layout-supported but fused-compute unsupported.
## Scale layout contract
There are three distinct scale states:
1. raw checkpoint scale: per-weight-group scale from model loading or synthetic
quantization, either float UE8M0 values or packed UE8M0 integers;
2. DeepGEMM runtime scale layout: output of `transform_sf_into_required_layout`,
packed, MN-major, TMA-aligned;
3. MegaMoE weight scale layout: runtime scale layout after the UTCCP transpose
required by the fused MegaMoE kernel.
`prepare_fp4_weights_for_mega_moe(...)` accepts state 1 and returns state 3.
`transform_weights_for_mega_moe(...)` remains a lower-level helper that accepts
state 2 and returns state 3.
## Requantization path
If the SM100 MXF4 MMA path cannot consume group16 scales directly, the fallback

11
tests/test_layout.py
View File

@@ -113,11 +113,22 @@ def test_fp4_requant_granularity() -> None:
repacked, rescaled = requant_fp4_to_gran_k(fp4_g16[0], fp4_g16[1], 16, 32)
restored = cast_back_from_fp4(fp4_g16[0], fp4_g16[1], gran_k=16)
ref_repacked, ref_rescaled = per_token_cast_to_fp4(restored, use_ue8m0=True, gran_k=32)
fp4_g16_packed = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=16, use_packed_ue8m0=True)
repacked_from_packed, rescaled_from_packed = requant_fp4_to_gran_k(
fp4_g16_packed[0], fp4_g16_packed[1], 16, 32, src_scale_packed_ue8m0=True)
ref_repacked_packed, ref_rescaled_packed = per_token_cast_to_fp4(
restored, use_ue8m0=True, gran_k=32, use_packed_ue8m0=True)
repacked_to_packed, rescaled_to_packed = requant_fp4_to_gran_k(
fp4_g16[0], fp4_g16[1], 16, 32, dst_scale_packed_ue8m0=True)
assert repacked.shape == fp4_g16[0].shape
assert rescaled.shape[-1] == ceil_div(n, 32)
assert torch.equal(repacked, ref_repacked)
assert torch.equal(rescaled, ref_rescaled)
assert torch.equal(repacked_from_packed, ref_repacked)
assert torch.equal(rescaled_from_packed, ref_rescaled)
assert torch.equal(repacked_to_packed, ref_repacked_packed)
assert torch.equal(rescaled_to_packed, ref_rescaled_packed)
print(f' > Requant ({m=}, {n=}): group16 -> group32')
print()

29
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
from deep_gemm.utils import per_token_cast_to_fp4, per_token_cast_to_fp8, requant_fp4_to_gran_k
from deep_gemm.utils.dist import dist_print, init_dist, uneven_all_gather
from deep_gemm.testing import bench_kineto
@@ -97,18 +97,25 @@ 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 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)
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)
def prepare_grouped_runtime_weights(weights: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
w, w_sf = weights
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)
num_groups, n, packed_k = w.shape
w_sf = deep_gemm.transform_sf_into_required_layout(
w_sf, n, packed_k * 2, (1, runtime_weight_gran_k), num_groups)
return w, w_sf
source_l1_weights = cast_grouped_weights_to_fp4(l1_weights)
source_l2_weights = cast_grouped_weights_to_fp4(l2_weights)
l1_weights = prepare_grouped_runtime_weights(source_l1_weights)
l2_weights = prepare_grouped_runtime_weights(source_l2_weights)
transformed_l1_weights, transformed_l2_weights = deep_gemm.prepare_fp4_weights_for_mega_moe(
source_l1_weights, source_l2_weights,
source_weight_gran_k=args.weight_gran_k,
runtime_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