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 14, 05:30 UTC
Add tests for chunked prefill and prefix cache with causal pooling models (#26526)
Signed-off-by: Max de Bayser <mbayser@br.ibm.com> Co-authored-by: Ayush Singh <ayush1009208@gmail.com>
Maximilien de Bayser committed
Oct 13, 2025 at 20:45 UTC
d8bebb008a8bac5a0095288815de5cc79753ed38
1 file changed
+167
tests/v1/e2e/test_pooling_chunked_prefill.py
new
+167
@@ -0,0 +1,167 @@
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
+
75
+ with monkeypatch.context() as m:
76
+ m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
77
+ model_id = "Qwen/Qwen3-Embedding-0.6B"
78
+
79
+ chunk_size = 10
80
+
81
+ # Set chunking parameters to force chunked prefill
82
+ # Note: Chunked prefill is automatically handled by vLLM
83
+ # internally based on the model size and prompt
84
+ with vllm_runner(
85
+ model_id,
86
+ runner="pooling",
87
+ long_prefill_token_threshold=chunk_size,
88
+ tensor_parallel_size=1,
89
+ enforce_eager=True,
90
+ enable_chunked_prefill=True,
91
+ ) as llm:
92
+ llm.get_llm().llm_engine.collective_rpc(inject_pooler)
93
+
94
+ tokenizer = llm.get_llm().get_tokenizer()
95
+ tokens = tokenizer(prompt)["input_ids"]
96
+ prompt_len = len(tokens)
97
+ full_chunks, last_chunk = divmod(prompt_len, chunk_size)
98
+ expected_chunks = [chunk_size] * full_chunks
99
+ if last_chunk:
100
+ expected_chunks.append(last_chunk)
101
+ llm.embed([prompt])
102
+ chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
103
+
104
+ # Check that PoolerWrapper was called and chunks were received
105
+ assert len(chunks) > 1
106
+ assert chunks == expected_chunks
107
+
108
+ # Disable chunked prefill
109
+ with vllm_runner(
110
+ model_id,
111
+ runner="pooling",
112
+ tensor_parallel_size=1,
113
+ enforce_eager=True,
114
+ ) as llm:
115
+ llm.get_llm().llm_engine.collective_rpc(inject_pooler)
116
+ llm.embed([prompt])
117
+ chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
118
+
119
+ # Check that PoolerWrapper was called and no chunks were received
120
+ assert len(chunks) == 1
121
+ assert chunks[0] == prompt_len
122
+
123
+
124
+@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available")
125
+def test_pooling_prefix_cache(vllm_runner, monkeypatch):
126
+ """Test chunked prefill for pooling models with LastPool."""
127
+
128
+ verses = prompt.split("\n\n")
129
+
130
+ with monkeypatch.context() as m:
131
+ m.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
132
+ model_id = "Qwen/Qwen3-Embedding-0.6B"
133
+
134
+ with vllm_runner(
135
+ model_id,
136
+ runner="pooling",
137
+ enable_prefix_caching=True,
138
+ tensor_parallel_size=1,
139
+ enforce_eager=True,
140
+ ) as llm:
141
+ llm.get_llm().llm_engine.collective_rpc(inject_pooler)
142
+ tokenizer = llm.get_llm().get_tokenizer()
143
+
144
+ prompt1 = "\n\n".join([verses[0], verses[1]])
145
+ prompt2 = "\n\n".join([verses[0], verses[2]])
146
+ tokens1 = tokenizer(prompt1)["input_ids"]
147
+ tokens2 = tokenizer(prompt2)["input_ids"]
148
+ prompt1_len = len(tokens1)
149
+ prompt2_len = len(tokens2)
150
+
151
+ llm.embed([prompt1])
152
+ chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
153
+
154
+ assert len(chunks) == 1
155
+ assert chunks[0] == prompt1_len
156
+
157
+ llm.embed([prompt2])
158
+ chunks = llm.get_llm().llm_engine.collective_rpc(retrieve_chunks)[0]
159
+
160
+ assert len(chunks) == 1
161
+ assert chunks[0] <= prompt1_len
162
+ assert chunks[0] < prompt2_len
163
+
164
+ cache_config = llm.get_llm().llm_engine.cache_config
165
+ print(f"{cache_config=}")
166
+ # Prefixes are cached in blocks
167
+ assert (prompt2_len - chunks[0]) % cache_config.block_size == 0