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 5, 11:00 UTC
[Feature] Batch invariant nvfp4 linear support (#39322)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Apr 8, 2026 at 16:29 UTC
201813724291da44c1b290211431f2a0a1ccf503
2 files changed
+113
-5
tests/v1/determinism/test_nvfp4_batch_invariant.py
new
+100
@@ -0,0 +1,100 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import contextlib
4
+import os
5
+import random
6
+
7
+import pytest
8
+import torch
9
+from utils import (
10
+ _extract_step_logprobs,
11
+ _random_prompt,
12
+ skip_unsupported,
13
+)
14
+
15
+from vllm import LLM, SamplingParams
16
+
17
+pytestmark = pytest.mark.skipif(
18
+ not hasattr(torch, "float8_e4m3fn"),
19
+ reason="NVFP4 tests require torch.float8_e4m3fn support.",
20
+)
21
+
22
+NVFP4_TEST_MODEL = os.getenv(
23
+ "VLLM_TEST_NVFP4_MODEL", "nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4"
24
+)
25
+
26
+
27
+def _make_llm(max_num_seqs: int, backend: str) -> LLM:
28
+ return LLM(
29
+ model=NVFP4_TEST_MODEL,
30
+ max_num_seqs=max_num_seqs,
31
+ gpu_memory_utilization=float(
32
+ os.getenv("VLLM_NVFP4_TEST_GPU_MEMORY_UTILIZATION", "0.05")
33
+ ),
34
+ max_model_len=int(os.getenv("VLLM_NVFP4_TEST_MAX_MODEL_LEN", "2048")),
35
+ dtype="auto",
36
+ tensor_parallel_size=int(os.getenv("VLLM_NVFP4_TEST_TP_SIZE", "1")),
37
+ enable_prefix_caching=False,
38
+ enforce_eager=True,
39
+ attention_config={"backend": backend},
40
+ )
41
+
42
+
43
+@skip_unsupported
44
+@pytest.mark.parametrize("backend", ["FLASH_ATTN"])
45
+def test_dense_nvfp4_generation_is_deterministic_across_batch_sizes_e2e(backend):
46
+ seed = int(os.getenv("VLLM_TEST_SEED", "12345"))
47
+ random.seed(seed)
48
+
49
+ num_trials = int(os.getenv("VLLM_NVFP4_NEEDLE_TRIALS", "2"))
50
+ max_batch_size = int(os.getenv("VLLM_NVFP4_NEEDLE_BATCH_SIZE", "8"))
51
+ min_random_prompt = int(os.getenv("VLLM_NVFP4_MIN_PROMPT", "32"))
52
+ max_random_prompt = int(os.getenv("VLLM_NVFP4_MAX_PROMPT", "96"))
53
+ assert max_batch_size >= 2, "Batch size should be >= 2 to test invariance."
54
+
55
+ sampling = SamplingParams(
56
+ temperature=float(os.getenv("VLLM_NVFP4_NEEDLE_TEMPERATURE", "0.6")),
57
+ top_p=float(os.getenv("VLLM_NVFP4_NEEDLE_TOP_P", "0.95")),
58
+ max_tokens=int(os.getenv("VLLM_NVFP4_NEEDLE_MAX_TOKENS", "16")),
59
+ seed=20240919,
60
+ logprobs=5,
61
+ )
62
+ needle_prompt = "Write one factual sentence about the moon."
63
+
64
+ llm = None
65
+ baseline_completion = None
66
+ baseline_logprobs = None
67
+ try:
68
+ llm = _make_llm(max_num_seqs=max_batch_size, backend=backend)
69
+ baseline_output = llm.generate([needle_prompt], sampling, use_tqdm=False)[0]
70
+ baseline_completion = baseline_output.outputs[0]
71
+ baseline_logprobs, baseline_token_ids = _extract_step_logprobs(baseline_output)
72
+ assert baseline_logprobs is not None
73
+ assert baseline_token_ids is not None
74
+ for _ in range(num_trials):
75
+ batch_size = random.randint(max_batch_size // 2, max_batch_size)
76
+ needle_pos = random.randint(0, batch_size - 1)
77
+ prompts: list[str] = []
78
+ for idx in range(batch_size):
79
+ if idx == needle_pos:
80
+ prompts.append(needle_prompt)
81
+ else:
82
+ prompts.append(_random_prompt(min_random_prompt, max_random_prompt))
83
+
84
+ outputs = llm.generate(prompts, sampling, use_tqdm=False)
85
+ needle_output = outputs[needle_pos]
86
+ needle_completion = needle_output.outputs[0]
87
+ needle_logprobs, needle_token_ids = _extract_step_logprobs(needle_output)
88
+ assert needle_logprobs is not None
89
+ assert needle_token_ids is not None
90
+
91
+ assert needle_output.prompt == needle_prompt
92
+ assert baseline_completion is not None
93
+ assert baseline_logprobs is not None
94
+ assert needle_completion.token_ids == baseline_completion.token_ids
95
+ assert needle_completion.text == baseline_completion.text
96
+ torch.testing.assert_close(needle_logprobs, baseline_logprobs)
97
+ finally:
98
+ if llm is not None:
99
+ with contextlib.suppress(Exception):
100
+ llm.shutdown()
vllm/model_executor/layers/quantization/utils/nvfp4_utils.py
+13
-5
@@ -109,6 +109,13 @@ def select_nvfp4_linear_backend() -> NvFp4LinearBackend:
109
Select the best available NVFP4 GEMM backend based on environment
110
configuration and platform capabilities.
111
"""
112
+ if envs.VLLM_BATCH_INVARIANT:
113
+ logger.info_once(
114
+ "VLLM_BATCH_INVARIANT forces NVFP4 linear to use the emulation "
115
+ "backend for deterministic execution."
116
+ )
117
+ return NvFp4LinearBackend.EMULATION
118
+
119
selected_backend: NvFp4LinearBackend | None = None
120
121
if envs.VLLM_USE_FBGEMM:
@@ -259,6 +266,8 @@ def apply_nvfp4_linear(
266
alpha = layer.alpha
267
output_size = layer.output_size_per_partition
268
input_size = layer.input_size_per_partition
269
+ output_dtype = x.dtype
270
+ output_shape = [*x.shape[:-1], output_size]
271
272
if backend == NvFp4LinearBackend.MARLIN:
273
return apply_fp4_marlin_linear(
@@ -272,20 +281,19 @@ def apply_nvfp4_linear(
281
bias=bias,
282
)
283
elif backend == NvFp4LinearBackend.EMULATION:
284
+ x_2d = x.reshape(-1, x.shape[-1])
285
out = run_nvfp4_emulations(
276
- x=x,
286
+ x=x_2d,
287
input_global_scale=input_global_scale_inv,
288
weight=weight,
289
weight_scale_swizzled=weight_scale,
290
weight_global_scale=weight_global_scale,
291
swizzle=swizzle,
292
)
293
+ out = out[:, :output_size]
294
if bias is not None:
295
out = out + bias
285
- return out
286
-
287
- output_dtype = x.dtype
288
- output_shape = [*x.shape[:-1], output_size]
296
+ return out.view(*output_shape)
297
298
# Quantize BF16 or FP16 to (FP4 and interleaved block scale)
299
x_fp4, x_blockscale = scaled_fp4_quant(