A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 05:00 UTC
[Misc] Fix flashinfer related tests (#33462)
Signed-off-by: esmeetu <jasonailu87@gmail.com>
Roy Wang committed
Feb 1, 2026 at 05:10 UTC
63c0889416f0d4c3979c4046c6bf41c43143080c
5 files changed
+9
-8
tests/kernels/moe/test_moe.py
+1
-1
@@ -412,7 +412,7 @@ def test_naive_block_assignment_moe(
412
monkeypatch,
413
workspace_init,
414
):
415
- current_platform.seed_everything(7)
415
+ set_random_seed(7)
416
417
monkeypatch.setenv("VLLM_FUSED_MOE_CHUNK_SIZE", str(chunk_size))
418
tests/kernels/quantization/test_flashinfer_nvfp4_scaled_mm.py
+1
-1
@@ -74,7 +74,7 @@ def get_ref_results(
74
@pytest.mark.parametrize("shape", SHAPES)
75
@pytest.mark.parametrize("seed", SEEDS)
76
@pytest.mark.parametrize("device", CUDA_DEVICES)
77
-@pytest.mark.parametrize("backend", ["cutlass", "trtllm"])
77
+@pytest.mark.parametrize("backend", ["cutlass", "cudnn", "trtllm"])
78
@pytest.mark.parametrize("autotune", [False, True])
79
@torch.inference_mode()
80
def test_flashinfer_nvfp4_gemm(
tests/kernels/quantization/test_fp8_quant.py
+2
-2
@@ -174,7 +174,7 @@ def test_static_fp8_quant_group_2d(
174
f"group_shape ({group_shape[0]}, {group_shape[1]})"
175
)
176
177
- current_platform.seed_everything(seed)
177
+ set_random_seed(seed)
178
179
x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda")
180
ref_out, scale = scaled_quantize(
@@ -202,7 +202,7 @@ def test_static_fp8_quant_1d_scale(
202
group_shape: tuple[int, int],
203
) -> None:
204
"""Test static FP8 quantization with 1D scale (per-token or per-channel)."""
205
- current_platform.seed_everything(seed)
205
+ set_random_seed(seed)
206
207
x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda")
208
ref_out, scale_2d = scaled_quantize(
vllm/model_executor/layers/quantization/utils/nvfp4_utils.py
+4
-3
@@ -154,9 +154,10 @@ def convert_to_nvfp4_linear_kernel_format(
154
)
155
layer.weight = torch.nn.Parameter(weight, requires_grad=False)
156
layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
157
- elif (
158
- backend == NvFp4LinearBackend.VLLM_CUTLASS
159
- or backend == NvFp4LinearBackend.FLASHINFER_CUTLASS
157
+ elif backend in (
158
+ NvFp4LinearBackend.VLLM_CUTLASS,
159
+ NvFp4LinearBackend.FLASHINFER_CUTLASS,
160
+ NvFp4LinearBackend.FLASHINFER_CUDNN,
161
):
162
weight, weight_scale, weights_padding_cols = prepare_weights_for_nvfp4_cutlass(
163
layer.weight.data, layer.weight_scale.data
vllm/utils/flashinfer.py
+1
-1
@@ -521,7 +521,7 @@ def flashinfer_scaled_fp4_mm(
521
assert a.stride(-1) == 1 and b.stride(-1) == 1
522
assert a.shape[1] == b.shape[1]
523
524
- if backend == "cutlass":
524
+ if backend in ("cutlass", "cudnn"):
525
block_scale_a = block_scale_a.view(torch.uint8)
526
block_scale_b = block_scale_b.view(torch.uint8)
527