[Test] Add focused hybrid MTP prefix-cache regressions (#53189)
Signed-off-by: mgoin <mgoin64@gmail.com>
Misha Goin committed
Aug 21, 2026 at 20:05 UTC
a34f2ab70a94d865e0ec01e200c73bcbff472e0b
2 files changed
+128
-22
.buildkite/test_areas/engine.yaml
+6
-5
@@ -133,14 +133,15 @@ steps:
133
depends_on:
134
- image-build-amd
135
136
-- label: ":nvidia: (H100) V1 E2E Hybrid Chunked Prefill"
137
- key: v1-e2e-4xh100
136
+- label: ":nvidia: (B200) V1 E2E Hybrid Chunked Prefill"
137
+ key: v1-e2e-hybrid-chunked-prefill-b200
138
timeout_in_minutes: 35
139
- device: h100
140
- num_devices: 4
139
+ device: b200-k8s
140
+ num_devices: 1
141
optional: true
142
source_file_dependencies:
143
- vllm/v1/attention/backends/utils.py
144
+ - vllm/v1/core/
145
- vllm/v1/worker/gpu_model_runner.py
146
- tests/v1/e2e/test_hybrid_chunked_prefill.py
147
commands:
@@ -149,7 +150,7 @@ steps:
150
amd:
151
label: ":amd: (MI300) V1 E2E Hybrid Chunked Prefill"
152
dind: false
152
- device: mi300_4
153
+ device: mi300_1
154
timeout_in_minutes: 35
155
depends_on:
156
- image-build-amd
tests/v1/e2e/test_hybrid_chunked_prefill.py
+122
-17
@@ -5,8 +5,30 @@ import pytest
5
6
from vllm import SamplingParams
7
from vllm.platforms import current_platform
8
+from vllm.v1.kv_cache_interface import MambaSpec
9
9
-from ...utils import large_gpu_mark, multi_gpu_marks
10
+from ...utils import large_gpu_mark
11
+
12
+QWEN_MODEL = "Qwen/Qwen3.5-4B"
13
+QWEN_KV_CACHE_BYTES = 12 << 30
14
+HYBRID_MTP_MODELS = [
15
+ pytest.param(
16
+ QWEN_MODEL,
17
+ marks=[large_gpu_mark(min_gb=30)],
18
+ id="qwen",
19
+ ),
20
+ pytest.param(
21
+ "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4",
22
+ marks=[large_gpu_mark(min_gb=80)]
23
+ + [
24
+ pytest.mark.skipif(
25
+ not current_platform.is_cuda(),
26
+ reason="modelopt quantization is supported only on CUDA",
27
+ )
28
+ ],
29
+ id="nemotron",
30
+ ),
31
+]
32
33
# A trivial request with a short prompt to ensure we run a mixed batch
34
SMALL_MESSAGE = [
@@ -38,20 +60,7 @@ MESSAGES = [
60
61
@pytest.mark.parametrize(
62
"model_name",
41
- [
42
- pytest.param("Qwen/Qwen3.5-4B", marks=[large_gpu_mark(min_gb=40)]),
43
- pytest.param(
44
- "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8",
45
- marks=[large_gpu_mark(min_gb=80)]
46
- + multi_gpu_marks(num_gpus=4)
47
- + [
48
- pytest.mark.skipif(
49
- not current_platform.is_cuda(),
50
- reason="modelopt quantization is supported only on CUDA",
51
- )
52
- ],
53
- ),
54
- ],
63
+ HYBRID_MTP_MODELS,
64
)
65
@pytest.mark.parametrize("enable_prefix_caching", [False, True])
66
def test_mtp_speculative_mixed_batch_short_prefill(
@@ -62,7 +71,7 @@ def test_mtp_speculative_mixed_batch_short_prefill(
71
72
# Set so large that both prefills will be classified as decodes in a mixed batch
73
# note, with prefix caching we require chunk_size >= mamba_block_size
65
- chunk_size = 256 if not enable_prefix_caching else 16384
74
+ chunk_size = 256 if not enable_prefix_caching else 2048
75
num_draft_tokens = 100
76
77
with vllm_runner(
@@ -72,9 +81,11 @@ def test_mtp_speculative_mixed_batch_short_prefill(
81
"num_speculative_tokens": num_draft_tokens,
82
},
83
max_num_batched_tokens=chunk_size,
84
+ max_num_seqs=4,
85
max_model_len=512,
86
+ kv_cache_memory_bytes=QWEN_KV_CACHE_BYTES if model_name == QWEN_MODEL else None,
87
enforce_eager=True,
77
- tensor_parallel_size=4,
88
+ tensor_parallel_size=1,
89
trust_remote_code=True,
90
enable_chunked_prefill=True,
91
enable_prefix_caching=enable_prefix_caching,
@@ -108,3 +119,97 @@ def test_mtp_speculative_mixed_batch_short_prefill(
119
assert "43" in responses[1], (
120
"The second response should contain the correct value of 42+1=43."
121
)
122
+
123
+
124
+def _get_mamba_block_size(llm) -> int:
125
+ scheduler = llm.llm_engine.engine_core.engine_core.scheduler
126
+ block_sizes = {
127
+ group.kv_cache_spec.block_size
128
+ for group in scheduler.kv_cache_config.kv_cache_groups
129
+ if isinstance(group.kv_cache_spec, MambaSpec)
130
+ }
131
+ assert len(block_sizes) == 1
132
+ block_size = block_sizes.pop()
133
+ assert scheduler.cache_config.block_size == block_size
134
+ return block_size
135
+
136
+
137
+def _build_access_code_manual(tokenizer, target_tokens: int) -> tuple[str, list[str]]:
138
+ codes = ["605341", "693278", "597596", "751982"]
139
+ header = "Memorize this facility manual and its access codes.\n\n"
140
+ facts = "".join(
141
+ f"The access code for vault-{i:02d} is {code}.\n"
142
+ for i, code in enumerate(codes)
143
+ )
144
+ filler = "Routine facility inspections are recorded in the audit ledger every day. "
145
+ manual = header + facts
146
+ while len(tokenizer.encode(manual)) < target_tokens:
147
+ manual = header + filler + manual.removeprefix(header)
148
+ return manual, codes
149
+
150
+
151
+@large_gpu_mark(min_gb=30)
152
+def test_qwen_mtp_mamba_prefix_cache_hit_is_bounded(vllm_runner, monkeypatch):
153
+ """MTP must not extend a hybrid Mamba hit past the attention hit."""
154
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
155
+
156
+ with vllm_runner(
157
+ QWEN_MODEL,
158
+ tensor_parallel_size=1,
159
+ max_model_len=8192,
160
+ kv_cache_memory_bytes=QWEN_KV_CACHE_BYTES,
161
+ block_size=None,
162
+ enforce_eager=True,
163
+ enable_chunked_prefill=True,
164
+ enable_prefix_caching=True,
165
+ mamba_cache_mode="align",
166
+ speculative_config={"method": "mtp", "num_speculative_tokens": 2},
167
+ ) as runner:
168
+ llm = runner.get_llm()
169
+ block_size = _get_mamba_block_size(llm)
170
+ manual, codes = _build_access_code_manual(
171
+ llm.get_tokenizer(), 2 * block_size - 192
172
+ )
173
+ manual_tokens = len(llm.get_tokenizer().encode(manual))
174
+ assert block_size < manual_tokens < 2 * block_size
175
+
176
+ wave1_prompts = [
177
+ manual
178
+ + f"\nDescribe the audit procedure for vault-{i:02d} in detail.\nAnswer:"
179
+ for i in range(len(codes))
180
+ ]
181
+ decode_tokens = 2 * block_size - manual_tokens + 96
182
+ wave1_params = SamplingParams(
183
+ temperature=0.0,
184
+ min_tokens=decode_tokens,
185
+ max_tokens=decode_tokens,
186
+ ignore_eos=True,
187
+ )
188
+ wave1_outputs = llm.generate(wave1_prompts, wave1_params)
189
+
190
+ wave2_prompts = [
191
+ prompt
192
+ + output.outputs[0].text
193
+ + f"\n\nWhat is the access code for vault-{i:02d}? Answer:"
194
+ for i, (prompt, output) in enumerate(zip(wave1_prompts, wave1_outputs))
195
+ ]
196
+ warm_params = SamplingParams(temperature=0.0, max_tokens=24, stop=["\n"])
197
+ cold_params = SamplingParams(
198
+ temperature=0.0,
199
+ max_tokens=24,
200
+ stop=["\n"],
201
+ skip_reading_prefix_cache=True,
202
+ )
203
+ warm_outputs = llm.generate(wave2_prompts, warm_params)
204
+ cold_outputs = llm.generate(wave2_prompts, cold_params)
205
+
206
+ cached_tokens = [output.num_cached_tokens for output in warm_outputs]
207
+ assert max(cached_tokens) == block_size
208
+
209
+ for arm, outputs in (("warm", warm_outputs), ("cold", cold_outputs)):
210
+ missed = [
211
+ i
212
+ for i, (code, output) in enumerate(zip(codes, outputs))
213
+ if code not in output.outputs[0].text
214
+ ]
215
+ assert not missed, f"{arm} cache missed access codes for prompts {missed}"