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, 03:30 UTC
[Bug] Fix batch invariant test issue, bs=1 with `max_seq_num = 1` (#39320)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Apr 15, 2026 at 16:28 UTC
39ac640490ee2e8f951d343ae1707dd9bdacaf70
1 file changed
+11
-25
tests/v1/determinism/test_batch_invariance.py
+11
-25
@@ -36,10 +36,10 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
36
using the high-level v1 LLM() API only (no manual batching).
37
38
Strategy:
39
- - Create two LLM engines with identical config except max_num_seqs: 1 vs N.
40
- - Compute a baseline output for the needle prompt with the bs=1 engine.
41
- - For many trials, generate a batch (size N) where the needle appears at a
42
- random position among random filler prompts using the bs=N engine.
39
+ - Create a single LLM engine configured for the larger batch limit (N).
40
+ - Compute a baseline output for the needle prompt when it is run alone.
41
+ - For many trials, generate a mixed batch (size N) where the needle appears
42
+ at a random position among random filler prompts using the same engine.
43
- Track how many trials match vs mismatch, and report totals at the end.
44
The test fails if any mismatches occur, but we still dump pass/fail
45
counts.
@@ -83,11 +83,9 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
83
84
needle_prompt = "There once was a "
85
86
- llm_bs1 = None
87
- llm_bsN = None
86
+ llm = None
87
try:
89
- # Engine with bs=1 behavior
90
- llm_bs1 = LLM_with_max_seqs(
88
+ llm = LLM_with_max_seqs(
89
model=model,
90
max_num_seqs=max_batch_size,
91
gpu_memory_utilization=gpu_mem_util,
@@ -96,20 +94,11 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
94
)
95
96
# Baseline generation for the needle prompt alone.
99
- baseline_out = llm_bs1.generate([needle_prompt], sampling)
97
+ baseline_out = llm.generate([needle_prompt], sampling)
98
assert len(baseline_out) == 1
99
assert len(baseline_out[0].outputs) >= 1
100
baseline_text = baseline_out[0].outputs[0].text
101
104
- # Engine with larger batch limit (e.g., 64)
105
- llm_bsN = LLM_with_max_seqs(
106
- model=model,
107
- max_num_seqs=max_batch_size,
108
- gpu_memory_utilization=gpu_mem_util,
109
- max_model_len=max_model_len,
110
- attention_config=attention_config,
111
- )
112
-
102
mismatches = 0
103
104
for trial in range(num_trials):
@@ -124,8 +113,8 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
113
else:
114
prompts.append(_random_prompt(min_random_prompt, max_random_prompt))
115
127
- # Generate with the larger-batch engine
128
- outputs = llm_bsN.generate(prompts, sampling)
116
+ # Generate with the same engine but in a larger batch.
117
+ outputs = llm.generate(prompts, sampling)
118
# Find the needle output by position
119
needle_output = outputs[needle_pos]
120
assert needle_output.prompt == needle_prompt
@@ -151,12 +140,9 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
140
141
finally:
142
# Ensure engines are shutdown to free GPU/VRAM across test sessions
154
- if llm_bs1 is not None:
155
- with contextlib.suppress(Exception):
156
- llm_bs1.shutdown()
157
- if llm_bsN is not None:
143
+ if llm is not None:
144
with contextlib.suppress(Exception):
159
- llm_bsN.shutdown()
145
+ llm.shutdown()
146
147
148
@skip_unsupported