@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 166 lines 5.33 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import pytest
4 import torch.nn as nn
5
6 from vllm.platforms import current_platform
7
8 prompt = """
9 Generals gathered in their masses
10 Just like witches at black masses
11 Evil minds that plot destruction
12 Sorcerer of death's construction
13 In the fields, the bodies burning
14 As the war machine keeps turning
15 Death and hatred to mankind
16 Poisoning their brainwashed minds
17 Oh, Lord, yeah
18
19 Politicians hide themselves away
20 They only started the war
21 Why should they go out to fight?
22 They leave that all to the poor, yeah
23 Time will tell on their power minds
24 Making war just for fun
25 Treating people just like pawns in chess
26 Wait till their judgment day comes, yeah
27
28 Now, in darkness, world stops turning
29 Ashes where their bodies burning
30 No more war pigs have the power
31 Hand of God has struck the hour
32 Day of Judgment, God is calling
33 On their knees, the war pigs crawling
34 Begging mercies for their sins
35 Satan, laughing, spreads his wings
36 Oh, Lord, yeah
37 """
38
39
40 class WrapperPooler(nn.Module):
41 def __init__(self, pooler):
42 super().__init__()
43 self.pooler = pooler
44 self.chunks = []
45
46 def get_pooling_updates(self, task):
47 return self.pooler.get_pooling_updates(task)
48
49 def forward(
50 self,
51 hidden_states,
52 pooling_metadata,
53 ):
54 self.chunks.append(hidden_states.shape[0])
55 return self.pooler(hidden_states, pooling_metadata)
56
57
58 def inject_pooler(self):
59 model = self.get_model()
60 wrapper = WrapperPooler(model.pooler)
61 model.pooler = wrapper
62
63
64 def retrieve_chunks(self):
65 model = self.get_model()
66 chunks = model.pooler.chunks
67 model.pooler.chunks = []
68 return chunks
69
70
71 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
72 def test_pooling_chunked_prefill(vllm_runner, monkeypatch):
73 """Test chunked prefill for pooling models with LastPool."""
74 with monkeypatch.context() as m:
75 m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
76 model_id = "Qwen/Qwen3-Embedding-0.6B"
77
78 chunk_size = 10
79
80 # Cap the token budget to force chunked prefill; the
81 # long_prefill_token_threshold is not applied to a lone request.
82 with vllm_runner(
83 model_id,
84 runner="pooling",
85 max_num_batched_tokens=chunk_size,
86 max_num_seqs=1,
87 tensor_parallel_size=1,
88 enforce_eager=True,
89 enable_chunked_prefill=True,
90 ) as llm:
91 llm.get_llm().llm_engine.collective_rpc(inject_pooler)
92
93 tokenizer = llm.get_llm().get_tokenizer()
94 tokens = tokenizer(prompt)["input_ids"]
95 prompt_len = len(tokens)
96 full_chunks, last_chunk = divmod(prompt_len, chunk_size)
97 expected_chunks = [chunk_size] * full_chunks
98 if last_chunk:
99 expected_chunks.append(last_chunk)
100 llm.embed([prompt])
101 chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
102
103 # Check that PoolerWrapper was called and chunks were received
104 assert len(chunks) > 1
105 assert chunks == expected_chunks
106
107 # Disable chunked prefill
108 with vllm_runner(
109 model_id,
110 runner="pooling",
111 tensor_parallel_size=1,
112 enforce_eager=True,
113 ) as llm:
114 llm.get_llm().llm_engine.collective_rpc(inject_pooler)
115 llm.embed([prompt])
116 chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
117
118 # Check that PoolerWrapper was called and no chunks were received
119 assert len(chunks) == 1
120 assert chunks[0] == prompt_len
121
122
123 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
124 def test_pooling_prefix_cache(vllm_runner, monkeypatch):
125 """Test chunked prefill for pooling models with LastPool."""
126 verses = prompt.split("\n\n")
127
128 with monkeypatch.context() as m:
129 m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
130 model_id = "Qwen/Qwen3-Embedding-0.6B"
131
132 with vllm_runner(
133 model_id,
134 runner="pooling",
135 enable_prefix_caching=True,
136 tensor_parallel_size=1,
137 enforce_eager=True,
138 ) as llm:
139 llm.get_llm().llm_engine.collective_rpc(inject_pooler)
140 tokenizer = llm.get_llm().get_tokenizer()
141
142 prompt1 = "\n\n".join([verses[0], verses[1]])
143 prompt2 = "\n\n".join([verses[0], verses[2]])
144 tokens1 = tokenizer(prompt1)["input_ids"]
145 tokens2 = tokenizer(prompt2)["input_ids"]
146 prompt1_len = len(tokens1)
147 prompt2_len = len(tokens2)
148
149 llm.embed([prompt1])
150 chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
151
152 assert len(chunks) == 1
153 assert chunks[0] == prompt1_len
154
155 llm.embed([prompt2])
156 chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
157
158 assert len(chunks) == 1
159 assert chunks[0] <= prompt1_len
160 assert chunks[0] < prompt2_len
161
162 vllm_config = llm.get_llm().llm_engine.vllm_config
163 cache_config = vllm_config.cache_config
164 print(f"{cache_config=}")
165 # Prefixes are cached in blocks
166 assert (prompt2_len - chunks[0]) % cache_config.block_size == 0