Skip to content

Add Mxfp4 trtllm-gen moe unit tests #1399

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
yzh119 merged 2 commits into from
Aug 12, 2025
Merged

Add Mxfp4 trtllm-gen moe unit tests #1399

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
b3e8a01
add mxfp4 x mxfp8/bf16 moe unit tests
IwakuraRein Aug 9, 2025
be3299c
address precommit
IwakuraRein Aug 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
37 changes: 27 additions & 10 deletions tests/test_fp4_quantize.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

import pytest
import torch
from utils_fp4 import cast_from_fp4, recover_swizzled_scales, ref_nvfp4_quant
from utils_fp4 import cast_from_fp4, recover_swizzled_scales, ref_fp4_quant

from flashinfer import (
e2m1_and_ufp8sf_scale_to_float,
Expand Down Expand Up @@ -88,30 +88,47 @@ def unswizzle_sf(
@pytest.mark.parametrize("shape", SHAPES)
@pytest.mark.parametrize("seed", SEEDS)
@pytest.mark.parametrize("device", CUDA_DEVICES)
@pytest.mark.parametrize("sf_use_ue8m0", [False, True])
@pytest.mark.parametrize("is_swizzled", [False, True])
@torch.inference_mode()
def test_fp4_quantization(
dtype: torch.dtype,
shape: tuple[int, int],
seed: int,
device: str,
sf_use_ue8m0: bool,
is_swizzled: bool,
) -> None:
if not is_sm100a_supported(torch.device(device)):
pytest.skip("Nvfp4 Requires compute capability of 10 or above")
torch.set_default_device(device)
torch.manual_seed(seed)
m, n = shape
sf_vec_size = 32 if sf_use_ue8m0 else 16
x = torch.randn((m, n), dtype=dtype)
tensor_amax = torch.abs(x).max().to(torch.float32)
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
out_ref, scale_ref = ref_nvfp4_quant(x, global_scale, BLOCK_SIZE)
out, out_scale = fp4_quantize(x, global_scale, BLOCK_SIZE, False)
assert n % BLOCK_SIZE == 0, f"cols needs to be {BLOCK_SIZE} divisible"
scale_ans = recover_swizzled_scales(
out_scale.reshape(-1, n // BLOCK_SIZE).view(torch.float8_e4m3fn),
m,
n,
BLOCK_SIZE,
if sf_use_ue8m0:
global_scale = torch.tensor(1.0, dtype=torch.float32)
else:
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
out_ref, scale_ref = ref_fp4_quant(x, global_scale, sf_vec_size, sf_use_ue8m0)
out, out_scale = fp4_quantize(
x, global_scale, sf_vec_size, sf_use_ue8m0, is_swizzled
)
assert n % sf_vec_size == 0, f"cols needs to be {sf_vec_size} divisible"
if sf_use_ue8m0:
out_scale = (out_scale.to(torch.int32) << 23).view(torch.float32)
else:
out_scale = out_scale.view(torch.float8_e4m3fn).to(torch.float32)
if is_swizzled:
scale_ans = recover_swizzled_scales(
out_scale.reshape(-1, n // sf_vec_size),
m,
n,
sf_vec_size,
)
else:
scale_ans = out_scale
out_ans = cast_from_fp4(out).reshape(m, n)
torch.testing.assert_close(out_ans, out_ref, rtol=1e0, atol=1e-1)
torch.testing.assert_close(scale_ans, scale_ref, rtol=1e-1, atol=1e-1)
Expand Down
4 changes: 2 additions & 2 deletions tests/test_trtllm_gen_context.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@

import pytest
import torch
from utils_fp4 import cast_from_fp4, recover_swizzled_scales, ref_nvfp4_quant
from utils_fp4 import cast_from_fp4, recover_swizzled_scales, ref_fp4_quant

import flashinfer
from flashinfer.utils import FP4Tensor
Expand Down Expand Up @@ -437,7 +437,7 @@ def test_trtllm_batch_prefill(

if o_dtype == "nvfp4":
output = cast_from_fp4(output)
output_ref, out_scale_factor_ref = ref_nvfp4_quant(output_ref, o_sf_scale, 16)
output_ref, out_scale_factor_ref = ref_fp4_quant(output_ref, o_sf_scale, 16)
out_scale_factor = recover_swizzled_scales(
out_scale_factor,
output.shape[0],
Expand Down
4 changes: 2 additions & 2 deletions tests/test_trtllm_gen_decode.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,7 @@
import pytest
import torch
import torch.nn.functional as F
from utils_fp4 import cast_from_fp4, recover_swizzled_scales, ref_nvfp4_quant
from utils_fp4 import cast_from_fp4, recover_swizzled_scales, ref_fp4_quant

import flashinfer
from flashinfer.utils import FP4Tensor
Expand Down Expand Up @@ -328,7 +328,7 @@ def test_trtllm_batch_decode_fmha(

if o_dtype == "nvfp4":
output = cast_from_fp4(output)
output_ref, out_scale_factor_ref = ref_nvfp4_quant(output_ref, o_sf_scale, 16)
output_ref, out_scale_factor_ref = ref_fp4_quant(output_ref, o_sf_scale, 16)
out_scale_factor = recover_swizzled_scales(
out_scale_factor,
output.shape[0],
Expand Down
Loading