[Bugfix][Pooling] Fix wrong scores for chunked prefill under torch.compile (#48901)
Signed-off-by: seewoo <seewoo@ucsc.edu> Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>
Sebastian Woo committed
Jul 16, 2026 at 22:19 UTC
3b6c96a101651c78737b23b2bd8f1026c02f22e2
2 files changed
+52
tests/models/language/pooling/test_all_pooling_plus_chunked_prefill.py
+51
@@ -5,6 +5,7 @@ import torch
5
from transformers import AutoModel
6
7
from tests.models.utils import check_embeddings_close
8
+from tests.utils import VLLM_PATH
9
from vllm import TokensPrompt
10
from vllm.config import PoolerConfig
11
@@ -53,3 +54,53 @@ def test_embed_models(hf_runner, vllm_runner, model: str):
54
name_1="vllm",
55
tol=1e-2,
56
)
57
+
58
+
59
+_RERANKER_HF_OVERRIDES = {
60
+ "architectures": ["Qwen3ForSequenceClassification"],
61
+ "classifier_from_token": ["no", "yes"],
62
+ "is_original_qwen3_reranker": True,
63
+}
64
+_RERANKER_TEMPLATE = VLLM_PATH / "examples/pooling/score/template/qwen3_reranker.jinja"
65
+
66
+
67
+@pytest.mark.parametrize("model", ["Qwen/Qwen3-Reranker-0.6B"])
68
+@torch.inference_mode
69
+def test_last_pool_score_chunked_prefill_matches_unchunked(vllm_runner, model: str):
70
+ """LAST-pooling score must not depend on whether the prompt is chunked.
71
+
72
+ Regression test for a wrong-score bug where, under ``torch.compile`` (i.e.
73
+ not ``enforce_eager``), a query+document pair long enough to be split across
74
+ prefill chunks produced a corrupted last-token hidden state and thus a
75
+ wrong relevance score, while the same input scored correctly unchunked.
76
+ The existing pooling+chunked-prefill coverage runs ``enforce_eager=True``
77
+ and so never exercised the compiled path.
78
+ """
79
+ chat_template = _RERANKER_TEMPLATE.read_text()
80
+ query = "What organelle produces energy in the cell?"
81
+ # A long, matching document so query + doc exceeds the chunk size below.
82
+ document = (
83
+ "The mitochondria is the powerhouse of the cell. It generates most of "
84
+ "the cell chemical energy through oxidative phosphorylation. "
85
+ ) * 400
86
+
87
+ def score_with(max_num_batched_tokens: int) -> float:
88
+ with vllm_runner(
89
+ model,
90
+ runner="pooling",
91
+ hf_overrides=_RERANKER_HF_OVERRIDES,
92
+ max_model_len=16384,
93
+ max_num_batched_tokens=max_num_batched_tokens,
94
+ enable_chunked_prefill=True,
95
+ enable_prefix_caching=False,
96
+ ) as vllm_model:
97
+ return vllm_model.score(query, document, chat_template=chat_template)[0]
98
+
99
+ # Chunk size forces the prompt across several prefill chunks; the large
100
+ # budget keeps it in a single chunk as the reference.
101
+ chunked = score_with(2048)
102
+ unchunked = score_with(16384)
103
+
104
+ assert chunked == pytest.approx(unchunked, abs=5e-2), (
105
+ f"chunked score {chunked} diverged from unchunked {unchunked}"
106
+ )
vllm/v1/worker/gpu_model_runner.py
+1
@@ -3437,6 +3437,7 @@ class GPUModelRunner(
3437
)
3438
3439
if raw_pooler_output is None or not any(finished_mask):
3440
+ self._sync_device()
3441
model_runner_output.pooler_output = [None] * num_reqs
3442
return model_runner_output
3443