@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 151 lines 5.42 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4 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
12
13 @pytest.mark.parametrize(
14 "model",
15 ["Qwen/Qwen3-Embedding-0.6B"],
16 )
17 @torch.inference_mode
18 def test_decoder_token_embed_model_runner_v2(
19 hf_runner, vllm_runner, monkeypatch, model: str
20 ):
21 monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
22 chunk_size = 10
23 n_prompt_tokens = [55, 56, 57]
24 token_prompts = [[1024 + i for i in range(n)] for n in n_prompt_tokens]
25
26 with vllm_runner(
27 model,
28 runner="pooling",
29 pooler_config=PoolerConfig(task="token_embed"),
30 max_model_len=128,
31 max_num_batched_tokens=chunk_size,
32 enforce_eager=True,
33 # `enable_chunked_prefill`: Set to `False` instead of `None` in VllmRunner
34 enable_chunked_prefill=True,
35 enable_prefix_caching=True,
36 ) as vllm_model:
37 assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
38 vllm_outputs = vllm_model.token_embed(
39 [TokensPrompt(prompt_token_ids=t) for t in token_prompts],
40 )
41
42 with hf_runner(
43 model,
44 auto_cls=AutoModel,
45 ) as hf_model:
46 hf_outputs = []
47 for token_prompt in token_prompts:
48 inputs = hf_model.wrap_device({"input_ids": torch.tensor([token_prompt])})
49 input_ids = inputs["input_ids"]
50 output = hf_model.model(input_ids)
51 hf_outputs.append(output.last_hidden_state.cpu().float()[0])
52
53 for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
54 check_embeddings_close(
55 embeddings_0_lst=hf_output,
56 embeddings_1_lst=vllm_output,
57 name_0="hf",
58 name_1="vllm",
59 tol=1e-2,
60 )
61
62
63 _RERANKER_HF_OVERRIDES = {
64 "architectures": ["Qwen3ForSequenceClassification"],
65 "classifier_from_token": ["no", "yes"],
66 "is_original_qwen3_reranker": True,
67 }
68 _RERANKER_TEMPLATE = VLLM_PATH / "examples/pooling/score/template/qwen3_reranker.jinja"
69
70
71 @pytest.mark.parametrize("model", ["Qwen/Qwen3-Reranker-0.6B"])
72 @torch.inference_mode
73 def test_last_pool_score_chunked_prefill_matches_unchunked(vllm_runner, model: str):
74 """LAST-pooling score must not depend on whether the prompt is chunked.
75
76 Regression test for a wrong-score bug where, under ``torch.compile`` (i.e.
77 not ``enforce_eager``), a query+document pair long enough to be split across
78 prefill chunks produced a corrupted last-token hidden state and thus a
79 wrong relevance score, while the same input scored correctly unchunked.
80 The existing pooling+chunked-prefill coverage runs ``enforce_eager=True``
81 and so never exercised the compiled path.
82 """
83 chat_template = _RERANKER_TEMPLATE.read_text()
84 query = "What organelle produces energy in the cell?"
85 # A long, matching document so query + doc exceeds the chunk size below.
86 document = (
87 "The mitochondria is the powerhouse of the cell. It generates most of "
88 "the cell chemical energy through oxidative phosphorylation. "
89 ) * 400
90
91 def score_with(max_num_batched_tokens: int) -> float:
92 with vllm_runner(
93 model,
94 runner="pooling",
95 hf_overrides=_RERANKER_HF_OVERRIDES,
96 max_model_len=16384,
97 max_num_batched_tokens=max_num_batched_tokens,
98 enable_chunked_prefill=True,
99 enable_prefix_caching=False,
100 ) as vllm_model:
101 return vllm_model.score(query, document, chat_template=chat_template)[0]
102
103 # Chunk size forces the prompt across several prefill chunks; the large
104 # budget keeps it in a single chunk as the reference.
105 chunked = score_with(2048)
106 unchunked = score_with(16384)
107
108 assert chunked == pytest.approx(unchunked, abs=5e-2), (
109 f"chunked score {chunked} diverged from unchunked {unchunked}"
110 )
111
112
113 @torch.inference_mode
114 def test_sequence_embed_model_runner_v2(hf_runner, vllm_runner, monkeypatch) -> None:
115 model = "Qwen/Qwen3-Embedding-0.6B"
116 chunk_size = 10
117 token_prompts = [[1024 + i for i in range(n)] for n in (25, 27)]
118 prompts = [TokensPrompt(prompt_token_ids=t) for t in token_prompts]
119
120 with hf_runner(model, auto_cls=AutoModel) as hf_model:
121 hf_outputs = []
122 for token_prompt in token_prompts:
123 inputs = hf_model.wrap_device({"input_ids": torch.tensor([token_prompt])})
124 output = hf_model.model(inputs["input_ids"])
125 embedding = torch.nn.functional.normalize(
126 output.last_hidden_state.float()[0, -1], dim=0
127 )
128 hf_outputs.append(embedding.cpu().tolist())
129
130 monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
131 with vllm_runner(
132 model,
133 runner="pooling",
134 pooler_config=PoolerConfig(task="embed"),
135 max_model_len=64,
136 max_num_batched_tokens=chunk_size,
137 max_num_seqs=2,
138 gpu_memory_utilization=0.25,
139 enforce_eager=True,
140 enable_chunked_prefill=True,
141 ) as vllm_model:
142 assert vllm_model.llm.llm_engine.vllm_config.use_v2_model_runner
143 vllm_outputs = vllm_model.embed(prompts)
144
145 check_embeddings_close(
146 embeddings_0_lst=hf_outputs,
147 embeddings_1_lst=vllm_outputs,
148 name_0="hf",
149 name_1="vllm",
150 tol=1e-2,
151 )