Skip to content

[https://nvbugs/6801229][fix] Fix NVFP4 gather test references and skip empty FP4 quantize launches #19761

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
zhangcl merged 2 commits into from
Oct 8, 2026
Merged

[https://nvbugs/6801229][fix] Fix NVFP4 gather test references and skip empty FP4 quantize launches #19761

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
4ef50d5
[None][fix] Fix NVFP4 gather test references and skip empty FP4 quant…
zhangcl Oct 1, 2026
4c6b146
[https://nvbugs/6801229][fix] Apply the Relu2 FP32 reference to the B…
zhangcl Oct 8, 2026
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
11 changes: 11 additions & 0 deletions cpp/tensorrt_llm/kernels/quantization.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,11 @@ template <typename T, int SF_VEC_SIZE>
void invokeFP4Quantization(int b, int m, int n, T const* input, float const* SFScale, int64_t* output, int32_t* SFOuput,
bool useUE8M0, QuantizationSFLayout layout, int multiProcessorCount, cudaStream_t stream)
{
// Skip empty inputs: a zero-size launch is invalid and leaves a pending CUDA error.
if (b <= 0 || m <= 0 || n <= 0)
{
return;
}
#ifdef ENABLE_FP8
if constexpr (std::is_same_v<T, __nv_fp8_e4m3>)
{
Expand Down Expand Up @@ -197,6 +202,12 @@ void invokeMxFP8Quantization(int b, int m, int n, int padded_n, T const* input,
static_assert(SF_OUTPUT_VEC_SIZE == 32 || SF_OUTPUT_VEC_SIZE == SF_VEC_SIZE,
"MXFP8 output SF vector size must be 32 or match the quantization SF vector size.");

// Skip empty inputs: a zero-size launch is invalid and leaves a pending CUDA error.
if (b <= 0 || m <= 0 || padded_n <= 0)
{
return;
}

// Grid, Block size.
// Each thread converts 8 values.
dim3 block(std::min(int(padded_n / CVT_ELTS_PER_THREAD), 512));
Expand Down
75 changes: 60 additions & 15 deletions tests/unittest/_torch/thop/parallel/test_cute_dsl_moe.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,29 @@
from tensorrt_llm._utils import get_sm_version


def _nvfp4_quantize_fp32_ref(
x: torch.Tensor, inverse_global_scale: torch.Tensor
) -> tuple[torch.Tensor, torch.Tensor]:
"""Quantize without the BF16 rounding absent from the fused FP32 epilogue."""
m, n = x.shape
blocks = x.reshape(m, n // 16, 16)
scales = (blocks.abs().amax(dim=-1) / 6 * inverse_global_scale).to(torch.float8_e4m3fn)
inverse_scales = (inverse_global_scale / scales.float()).clamp(
max=torch.finfo(torch.float32).max
)
normalized = (blocks * inverse_scales.unsqueeze(-1)).reshape(m, n)
magnitudes = normalized.abs().contiguous()
# E2M1 magnitudes are 0, 0.5, 1, 1.5, 2, 3, 4, 6. At midpoint ties,
# choose the even code, matching round-to-nearest-even conversion.
midpoints = torch.tensor([0.25, 0.75, 1.25, 1.75, 2.5, 3.5, 5.0], device=x.device)
codes = torch.bucketize(magnitudes, midpoints)
for index in (1, 3, 5):
codes += (magnitudes == midpoints[index]).to(codes.dtype)
codes = codes.to(torch.uint8) | (torch.signbit(normalized).to(torch.uint8) << 3)
packed = codes[:, 0::2] | (codes[:, 1::2] << 4)
return packed, swizzle_sf(scales.view(torch.uint8), m, n)


def swiglu_ref(x: torch.Tensor, swiglu_limit: float = float("inf")) -> torch.Tensor:
x, gate = x.chunk(2, dim=-1)
if swiglu_limit != float("inf"):
Expand Down Expand Up @@ -1276,8 +1299,12 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_blackwell(
permuted_idx_to_expanded_idx_list = permuted_idx_to_expanded_idx.cpu().tolist()
tile_idx_to_mn_limit_list = tile_idx_to_mn_limit.cpu().tolist()

a_gathered = torch.empty(max_num_permuted_tokens, hidden_size // 2, dtype=a.dtype)
a_sf_gathered = torch.empty(
# Padding participates in the reference GEMM and output-scale reduction.
# Initialize both FP4 values and scales so padding cannot introduce NaNs.
a_gathered = torch.zeros(max_num_permuted_tokens, hidden_size // 2, dtype=torch.uint8).view(
a.dtype
)
a_sf_gathered = torch.zeros(
max_num_permuted_tokens, hidden_size // sf_vec_size, dtype=a_sf.dtype
)
for i in range(num_valid_permuted_tokens):
Expand Down Expand Up @@ -1306,12 +1333,15 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_blackwell(
tile_idx_to_group_idx,
num_non_exiting_tiles,
tile_size=tile_size,
output_dtype=torch.bfloat16,
output_dtype=torch.float32 if activation_type == ActivationType.Relu2 else torch.bfloat16,
scaling_vector_size=sf_vec_size,
)
c_ref = apply_activation_ref(c_ref, activation_type, swiglu_limit)
global_sf = c_ref[:num_valid_permuted_tokens].abs().max().float() / (448 * 6)
c_ref, c_sf_ref = torch.ops.trtllm.fp4_quantize(c_ref, 1 / global_sf, sf_vec_size, False)
if activation_type == ActivationType.Relu2:
c_ref, c_sf_ref = _nvfp4_quantize_fp32_ref(c_ref, 1 / global_sf)
else:
c_ref, c_sf_ref = torch.ops.trtllm.fp4_quantize(c_ref, 1 / global_sf, sf_vec_size, False)
Comment thread
zhangcl marked this conversation as resolved.
Resolved

# Call gather kernel (single-B)
c, c_sf = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_blackwell(
Expand Down Expand Up @@ -1482,8 +1512,12 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_rubin(
permuted_idx_to_expanded_idx_list = permuted_idx_to_expanded_idx.cpu().tolist()
tile_idx_to_mn_limit_list = tile_idx_to_mn_limit.cpu().tolist()

a_gathered = torch.empty(max_num_permuted_tokens, hidden_size // 2, dtype=a.dtype)
a_sf_gathered = torch.empty(
# Padding participates in the reference GEMM and output-scale reduction.
# Initialize both FP4 values and scales so padding cannot introduce NaNs.
a_gathered = torch.zeros(max_num_permuted_tokens, hidden_size // 2, dtype=torch.uint8).view(
a.dtype
)
a_sf_gathered = torch.zeros(
max_num_permuted_tokens, hidden_size // sf_vec_size, dtype=a_sf.dtype
)
for i in range(num_valid_permuted_tokens):
Expand Down Expand Up @@ -1511,12 +1545,15 @@ def test_nvfp4_gather_grouped_gemm_act_fusion_rubin(
tile_idx_to_group_idx,
num_non_exiting_tiles,
tile_size=tile_size,
output_dtype=torch.bfloat16,
output_dtype=torch.float32 if activation_type == ActivationType.Relu2 else torch.bfloat16,
scaling_vector_size=sf_vec_size,
)
c_ref = apply_activation_ref(c_ref, activation_type)
global_sf = c_ref[:num_valid_permuted_tokens].abs().max().float() / (448 * 6)
c_ref, c_sf_ref = torch.ops.trtllm.fp4_quantize(c_ref, 1 / global_sf, sf_vec_size, False)
if activation_type == ActivationType.Relu2:
c_ref, c_sf_ref = _nvfp4_quantize_fp32_ref(c_ref, 1 / global_sf)
else:
c_ref, c_sf_ref = torch.ops.trtllm.fp4_quantize(c_ref, 1 / global_sf, sf_vec_size, False)

# Call Rubin gather kernel
c, c_sf = torch.ops.trtllm.cute_dsl_nvfp4_gather_grouped_gemm_act_fusion_rubin(
Expand Down Expand Up @@ -1653,8 +1690,12 @@ def test_nvfp4_gather_grouped_gemm_situ_rubin(tile_size: int):
permuted_idx_to_expanded_idx_list = permuted_idx_to_expanded_idx.cpu().tolist()
tile_idx_to_mn_limit_list = tile_idx_to_mn_limit.cpu().tolist()

a_gathered = torch.empty(max_num_permuted_tokens, hidden_size // 2, dtype=a.dtype)
a_sf_gathered = torch.empty(
# Padding participates in the reference GEMM and output-scale reduction.
# Initialize both FP4 values and scales so padding cannot introduce NaNs.
a_gathered = torch.zeros(max_num_permuted_tokens, hidden_size // 2, dtype=torch.uint8).view(
a.dtype
)
a_sf_gathered = torch.zeros(
max_num_permuted_tokens, hidden_size // sf_vec_size, dtype=a_sf.dtype
)
for i in range(num_valid_permuted_tokens):
Expand Down Expand Up @@ -1876,12 +1917,13 @@ def test_nvfp4_gather_grouped_gemm_swiglu_rubin_small_tokens(
)
num_valid_permuted_tokens = n_tiles * tile_size
for i in range(min(num_valid_permuted_tokens, max_num_permuted_tokens)):
# Routing metadata defines valid rows; expanded index zero is valid.
if i >= tile_idx_to_mn_limit[i // tile_size].item():
continue
expanded_idx = permuted_idx_to_expanded_idx[i].item()
if expanded_idx > 0 or i == 0:
token_id = expanded_idx // top_k
if token_id < num_tokens:
a_gathered[i] = a[token_id]
a_sf_gathered[i] = a_sf_unswizzled[token_id]
token_id = expanded_idx // top_k
a_gathered[i] = a[token_id]
a_sf_gathered[i] = a_sf_unswizzled[token_id]

a_sf_gathered_swizzled = swizzle_sf(
a_sf_gathered.view(max_num_permuted_tokens, hidden_size // sf_vec_size),
Expand Down Expand Up @@ -2741,6 +2783,9 @@ def test_cute_dsl_nvfp4_quantize_empty_input_rubin():

assert x.shape == (0, hidden_size // 2)
assert x_sf.shape == (0, hidden_size // scaling_vector_size)
# The empty input must not leave a pending CUDA error for the next launch.
torch.ones(1, device="cuda").add_(1)
torch.cuda.synchronize()


def test_sm107_nvfp4_tile512_fallback_uses_two_cta_cluster():
Expand Down