feat(megamoe): add fp4 group16 to group32 requant path
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:
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user