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

feat(megamoe): add fp4 group16 to group32 requant path

Browse Source 
Add a utility and synthetic benchmark path for converting FP4 group16 tensors into group32 tensors that the existing SM100 block32 MegaMoE kernels can consume.

Document that this is a requantization path rather than a lossless metadata rewrite, so GLM-5.2 accuracy validation is still required before production use.

Tested: PYTHONPYCACHEPREFIX=/private/tmp/deepgemm_pycache python3 -m py_compile 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:34:33 +08:00
parent 79fcfd6abf
commit 8ad348fb11
4 changed files with 72 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
tests/test_layout.py
View File

@@ -4,6 +4,7 @@ from deep_gemm.testing import bench_kineto, count_bytes, get_arch_major
from deep_gemm.utils import (
align, ceil_div,
per_token_cast_to_fp8, per_channel_cast_to_fp8,
per_token_cast_to_fp4, cast_back_from_fp4, requant_fp4_to_gran_k,
get_tma_aligned_size,
get_mn_major_tma_aligned_tensor,
get_mn_major_tma_aligned_packed_ue8m0_tensor,
@@ -104,9 +105,27 @@ def test_k_grouped_sf_layout_kernels() -> None:
print()
def test_fp4_requant_granularity() -> None:
print('Testing FP4 requant granularity:')
for m, n in ((128, 7168), (257, 3072)):
x = torch.randn((m, n), dtype=torch.bfloat16, device='cuda')
fp4_g16 = per_token_cast_to_fp4(x, use_ue8m0=True, gran_k=16)
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)
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)
print(f' > Requant ({m=}, {n=}): group16 -> group32')
print()
if __name__ == '__main__':
torch.manual_seed(1)
random.seed(1)
test_sf_layout_kernels()
test_k_grouped_sf_layout_kernels()
test_fp4_requant_granularity()

19
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
@@ -47,7 +47,9 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
hidden, intermediate_hidden = args.hidden, args.intermediate_hidden
num_experts, num_topk = args.num_experts, args.num_topk
num_experts_per_rank = num_experts // num_ranks
runtime_weight_gran_k = 32 if args.requant_group16_to_group32 else args.weight_gran_k
assert num_tokens <= num_max_tokens_per_rank
assert not args.requant_group16_to_group32 or args.weight_gran_k == 16
# Allocate symmetric memory
buffer = deep_gemm.get_symm_buffer_for_mega_moe(
@@ -95,13 +97,15 @@ 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)
w_sf = deep_gemm.transform_sf_into_required_layout(w_sf, n, k, (1, args.weight_gran_k), num_groups)
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)
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=args.weight_gran_k)
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
@@ -118,7 +122,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),
recipe=(1, 1, runtime_weight_gran_k),
activation_clamp=args.activation_clamp,
fast_math=bool(args.fast_math)
)
@@ -129,6 +133,7 @@ def test(local_rank: int, num_local_ranks: int, args: argparse.Namespace):
dist_print(f' > Hidden: {hidden}', once_in_node=True)
dist_print(f' > Intermediate: {intermediate_hidden}', once_in_node=True)
dist_print(f' > Experts: {num_topk}/{num_experts}', once_in_node=True)
dist_print(f' > Weight granularity: source={args.weight_gran_k}, runtime={runtime_weight_gran_k}', once_in_node=True)
dist_print(f' > Buffer: {buffer.buffer.nbytes / 2 ** 30:.3f} GiB', once_in_node=True)
dist_print(once_in_node=True)
@@ -170,7 +175,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, args.weight_gran_k))
use_psum_layout=True, recipe=(1, 1, runtime_weight_gran_k))
# noinspection PyCallingNonCallable
l1_y = tilelang_ops.swiglu_apply_weight_to_fp8(
x=l1_y,
@@ -187,7 +192,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, args.weight_gran_k))
use_psum_layout=True, recipe=(1, 1, runtime_weight_gran_k))
return ep_buffer.combine(l2_y, handle=handle)[0], cumulative_local_expert_recv_stats_baseline
# Check correctness (must be bitwise identical)
@@ -281,6 +286,8 @@ if __name__ == '__main__':
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')
parser.add_argument('--requant-group16-to-group32', action='store_true',
help='Simulate loading a group16 FP4 checkpoint and requantizing it to DeepGEMM group32')
# Test settings
parser.add_argument('--num-correctness-tests', type=int, default=None, help='Pressure test')