[Bugfix][CI] Fix the flaky sharded-sampling tests, and the engine teardown need (#58342)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com> Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Andreas Karatzas <Andreas.Karatzas@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
stefankoncarevic committed
Sep 24, 2026 at 11:05 UTC
c0c35ddd8e9c6dadd53a8f92e20e9a03a210da15
3 files changed
+132
-37
tests/v1/e2e/general/test_sharded_sampling.py
+115
-27
@@ -2,26 +2,30 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Batch-sharded sampling must be a drop-in for replicated sampling.
4
5
-The comparison is tolerance-based, not bitwise: engine boots land in
6
-slightly different kernel/collective states (measured: same-state boots are
7
-bitwise identical across the two modes; cross-state boots differ by up to
8
-~0.25 logprob and can flip a near-tie token). A real sharding bug — rows
9
-routed to the wrong request, misaligned gathers — produces many-nats
10
-divergence and scrambled top-k sets, far outside these bounds."""
5
+The two boots are compared token for token, which is only meaningful because
6
+both schedule the same batches. With the engine in the caller's process every
7
+request is queued before the first step, so the first prefill has the same
8
+shape in both runs. Give the engine its own process and the comparison turns
9
+into a coin toss: prefill starts before the last requests arrive over IPC, so
10
+one boot batches two requests where the other batches one. The batch shape
11
+sets the GEMM reduction order, logprobs shift by up to ~0.25, and near-ties
12
+flip — measured at up to 5 of 16 completions, with sharding disabled on both
13
+sides. A real sharding bug — rows routed to the wrong request, misaligned
14
+gathers — produces many-nats divergence and scrambled top-k sets."""
15
16
import pytest
17
18
from tests.utils import multi_gpu_test
19
from vllm import LLM, SamplingParams
20
from vllm.distributed import cleanup_dist_env_and_memory
21
+from vllm.logprobs import Logprob, LogprobsOnePosition
22
23
MODEL = "meta-llama/Llama-3.2-1B-Instruct"
24
20
-# Boot-state numeric noise bound (measured <= 0.25 on identical contexts).
25
+# Misrouted rows move logprobs by nats; anything under this is float noise.
26
+# With matched batches the two boots agreed bitwise over 320 positions, so the
27
+# bound is not expected to be reached.
28
LOGPROB_TOL = 0.5
22
-# Near-tie greedy/seeded flips from boot-state noise; measured <= 1 per
23
-# engine boot across ~16 completions, plus one of headroom.
24
-MAX_DIVERGENT_PROMPTS = 3
29
30
PROMPTS = [
31
"The capital of France is",
@@ -54,6 +58,10 @@ def _is_sharded_sampling_active(worker) -> bool:
58
return worker.model_runner.batch_sharder is not None
59
60
61
+def _sampling_params() -> list[SamplingParams]:
62
+ return [SAMPLING_PARAMS[i % len(SAMPLING_PARAMS)] for i in range(len(PROMPTS))]
63
+
64
+
65
def _generate(monkeypatch: pytest.MonkeyPatch, disable_sharding: bool):
66
llm = LLM(
67
model=MODEL,
@@ -75,9 +83,7 @@ def _generate(monkeypatch: pytest.MonkeyPatch, disable_sharding: bool):
83
# intended sampling mode.
84
modes = llm.llm_engine.collective_rpc(_is_sharded_sampling_active)
85
assert all(mode == (not disable_sharding) for mode in modes), modes
78
- params = [
79
- SAMPLING_PARAMS[i % len(SAMPLING_PARAMS)] for i in range(len(PROMPTS))
80
- ]
86
+ params = _sampling_params()
87
# Two waves: the second reuses the request slots freed by the first
88
# (in whatever order requests finished), covering the cross-rank
89
# slot-recycling path that request ownership derives from.
@@ -89,13 +95,36 @@ def _generate(monkeypatch: pytest.MonkeyPatch, disable_sharding: bool):
95
cleanup_dist_env_and_memory()
96
97
92
-def _assert_logprob_dicts_close(ref_lps, out_lps, what: str) -> None:
93
- # Allow one boundary entry of the top-k set to swap at a near-tie.
94
- common = set(ref_lps) & set(out_lps)
95
- assert len(common) >= max(len(ref_lps), len(out_lps)) - 1, (
96
- f"{what}: top-k token sets diverge: {sorted(ref_lps)} vs {sorted(out_lps)}"
98
+def _requested_top_k(lps: LogprobsOnePosition, num_logprobs: int) -> set[int]:
99
+ """Exclude the extra sampled/prompt token, including at rank k+1."""
100
+ top = {
101
+ token_id
102
+ for token_id, logprob in lps.items()
103
+ if logprob.rank is not None and 1 <= logprob.rank <= num_logprobs
104
+ }
105
+ assert len(top) == num_logprobs, (
106
+ f"expected {num_logprobs} ranked top-k entries, got {len(top)}"
107
+ )
108
+ return top
109
+
110
+
111
+def _assert_logprob_dicts_close(
112
+ ref_lps: LogprobsOnePosition,
113
+ out_lps: LogprobsOnePosition,
114
+ what: str,
115
+ num_logprobs: int,
116
+) -> None:
117
+ # Allow one boundary entry of the top-k set to swap at a near-tie. Compare
118
+ # the top-k alone: counting the sampled token would spend that allowance
119
+ # on the extra entry itself whenever only one side carries one.
120
+ ref_top = _requested_top_k(ref_lps, num_logprobs)
121
+ out_top = _requested_top_k(out_lps, num_logprobs)
122
+ shared_top = ref_top & out_top
123
+ assert len(shared_top) >= num_logprobs - 1, (
124
+ f"{what}: top-k token sets diverge: {sorted(ref_top)} vs {sorted(out_top)}"
125
)
98
- for token_id in common:
126
+ # Values are comparable for every shared token, sampled one included.
127
+ for token_id in set(ref_lps) & set(out_lps):
128
diff = abs(ref_lps[token_id].logprob - out_lps[token_id].logprob)
129
assert diff <= LOGPROB_TOL, (
130
f"{what}[{token_id}]: logprob diff {diff} "
@@ -103,20 +132,72 @@ def _assert_logprob_dicts_close(ref_lps, out_lps, what: str) -> None:
132
)
133
134
135
+@pytest.mark.cpu_test
136
+@pytest.mark.parametrize(
137
+ ("num_logprobs", "ref_ranks", "out_ranks"),
138
+ [
139
+ (2, {0: 1, 1: 2, 2: 3}, {0: 1, 3: 2, 2: 4}),
140
+ (1, {0: 1, 1: 2}, {2: 1, 1: 3}),
141
+ (2, {0: 1, 1: 2}, {0: 1, 2: 2, 1: 3}),
142
+ (2, {0: 1, 1: 2, 2: 4}, {0: 1, 1: 2, 2: 5}),
143
+ (0, {0: 2}, {0: 3}),
144
+ ],
145
+ ids=["decode-boundary", "prompt-boundary", "sampled-in-top-k", "gap", "zero"],
146
+)
147
+def test_logprob_comparison_allows_sampled_token_rank_changes(
148
+ num_logprobs: int, ref_ranks: dict[int, int], out_ranks: dict[int, int]
149
+):
150
+ ref_lps, out_lps = (
151
+ {token: Logprob(-2.0 - rank / 100, rank=rank) for token, rank in ranks.items()}
152
+ for ranks in (ref_ranks, out_ranks)
153
+ )
154
+ _assert_logprob_dicts_close(ref_lps, out_lps, "boundary swap", num_logprobs)
155
+
156
+
157
+@pytest.mark.cpu_test
158
+def test_logprob_comparison_rejects_multiple_top_k_swaps():
159
+ ref = {0: Logprob(-2.0, rank=1), 1: Logprob(-2.01, rank=2)}
160
+ out = {2: Logprob(-2.0, rank=1), 3: Logprob(-2.01, rank=2)}
161
+ with pytest.raises(AssertionError, match="top-k token sets diverge"):
162
+ _assert_logprob_dicts_close(ref, out, "misrouted top-k", num_logprobs=2)
163
+
164
+
165
+@pytest.mark.cpu_test
166
+def test_logprob_comparison_checks_extra_sampled_token_value():
167
+ ref = {0: Logprob(-1.0, rank=1), 1: Logprob(-3.0, rank=3)}
168
+ out = {0: Logprob(-1.0, rank=1), 1: Logprob(-4.0, rank=3)}
169
+ with pytest.raises(AssertionError, match="logprob diff"):
170
+ _assert_logprob_dicts_close(ref, out, "sampled token", num_logprobs=1)
171
+
172
+
173
+@pytest.mark.cpu_test
174
+@pytest.mark.parametrize("rank", [None, 2])
175
+def test_logprob_comparison_rejects_missing_top_k_rank(rank: int | None):
176
+ lps = {0: Logprob(-1.0, rank=rank)}
177
+ with pytest.raises(AssertionError, match="expected 1 ranked top-k entries"):
178
+ _assert_logprob_dicts_close(lps, lps, "missing rank", num_logprobs=1)
179
+
180
+
181
@multi_gpu_test(num_gpus=2)
182
def test_sharded_sampling_outputs_match(monkeypatch: pytest.MonkeyPatch):
183
"""Generation, logprobs, and prompt logprobs match between batch-sharded
109
- sampling and the replicated fallback, up to measured boot-state noise."""
184
+ sampling and the replicated fallback."""
185
monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
186
# Required for the collective_rpc mode check below.
187
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
188
+ # Queue every request before the first step, so both boots schedule the
189
+ # same batches -- see the module docstring.
190
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
191
192
ref_outputs = _generate(monkeypatch, disable_sharding=True)
193
shard_outputs = _generate(monkeypatch, disable_sharding=False)
194
195
assert len(ref_outputs) == len(shard_outputs) == 2 * len(PROMPTS)
118
- num_divergent = 0
119
- for i, (ref, out) in enumerate(zip(ref_outputs, shard_outputs)):
196
+ params = _sampling_params()
197
+ divergent = []
198
+ for i, (ref, out, sampling_params) in enumerate(
199
+ zip(ref_outputs, shard_outputs, params + params[::-1])
200
+ ):
201
ref_completion = ref.outputs[0]
202
out_completion = out.outputs[0]
203
ref_ids = list(ref_completion.token_ids)
@@ -131,21 +212,24 @@ def test_sharded_sampling_outputs_match(monkeypatch: pytest.MonkeyPatch):
212
break
213
prefix += 1
214
if ref_ids != out_ids:
134
- num_divergent += 1
215
+ divergent.append((i, prefix))
216
217
if ref_completion.logprobs is not None:
218
assert out_completion.logprobs is not None
219
+ assert sampling_params.logprobs is not None
220
comparable = min(prefix + 1, len(ref_ids), len(out_ids))
221
for pos in range(comparable):
222
_assert_logprob_dicts_close(
223
ref_completion.logprobs[pos],
224
out_completion.logprobs[pos],
225
f"prompt {i} logprobs[{pos}]",
226
+ sampling_params.logprobs,
227
)
228
229
# The prompt is fixed, so every prompt-logprob position is comparable.
230
if ref.prompt_logprobs is not None:
231
assert out.prompt_logprobs is not None
232
+ assert sampling_params.prompt_logprobs is not None
233
for pos, (ref_lps, out_lps) in enumerate(
234
zip(ref.prompt_logprobs, out.prompt_logprobs)
235
):
@@ -153,10 +237,14 @@ def test_sharded_sampling_outputs_match(monkeypatch: pytest.MonkeyPatch):
237
assert ref_lps is None and out_lps is None
238
continue
239
_assert_logprob_dicts_close(
156
- ref_lps, out_lps, f"prompt {i} prompt_logprobs[{pos}]"
240
+ ref_lps,
241
+ out_lps,
242
+ f"prompt {i} prompt_logprobs[{pos}]",
243
+ sampling_params.prompt_logprobs,
244
)
245
159
- assert num_divergent <= MAX_DIVERGENT_PROMPTS, (
160
- f"{num_divergent}/{2 * len(PROMPTS)} prompts diverged: beyond near-tie "
161
- "boot-state noise; sharded sampling is likely misrouting requests"
246
+ assert not divergent, (
247
+ f"{len(divergent)}/{2 * len(PROMPTS)} completions diverged at "
248
+ f"(prompt, position) {divergent}; the sampling modes did not agree "
249
+ "token for token with in-process scheduling"
250
)
tests/v1/e2e/spec_decode/test_sharded_sampling.py
+10
-9
@@ -39,6 +39,9 @@ def test_mtp_sharded_sampling_equivalence(monkeypatch: pytest.MonkeyPatch):
39
m.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
40
# Required for the collective_rpc mode check below.
41
m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
42
+ # Queue every request before the first step, so both boots schedule
43
+ # the same batches -- see tests/v1/e2e/general/test_sharded_sampling.py.
44
+ m.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
45
46
def run(disable_sharding: bool):
47
llm = LLM(
@@ -79,24 +82,22 @@ def test_mtp_sharded_sampling_equivalence(monkeypatch: pytest.MonkeyPatch):
82
ref_greedy, ref_seeded = run(disable_sharding=True)
83
shard_greedy, shard_seeded = run(disable_sharding=False)
84
82
- # Engine boots are not bitwise deterministic (kernel/collective
83
- # selection shifts logits by ~1 ulp, flipping near-tie tokens), so
84
- # allow a small number of divergent prompts. A sharding bug produces
85
- # wholesale divergence, not isolated near-tie flips.
85
+ # A sharding bug shows up as wholesale divergence, not isolated flips.
86
for name, ref_outputs, shard_outputs in (
87
("greedy", ref_greedy, shard_greedy),
88
("seeded", ref_seeded, shard_seeded),
89
):
90
- num_divergent = 0
90
+ divergent = []
91
for i, (ref, out) in enumerate(zip(ref_outputs, shard_outputs)):
92
if list(ref.outputs[0].token_ids) != list(out.outputs[0].token_ids):
93
- num_divergent += 1
93
+ divergent.append(i)
94
print(
95
f"{name} prompt {i} diverged:\n"
96
f" replicated: {ref.outputs[0].text!r}\n"
97
f" sharded: {out.outputs[0].text!r}"
98
)
99
- assert num_divergent <= 2, (
100
- f"{name}: {num_divergent}/{len(ref_outputs)} prompts diverged, "
101
- "beyond near-tie boot noise"
99
+ assert not divergent, (
100
+ f"{name}: {len(divergent)}/{len(ref_outputs)} prompts diverged "
101
+ f"at {divergent}; the sampling modes did not agree token for "
102
+ "token with in-process scheduling"
103
)
vllm/v1/engine/core_client.py
+7
-1
@@ -366,6 +366,11 @@ class InprocClient(EngineCoreClient):
366
log_stats,
367
executor_fail_callback=executor_fail_callback,
368
)
369
+ # Release the executor's workers when the client is collected, as
370
+ # MPClient does; they hold GPU memory for the life of the process
371
+ # otherwise. The callback must not reference self, or self would stay
372
+ # reachable and the finalizer would never run.
373
+ self._finalizer = weakref.finalize(self, self.engine_core.shutdown)
374
375
def get_output(self) -> EngineCoreOutputs:
376
outputs, model_executed = self.engine_core.step_fn()
@@ -384,7 +389,8 @@ class InprocClient(EngineCoreClient):
389
self.engine_core.abort_requests(request_ids)
390
391
def shutdown(self, timeout: float | None = None) -> None:
387
- self.engine_core.shutdown()
392
+ if self._finalizer.detach() is not None:
393
+ self.engine_core.shutdown()
394
395
def profile(self, is_start: bool = True, profile_prefix: str | None = None) -> None:
396
self.engine_core.profile(is_start, profile_prefix)