diff --git a/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py b/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py index 88c9e5f867c..5a737743961 100644 --- a/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py +++ b/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py @@ -1,13 +1,55 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import torch +from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend +from vllm.model_executor.layers.quantization.utils import flashinfer_fp4_moe +from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( + prepare_nvfp4_moe_layer_for_fi_or_cutlass, +) from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( align_trtllm_fp4_moe_hidden_dim_for_fi, ) +def test_shared_nvfp4_input_scales_have_writable_storage(monkeypatch): + monkeypatch.setattr(flashinfer_fp4_moe, "swizzle_blockscale", lambda x: x) + + num_experts = 3 + layer = SimpleNamespace(activation=SimpleNamespace(is_gated=False)) + w13 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8) + w2 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8) + w13_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn) + w2_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn) + weight_scale = torch.ones(num_experts) + + outputs = prepare_nvfp4_moe_layer_for_fi_or_cutlass( + backend=NvFp4MoeBackend.FLASHINFER_CUTLASS, + layer=layer, + w13=w13, + w13_scale=w13_scale, + w13_scale_2=weight_scale, + a13_scale=torch.tensor([1.0, 2.0, 3.0]), + w2=w2, + w2_scale=w2_scale, + w2_scale_2=weight_scale, + a2_scale=torch.tensor([4.0, 5.0, 6.0]), + is_act_and_mul=False, + ) + a13_scale, a2_scale = outputs[3], outputs[7] + + torch.testing.assert_close(a13_scale, torch.full((num_experts,), 3.0)) + torch.testing.assert_close(a2_scale, torch.full((num_experts,), 6.0)) + distinct_values = torch.arange(num_experts, dtype=torch.float32) + a13_scale.copy_(distinct_values) + a2_scale.copy_(distinct_values) + torch.testing.assert_close(a13_scale, distinct_values) + torch.testing.assert_close(a2_scale, distinct_values) + + def test_align_trtllm_fp4_moe_hidden_dim_noop(): w13 = torch.arange(2 * 8 * 256, dtype=torch.uint8).reshape(2, 8, 256) w13_scale = torch.arange(2 * 8 * 32, dtype=torch.uint8).reshape(2, 8, 32) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 6f0e237785e..bab3dee649b 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -109,8 +109,8 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( # Global scaling factors (same as other FlashInfer backends). num_experts = w13.shape[0] - a13_scale = a13_scale.max().to(torch.float32).expand(num_experts) - a2_scale = a2_scale.max().to(torch.float32).expand(num_experts) + a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts) + a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts) half = w13.shape[1] // 2 w13 = torch.cat([w13[:, half:], w13[:, :half]], dim=1) @@ -338,8 +338,8 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( # For some FI kernels, the input scales are shared by all experts. if is_global_sf_supported_for_nvfp4_backend(backend): num_experts = w13.shape[0] - a13_scale = a13_scale.max().to(torch.float32).expand(num_experts) - a2_scale = a2_scale.max().to(torch.float32).expand(num_experts) + a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts) + a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts) else: a13_scale = a13_scale.max(dim=1).values.to(torch.float32)