@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 1,231 lines 47.2 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import multiprocessing as mp
4 import os
5 import traceback
6 from collections.abc import Callable
7 from dataclasses import dataclass
8 from typing import Any
9
10 import datasets
11 import pytest
12 import torch
13
14 import vllm.envs as envs
15 from tests.utils import create_new_process_for_each_test
16 from vllm import LLM, SamplingParams, TokensPrompt
17 from vllm.config import CacheConfig
18 from vllm.distributed import cleanup_dist_env_and_memory
19 from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
20 from vllm.platforms import current_platform
21 from vllm.sequence import IntermediateTensors
22 from vllm.utils.gpu_sync_debug import gpu_sync_allowed
23 from vllm.utils.torch_utils import async_tensor_h2d
24 from vllm.v1.attention.backends.utils import CommonAttentionMetadata
25 from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
26 from vllm.v1.core.sched.output import SchedulerOutput
27 from vllm.v1.engine.core_client import InprocClient
28 from vllm.v1.kv_cache_interface import KVCacheConfig
29 from vllm.v1.outputs import SamplerOutput
30 from vllm.v1.request import Request
31 from vllm.v1.sample.metadata import SamplingMetadata
32 from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
33 from vllm.v1.worker import mamba_utils
34 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
35 from vllm.v1.worker.gpu_model_runner import GPUModelRunner
36 from vllm.v1.worker.mamba_utils import get_mamba_groups
37
38
39 @dataclass
40 class StepAction:
41 num_computed_tokens_start: int
42 num_scheduled_tokens: int
43 kv_cache_block_ids: list[int] # per-block mask: 1=held, 0=freed/nulled
44 preprocess_copy_idx: tuple[int, int] # -1, -1 for no copy
45 postprocess_copy_idx: tuple[int, int] # -1, -1 for no copy
46
47
48 num_speculative_tokens = 3
49
50 # Whether the run under test uses async scheduling. Set by each test entrypoint
51 # before generation; consulted where the scheduler's optimistic token count must
52 # be corrected for in-flight (possibly-rejected) speculative tokens.
53 async_scheduling_mode = False
54
55 num_accepted_tokens = 1
56 prompt_token_ids: list[int] = []
57 MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
58 # Must be a value no platform rewrites, since the step actions below encode
59 # where block boundaries fall. XPU rounds GDN block sizes up to a multiple of 64
60 # (XPUPlatform.update_block_size_for_backend), so use one.
61 BLOCK_SIZE = 576
62 DEVICE_TYPE = current_platform.device_type
63 NUM_HIDDEN_LAYERS = 1
64 cur_step_action_idx = 0
65 cur_step_action: StepAction | None = None
66 step_actions: list[StepAction] = []
67
68
69 def get_fake_sample_fn() -> SamplerOutput:
70 def fake_sample_fn(
71 self: GPUModelRunner,
72 logits: torch.Tensor | None,
73 spec_decode_metadata: SpecDecodeMetadata | None,
74 ) -> SamplerOutput:
75 assert logits is not None
76 num_computed_tokens_cpu_tensor = self.input_batch.num_computed_tokens_cpu_tensor
77 num_computed_tokens = num_computed_tokens_cpu_tensor[0].item()
78 if num_computed_tokens < self.input_batch.num_prompt_tokens[0].item():
79 first_token_id_index = self.input_batch.num_prompt_tokens[0].item()
80 else:
81 first_token_id_index = num_computed_tokens + 1
82 if spec_decode_metadata is None:
83 return SamplerOutput(
84 sampled_token_ids=async_tensor_h2d(
85 [[prompt_token_ids[first_token_id_index]]],
86 device=DEVICE_TYPE,
87 dtype=torch.int32,
88 ),
89 logprobs_tensors=None,
90 )
91 accepted_tokens = prompt_token_ids[
92 first_token_id_index : first_token_id_index
93 + min(num_accepted_tokens, logits.shape[0])
94 ]
95 sampled_token_ids = accepted_tokens
96 return SamplerOutput(
97 sampled_token_ids=async_tensor_h2d(
98 [sampled_token_ids],
99 device=DEVICE_TYPE,
100 dtype=torch.int32,
101 ),
102 logprobs_tensors=None,
103 )
104
105 return fake_sample_fn
106
107
108 def get_fake_propose_draft_token_ids_fn():
109 def fake_propose_draft_token_ids_fn(
110 self: GPUModelRunner,
111 scheduler_output: SchedulerOutput,
112 sampled_token_ids: torch.Tensor | list[list[int]],
113 sampling_metadata: SamplingMetadata,
114 hidden_states: torch.Tensor,
115 sample_hidden_states: torch.Tensor,
116 aux_hidden_states: list[torch.Tensor] | None,
117 spec_decode_metadata: SpecDecodeMetadata | None,
118 common_attn_metadata: CommonAttentionMetadata,
119 slot_mappings: dict[str, torch.Tensor] | list[dict[str, torch.Tensor]] | None,
120 ) -> list[list[int]]:
121 num_computed_tokens_cpu_tensor = self.input_batch.num_computed_tokens_cpu_tensor
122 num_computed_tokens = num_computed_tokens_cpu_tensor[0].item()
123 if (
124 self.input_batch.num_tokens_no_spec[0].item()
125 <= self.input_batch.num_prompt_tokens[0].item()
126 ):
127 first_token_id_index = self.input_batch.num_prompt_tokens[0].item()
128 else:
129 first_token_id_index = (
130 num_computed_tokens + 1
131 ) # bonus token isn't considered as computed
132 first_token_id_index += self.input_batch.num_accepted_tokens_cpu[0].item()
133 proposed_draft_token_ids = [
134 prompt_token_ids[
135 first_token_id_index : first_token_id_index + num_speculative_tokens
136 ]
137 ]
138
139 next_token_ids = async_tensor_h2d(
140 prompt_token_ids[
141 first_token_id_index - 1 : first_token_id_index
142 - 1
143 + num_accepted_tokens
144 ],
145 device=DEVICE_TYPE,
146 dtype=torch.int32,
147 )
148
149 valid_sampled_tokens_count = async_tensor_h2d(
150 [num_accepted_tokens],
151 device=DEVICE_TYPE,
152 dtype=torch.int32,
153 )
154
155 self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count)
156
157 return async_tensor_h2d(
158 proposed_draft_token_ids,
159 device=DEVICE_TYPE,
160 dtype=torch.int32,
161 )
162
163 return fake_propose_draft_token_ids_fn
164
165
166 def get_fake_step_action_fn(original_step_action_fn: Callable):
167 def fake_get_output(self: InprocClient):
168 global cur_step_action_idx
169 global cur_step_action
170 if cur_step_action_idx < len(step_actions):
171 cur_step_action = step_actions[cur_step_action_idx]
172 cur_step_action_idx += 1
173 else:
174 cur_step_action = None
175 print(f"cur_step_action: {cur_step_action_idx=} {cur_step_action=}")
176 return original_step_action_fn(self)
177
178 return fake_get_output
179
180
181 def get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable):
182 def fake_allocate_slots_fn(
183 self: KVCacheManager,
184 request: Request,
185 num_new_tokens: int,
186 num_new_computed_tokens: int = 0,
187 new_computed_blocks: KVCacheBlocks | None = None,
188 num_lookahead_tokens: int = 0,
189 num_external_computed_tokens: int = 0,
190 delay_cache_blocks: bool = False,
191 num_encoder_tokens: int = 0,
192 full_sequence_must_fit: bool = False,
193 reserved_blocks: int = 0,
194 has_scheduled_reqs: bool = True,
195 ):
196 ret = original_allocate_slots_fn(
197 self,
198 request,
199 num_new_tokens,
200 num_new_computed_tokens,
201 new_computed_blocks,
202 num_lookahead_tokens,
203 num_external_computed_tokens,
204 delay_cache_blocks,
205 num_encoder_tokens,
206 full_sequence_must_fit,
207 reserved_blocks,
208 has_scheduled_reqs,
209 )
210 if cur_step_action is not None:
211 cur_block_ids = self.coordinator.single_type_managers[0].req_to_blocks[
212 request.request_id
213 ]
214 not_null_block_flags = [not block.is_null for block in cur_block_ids]
215 block_ids = [1 if block else 0 for block in not_null_block_flags]
216 assert block_ids == cur_step_action.kv_cache_block_ids
217 return ret
218
219 return fake_allocate_slots_fn
220
221
222 mamba_kv_cache_dict = {}
223
224
225 def get_fake_execute_model_fn(original_execute_model_fn: Callable):
226 last_num_computed_tokens = 0
227 num_prompt_tokens = None
228
229 def fake_execute_model_fn(
230 self: GPUModelRunner,
231 scheduler_output: SchedulerOutput,
232 intermediate_tensors: IntermediateTensors | None = None,
233 ):
234 if cur_step_action is not None:
235 num_scheduled_tokens = next(
236 iter(scheduler_output.num_scheduled_tokens.values())
237 )
238 assert num_scheduled_tokens == cur_step_action.num_scheduled_tokens
239 mamba_groups = get_mamba_groups(self.kv_cache_config)
240 mamba_spec, mamba_group_ids = next(iter(mamba_groups.items()))
241 mamba_group_id = mamba_group_ids[0]
242 mamba_layer_name = self.kv_cache_config.kv_cache_groups[
243 mamba_group_id
244 ].layer_names[0]
245 nonlocal last_num_computed_tokens
246 nonlocal num_prompt_tokens
247
248 if (
249 len(scheduler_output.scheduled_new_reqs) > 0
250 and scheduler_output.scheduled_new_reqs[0].prompt_token_ids is not None
251 ):
252 # record number of prompt tokens
253 num_prompt_tokens = len(
254 scheduler_output.scheduled_new_reqs[0].prompt_token_ids
255 )
256
257 if len(scheduler_output.scheduled_cached_reqs.req_ids) > 0:
258 num_computed_tokens = (
259 scheduler_output.scheduled_cached_reqs.num_computed_tokens[0]
260 )
261 if (
262 async_scheduling_mode
263 and self.num_spec_tokens
264 and num_prompt_tokens is not None
265 and num_computed_tokens > num_prompt_tokens
266 ):
267 # NOTE (tdoublep) with async scheduling, the scheduler does not have an
268 # accurate measure of the number of computed tokens; we need to subtract
269 # the number of reject tokens from the previous timestep.
270 num_computed_tokens -= num_speculative_tokens + 1 - num_accepted_tokens
271 if (
272 num_computed_tokens // BLOCK_SIZE
273 > last_num_computed_tokens // BLOCK_SIZE
274 ):
275 # generated a new aligned block in this step
276 block_idx = num_computed_tokens // mamba_spec.block_size - 1
277 block_id = (
278 self.input_batch.block_table.block_tables[mamba_group_id]
279 .block_table.cpu[0, block_idx]
280 .item()
281 )
282 if block_id != 0:
283 kv_cache = self.compilation_config.static_forward_context[
284 mamba_layer_name
285 ].kv_cache
286 mamba_kv_cache_dict[
287 num_computed_tokens - num_computed_tokens % BLOCK_SIZE
288 ] = (
289 kv_cache[0][block_id].clone(),
290 kv_cache[1][block_id].clone(),
291 )
292
293 last_num_computed_tokens = num_computed_tokens
294 else:
295 last_num_computed_tokens = 0
296
297 ret = original_execute_model_fn(self, scheduler_output, intermediate_tensors)
298
299 if cur_step_action is not None:
300 assert (
301 cur_step_action.num_computed_tokens_start
302 == self.input_batch.num_computed_tokens_cpu[0].item()
303 )
304
305 return ret
306
307 return fake_execute_model_fn
308
309
310 def get_fake_process_mamba_fn(
311 original_preprocess_mamba_fn: Callable,
312 original_copy_fn: Callable,
313 ):
314 copy_info: tuple[list[int], list[int], list[int]] | None = None
315
316 def check_copy_info(
317 action: tuple[int, int],
318 kv_cache_config: KVCacheConfig,
319 forward_context: dict[str, Any],
320 input_batch: InputBatch,
321 ):
322 assert copy_info is not None
323 if action == (-1, -1):
324 assert len(copy_info[0]) == len(copy_info[1]) == len(copy_info[2]) == 0
325 else:
326 assert len(copy_info[0]) == len(copy_info[1]) == len(copy_info[2]) == 2
327 mamba_groups = get_mamba_groups(kv_cache_config)
328 mamba_spec, mamba_group_ids = next(iter(mamba_groups.items()))
329 mamba_group_id = mamba_group_ids[0]
330 mamba_layer_name = kv_cache_config.kv_cache_groups[
331 mamba_group_id
332 ].layer_names[0]
333 mamba_kv_cache = forward_context[mamba_layer_name].kv_cache[-1]
334 mamba_block_table = input_batch.block_table.block_tables[
335 mamba_group_id
336 ].block_table.cpu[0]
337 expected_temporal_src = mamba_kv_cache[
338 mamba_block_table[action[0]]
339 ].data_ptr()
340 expected_temporal_dest = mamba_kv_cache[
341 mamba_block_table[action[1]]
342 ].data_ptr()
343 # -1 is qwen3-next's temporal. We skip checking conv as it is more complex.
344 assert copy_info[0][-1] == expected_temporal_src
345 assert copy_info[1][-1] == expected_temporal_dest
346
347 def check_fused_copy_info(
348 action: tuple[int, int],
349 align_ctx: mamba_utils.MambaSpecDecodeGPUContext,
350 ):
351 # Align + spec-decode on a hybrid model routes the pre-copy through the
352 # fused kernel (preprocess_mamba -> run_fused_precopy) instead of
353 # do_mamba_copy_block, so copy_info is never populated. Verify from the
354 # fused buffers (req-0 scope, mirroring check_copy_info):
355 # - the copy DECISION: src_col is -1 iff no pre-copy is scheduled;
356 # - the DESTINATION column: state_idx == action[1] (curr_state_idx;
357 # maps directly through block_ids, exactly as check_copy_info's dst).
358 # The source column is NOT asserted here: on the scalar path the source
359 # address is produced by the per-state copy func with an accept-token
360 # bias offset (collect_mamba_copy_meta), so prev_state_idx does not map
361 # to action[0] by plain equality. Source block-level exactness (incl.
362 # the accept-bias) is covered by test_precopy_mamba_align.py.
363 src_col = int(align_ctx.precopy_src_col_buf.np[0])
364 state_idx = int(align_ctx.mamba_state_idx_buf.np[0])
365 if action == (-1, -1):
366 assert src_col == -1
367 else:
368 assert src_col != -1
369 assert state_idx == action[1]
370
371 def fake_preprocess_mamba_fn(
372 scheduler_output: SchedulerOutput,
373 kv_cache_config: KVCacheConfig,
374 cache_config: CacheConfig,
375 mamba_state_idx: dict[str, int],
376 input_batch: InputBatch,
377 requests: dict[str, CachedRequestState],
378 forward_context: dict[str, Any],
379 mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
380 copy_bufs: mamba_utils.MambaCopyBuffers,
381 align_ctx: mamba_utils.MambaSpecDecodeGPUContext | None = None,
382 ):
383 nonlocal copy_info
384 copy_info = None
385 ret = original_preprocess_mamba_fn(
386 scheduler_output,
387 kv_cache_config,
388 cache_config,
389 mamba_state_idx,
390 input_batch,
391 requests,
392 forward_context,
393 mamba_state_copy_funcs,
394 copy_bufs,
395 align_ctx,
396 )
397 if cur_step_action is not None:
398 if align_ctx is not None:
399 check_fused_copy_info(
400 cur_step_action.preprocess_copy_idx,
401 align_ctx,
402 )
403 else:
404 check_copy_info(
405 cur_step_action.preprocess_copy_idx,
406 kv_cache_config,
407 forward_context,
408 input_batch,
409 )
410 return ret
411
412 def fake_copy_fn(copy_bufs: mamba_utils.MambaCopyBuffers):
413 nonlocal copy_info
414 assert copy_info is None
415 n = copy_bufs.offset
416 src_state_list = copy_bufs.src_ptrs.cpu[:n].tolist()
417 dest_state_list = copy_bufs.dst_ptrs.cpu[:n].tolist()
418 num_elements_list = copy_bufs.sizes.cpu[:n].tolist()
419 copy_info = (src_state_list, dest_state_list, num_elements_list)
420 return original_copy_fn(copy_bufs)
421
422 return fake_preprocess_mamba_fn, fake_copy_fn
423
424
425 def run_ref_mamba_state_in_subprocess() -> None:
426 ctx = mp.get_context("spawn")
427 proc = ctx.Process(target=_run_ref_mamba_state_worker)
428 proc.start()
429 proc.join(timeout=600)
430 if proc.exitcode != 0:
431 raise RuntimeError(f"Ref mamba state process exited with code {proc.exitcode}.")
432
433
434 def _run_ref_mamba_state_worker():
435 try:
436 os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
437 num_generated_tokens = 8000
438 num_prompt_tokens = 500
439 sampling_params = SamplingParams(
440 temperature=0.0, max_tokens=num_generated_tokens
441 )
442 prompt_dataset = datasets.load_dataset("heheda/a_long_article")
443 full_prompt = prompt_dataset["train"][0]["text"]
444 fake_execute_model_fn = get_fake_execute_model_fn(GPUModelRunner.execute_model)
445 GPUModelRunner.execute_model = fake_execute_model_fn
446 fake_sample_fn = get_fake_sample_fn()
447 GPUModelRunner._sample = fake_sample_fn
448 engine = LLM(
449 model=MODEL,
450 load_format="dummy",
451 block_size=BLOCK_SIZE,
452 hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS},
453 seed=42,
454 )
455 global prompt_token_ids
456 prompt_token_ids = engine.get_tokenizer().encode(full_prompt)
457 print(f"Token IDs length: {len(prompt_token_ids)}")
458
459 _outputs = engine.generate(
460 [TokensPrompt(prompt_token_ids=prompt_token_ids[:num_prompt_tokens])],
461 sampling_params,
462 )
463 # ref_mamba_kv_cache_dict = torch.load("mamba_kv_cache_dict.pth")
464 # check_mamba_state_equal(ref_mamba_kv_cache_dict, mamba_kv_cache_dict)
465 # torch.save(mamba_kv_cache_dict, "mamba_kv_cache_dict.pth")
466 cpu_state_ref = {
467 key: tuple(tensor.detach().cpu() for tensor in tensors)
468 for key, tensors in mamba_kv_cache_dict.items()
469 }
470 torch.save(cpu_state_ref, "mamba_kv_cache_dict_ref.pth")
471 mamba_kv_cache_dict.clear()
472 del engine
473 torch.accelerator.empty_cache()
474 cleanup_dist_env_and_memory()
475 except Exception:
476 traceback.print_exc()
477 raise
478
479
480 def check_mamba_state_equal(
481 mamba_state_ref: dict, mamba_state_new: dict, keys_to_check: list[int]
482 ):
483 atol = 1e-2
484 rtol = 1e-2
485 for key in keys_to_check:
486 assert key in mamba_state_new
487 assert key in mamba_state_ref
488 # mamba state new is a subset of mamba state ref
489 for i, (ref, new) in enumerate(zip(mamba_state_ref[key], mamba_state_new[key])):
490 if ref.device != new.device:
491 new = new.to(ref.device)
492 new = new[: ref.shape[0]]
493 if not torch.allclose(ref, new, atol=atol, rtol=rtol):
494 diff_mask = ~torch.isclose(ref, new, atol=atol, rtol=rtol)
495 diff_idx = torch.nonzero(diff_mask)
496 if diff_idx.shape[0] * 100 < ref.numel():
497 print(
498 f"[WARNING] found {diff_idx.shape[0] * 100 / ref.numel()}% of the elements are different" # noqa: E501
499 )
500 continue
501 raise ValueError(
502 f"Mamba state is not equal for key: {key} at index {i}"
503 )
504 return True
505
506
507 @dataclass
508 class TestConfig:
509 num_prompt_tokens: int
510 num_generated_tokens: int
511 num_accepted_tokens: int
512 step_actions: list[StepAction]
513
514
515 def apply_patch(monkeypatch: pytest.MonkeyPatch):
516 monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
517
518 fake_sample_fn = get_fake_sample_fn()
519 monkeypatch.setattr(GPUModelRunner, "_sample", fake_sample_fn)
520
521 fake_propose_draft_token_ids_fn = get_fake_propose_draft_token_ids_fn()
522 monkeypatch.setattr(
523 GPUModelRunner, "propose_draft_token_ids", fake_propose_draft_token_ids_fn
524 )
525
526 fake_execute_model_fn = get_fake_execute_model_fn(GPUModelRunner.execute_model)
527 monkeypatch.setattr(GPUModelRunner, "execute_model", fake_execute_model_fn)
528
529 fake_step_action_fn = get_fake_step_action_fn(InprocClient.get_output)
530 monkeypatch.setattr(InprocClient, "get_output", fake_step_action_fn)
531
532 fake_allocate_slots_fn = get_fake_allocate_slots_fn(KVCacheManager.allocate_slots)
533 monkeypatch.setattr(KVCacheManager, "allocate_slots", fake_allocate_slots_fn)
534
535 fake_preprocess_mamba_fn, fake_copy_fn = get_fake_process_mamba_fn(
536 mamba_utils.preprocess_mamba,
537 mamba_utils.do_mamba_copy_block,
538 )
539 monkeypatch.setattr(mamba_utils, "preprocess_mamba", fake_preprocess_mamba_fn)
540 monkeypatch.setattr(mamba_utils, "do_mamba_copy_block", fake_copy_fn)
541
542
543 def get_mamba_prefix_cache_step_configs(
544 async_scheduling: bool = False,
545 ) -> dict[str, TestConfig]:
546 a = async_scheduling
547 B = BLOCK_SIZE
548 tests = {
549 "accept_1": TestConfig(
550 num_prompt_tokens=B - 6,
551 num_generated_tokens=20,
552 num_accepted_tokens=1,
553 step_actions=[
554 StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)),
555 StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
556 StepAction(
557 B - 5, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
558 ),
559 StepAction(
560 B - 4, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
561 ),
562 StepAction(B - 3, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
563 StepAction(B - 2, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
564 StepAction(B - 1, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)),
565 StepAction(B, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
566 StepAction(
567 B + 1,
568 4,
569 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
570 (-1, -1),
571 (-1, -1),
572 ),
573 ],
574 ),
575 # test case 2.1: no hit, accept 2 tokens
576 "accept_2_1": TestConfig(
577 num_prompt_tokens=B - 6,
578 num_generated_tokens=20,
579 num_accepted_tokens=2,
580 step_actions=[
581 StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)),
582 StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
583 StepAction(
584 B - 4, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
585 ),
586 StepAction(B - 2, 4, [1, 1, 1, 1, 1], (1, 1), (2, 0)),
587 StepAction(B, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
588 StepAction(
589 B + 2,
590 4,
591 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
592 (-1, -1),
593 (-1, -1),
594 ),
595 ],
596 ),
597 # test case 2.2: no hit, accept 2 tokens
598 "accept_2_2": TestConfig(
599 num_prompt_tokens=B - 5,
600 num_generated_tokens=20,
601 num_accepted_tokens=2,
602 step_actions=[
603 StepAction(0, B - 5, [1, 1, 1, 1], (-1, -1), (-1, -1)),
604 StepAction(B - 5, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
605 StepAction(B - 3, 4, [1, 1, 1, 1, 1], (1, 1), (-1, -1)),
606 StepAction(B - 1, 4, [1, 1, 1, 1, 1], (-1, -1), (1, 0)),
607 StepAction(
608 B + 1,
609 4,
610 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
611 (-1, -1),
612 (-1, -1),
613 ),
614 StepAction(B + 3, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
615 ],
616 ),
617 "accept_3_1": TestConfig(
618 num_prompt_tokens=B - 7,
619 num_generated_tokens=20,
620 num_accepted_tokens=3,
621 step_actions=[
622 StepAction(0, B - 7, [1, 1, 1, 1], (-1, -1), (-1, -1)),
623 StepAction(B - 7, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
624 StepAction(
625 B - 4, 4, [1, 1, 1, 1, 1] if a else [1, 1, 1, 1], (-1, -1), (-1, -1)
626 ),
627 StepAction(B - 1, 4, [1, 1, 1, 1, 1], (2, 1), (1, 0)),
628 StepAction(
629 B + 2,
630 4,
631 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
632 (-1, -1),
633 (-1, -1),
634 ),
635 StepAction(B + 5, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
636 ],
637 ),
638 "accept_3_2": TestConfig(
639 num_prompt_tokens=B - 6,
640 num_generated_tokens=20,
641 num_accepted_tokens=3,
642 step_actions=[
643 StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)),
644 StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
645 StepAction(B - 3, 4, [1, 1, 1, 1, 1], (2, 1), (3, 0)),
646 StepAction(B, 4, [1, 1, 1, 1, 1], (-1, -1), (-1, -1)),
647 StepAction(
648 B + 3,
649 4,
650 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
651 (-1, -1),
652 (-1, -1),
653 ),
654 ],
655 ),
656 "accept_3_3": TestConfig(
657 num_prompt_tokens=B - 5,
658 num_generated_tokens=20,
659 num_accepted_tokens=3,
660 step_actions=[
661 StepAction(0, B - 5, [1, 1, 1, 1], (-1, -1), (-1, -1)),
662 StepAction(B - 5, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
663 StepAction(B - 2, 4, [1, 1, 1, 1, 1], (2, 1), (2, 0)),
664 StepAction(
665 B + 1,
666 4,
667 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
668 (-1, -1),
669 (-1, -1),
670 ),
671 StepAction(B + 4, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
672 ],
673 ),
674 "accept_4_1": TestConfig(
675 num_prompt_tokens=B - 7,
676 num_generated_tokens=20,
677 num_accepted_tokens=4,
678 step_actions=[
679 StepAction(0, B - 7, [1, 1, 1, 1], (-1, -1), (-1, -1)),
680 StepAction(B - 7, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
681 StepAction(B - 3, 4, [1, 1, 1, 1, 1], (3, 1), (3, 0)),
682 StepAction(
683 B + 1,
684 4,
685 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
686 (-1, -1),
687 (-1, -1),
688 ),
689 StepAction(B + 5, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
690 ],
691 ),
692 "accept_4_2": TestConfig(
693 num_prompt_tokens=B - 6,
694 num_generated_tokens=25,
695 num_accepted_tokens=4,
696 step_actions=[
697 StepAction(0, B - 6, [1, 1, 1, 1], (-1, -1), (-1, -1)),
698 StepAction(B - 6, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
699 StepAction(B - 2, 4, [1, 1, 1, 1, 1], (3, 1), (2, 0)),
700 StepAction(
701 B + 2,
702 4,
703 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
704 (-1, -1),
705 (-1, -1),
706 ),
707 StepAction(B + 6, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
708 ],
709 ),
710 "accept_4_3": TestConfig(
711 num_prompt_tokens=B - 5,
712 num_generated_tokens=25,
713 num_accepted_tokens=4,
714 step_actions=[
715 StepAction(0, B - 5, [1, 1, 1, 1], (-1, -1), (-1, -1)),
716 StepAction(B - 5, 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
717 StepAction(B - 1, 4, [1, 1, 1, 1, 1], (3, 1), (1, 0)),
718 StepAction(
719 B + 3,
720 4,
721 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
722 (-1, -1),
723 (-1, -1),
724 ),
725 StepAction(B + 7, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
726 ],
727 ),
728 "accept_4_4": TestConfig(
729 num_prompt_tokens=B - 4,
730 num_generated_tokens=25,
731 num_accepted_tokens=4,
732 step_actions=[
733 StepAction(0, B - 4, [1, 1, 1, 1], (-1, -1), (-1, -1)),
734 StepAction(B - 4, 4, [1, 1, 1, 1], (-1, -1), (3, 0)),
735 StepAction(B, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
736 StepAction(
737 B + 4,
738 4,
739 [1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1],
740 (-1, -1),
741 (-1, -1),
742 ),
743 ],
744 ),
745 "prompt_block_size": TestConfig(
746 num_prompt_tokens=B,
747 num_generated_tokens=10,
748 num_accepted_tokens=4,
749 step_actions=[
750 StepAction(0, B, [1, 1, 1, 1], (-1, -1), (-1, -1)),
751 StepAction(B, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
752 ],
753 ),
754 "prompt_2_block_size": TestConfig(
755 num_prompt_tokens=B * 2,
756 num_generated_tokens=10,
757 num_accepted_tokens=4,
758 step_actions=[
759 StepAction(0, B, [1, 1, 1, 1], (-1, -1), (-1, -1)),
760 StepAction(B, B, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
761 StepAction(
762 B * 2,
763 4,
764 [1, 1, 1, 1, 1, 1] if a else [0, 1, 1, 1, 1, 1],
765 (1, 2),
766 (-1, -1),
767 ),
768 ],
769 ),
770 "prompt_2_block_size_10": TestConfig(
771 num_prompt_tokens=B * 2 + 10,
772 num_generated_tokens=10,
773 num_accepted_tokens=4,
774 step_actions=[
775 StepAction(0, B, [1, 1, 1, 1], (-1, -1), (-1, -1)),
776 StepAction(B, B + 10, [1, 0, 1, 1, 1, 1], (0, 2), (-1, -1)),
777 StepAction(
778 B * 2 + 10,
779 4,
780 [1, 0, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1],
781 (-1, -1),
782 (-1, -1),
783 ),
784 ],
785 ),
786 "prompt_3_block_size": TestConfig(
787 num_prompt_tokens=B * 3,
788 num_generated_tokens=10,
789 num_accepted_tokens=4,
790 step_actions=[
791 StepAction(0, B * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
792 StepAction(B * 2, B, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
793 StepAction(
794 B * 3,
795 4,
796 [0, 1, 1, 1, 1, 1, 1] if a else [0, 0, 1, 1, 1, 1, 1],
797 (2, 3),
798 (-1, -1),
799 ),
800 ],
801 ),
802 "prompt_3_block_size_10": TestConfig(
803 num_prompt_tokens=B * 3 + 10,
804 num_generated_tokens=10,
805 num_accepted_tokens=4,
806 step_actions=[
807 StepAction(0, B * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
808 StepAction(B * 2, B + 10, [0, 1, 0, 1, 1, 1, 1], (1, 3), (-1, -1)),
809 StepAction(
810 B * 3 + 10,
811 4,
812 [0, 1, 0, 1, 1, 1, 1] if a else [0, 0, 0, 1, 1, 1, 1],
813 (-1, -1),
814 (-1, -1),
815 ),
816 ],
817 ),
818 "prompt_10_block_size": TestConfig(
819 num_prompt_tokens=B * 10,
820 num_generated_tokens=10,
821 num_accepted_tokens=4,
822 step_actions=[
823 StepAction(0, B * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
824 StepAction(
825 B * 5,
826 B * 4,
827 [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1],
828 (4, 8),
829 (-1, -1),
830 ),
831 StepAction(
832 B * 9,
833 B,
834 [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1, 1]
835 if a
836 else [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
837 (8, 9),
838 (-1, -1),
839 ),
840 StepAction(
841 B * 10,
842 4,
843 [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1]
844 if a
845 else [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
846 (9, 10),
847 (-1, -1),
848 ),
849 ],
850 ),
851 "prompt_10_block_size_10": TestConfig(
852 num_prompt_tokens=B * 10 + 10,
853 num_generated_tokens=10,
854 num_accepted_tokens=4,
855 step_actions=[
856 StepAction(0, B * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
857 StepAction(
858 B * 5,
859 B * 4,
860 [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1],
861 (4, 8),
862 (-1, -1),
863 ),
864 StepAction(
865 B * 9,
866 B + 10,
867 [0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 1, 1, 1, 1]
868 if a
869 else [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1],
870 (8, 10),
871 (-1, -1),
872 ),
873 ],
874 ),
875 }
876 return tests
877
878
879 def _run_mamba_prefix_cache_mrv1(
880 monkeypatch: pytest.MonkeyPatch, async_scheduling: bool
881 ):
882 # This test patches the V1 model runner, so pin V1 explicitly: MoE/hybrid
883 # models like Qwen3-Next now default to the V2 runner.
884 monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0")
885 envs.disable_envs_cache()
886 global async_scheduling_mode
887 async_scheduling_mode = async_scheduling
888 run_ref_mamba_state_in_subprocess()
889 apply_patch(monkeypatch)
890 prompt_dataset = datasets.load_dataset("heheda/a_long_article")
891 full_prompt = prompt_dataset["train"][0]["text"]
892 tests = get_mamba_prefix_cache_step_configs(async_scheduling)
893
894 engine = LLM(
895 model=MODEL,
896 load_format="dummy",
897 enable_prefix_caching=True,
898 block_size=BLOCK_SIZE,
899 mamba_cache_mode="align",
900 speculative_config={
901 "method": "qwen3_next_mtp",
902 "num_speculative_tokens": num_speculative_tokens,
903 },
904 max_num_batched_tokens=3072,
905 hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS},
906 async_scheduling=async_scheduling,
907 seed=42,
908 )
909 global prompt_token_ids
910 prompt_token_ids = engine.get_tokenizer().encode(full_prompt)
911 print(f"Token IDs length: {len(prompt_token_ids)}")
912 for test_case_name, test_config in tests.items():
913 print(f"Running test case: {test_case_name}")
914 num_generated_tokens = test_config.num_generated_tokens
915 num_prompt_tokens = test_config.num_prompt_tokens
916 global num_accepted_tokens
917 num_accepted_tokens = test_config.num_accepted_tokens
918 sampling_params = SamplingParams(
919 temperature=0.0, max_tokens=num_generated_tokens
920 )
921 global cur_step_action_idx
922 cur_step_action_idx = 0
923 global step_actions
924 step_actions = test_config.step_actions
925 _ = engine.generate(
926 [TokensPrompt(prompt_token_ids=prompt_token_ids[:num_prompt_tokens])],
927 sampling_params,
928 )
929 assert engine.llm_engine.engine_core.engine_core.scheduler.reset_prefix_cache()
930 print(f"End test case: {test_case_name}")
931 keys_to_check = [
932 (action.postprocess_copy_idx[1] + 1) * BLOCK_SIZE
933 for action in test_config.step_actions
934 if action.postprocess_copy_idx and action.postprocess_copy_idx[0] != -1
935 ]
936 mamba_state_ref = torch.load("mamba_kv_cache_dict_ref.pth")
937 check_mamba_state_equal(mamba_state_ref, mamba_kv_cache_dict, keys_to_check)
938 mamba_kv_cache_dict.clear()
939 del engine
940 torch.accelerator.empty_cache()
941 cleanup_dist_env_and_memory()
942
943
944 @create_new_process_for_each_test("spawn")
945 def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch):
946 _run_mamba_prefix_cache_mrv1(monkeypatch, async_scheduling=False)
947
948
949 @create_new_process_for_each_test("spawn")
950 def test_mamba_prefix_cache_mrv1_async(monkeypatch: pytest.MonkeyPatch):
951 _run_mamba_prefix_cache_mrv1(monkeypatch, async_scheduling=True)
952
953
954 def _run_mamba_prefix_cache_mrv2(
955 monkeypatch: pytest.MonkeyPatch, async_scheduling: bool
956 ):
957 global async_scheduling_mode
958 async_scheduling_mode = async_scheduling
959 monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
960 monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
961 envs.disable_envs_cache()
962
963 from vllm.v1.worker.gpu.model_runner import GPUModelRunner as MRV2GPUModelRunner
964 from vllm.v1.worker.gpu.model_states.mamba_hybrid import (
965 MambaHybridModelState,
966 )
967 from vllm.v1.worker.gpu.sample.output import SamplerOutput as MRV2SamplerOutput
968
969 events: list[int] = []
970 original_execute_model = MRV2GPUModelRunner.execute_model
971 original_sample = MRV2GPUModelRunner.sample
972 original_preprocess_state = MambaHybridModelState.preprocess_state
973 original_postprocess_state = MambaHybridModelState.postprocess_state
974 original_step_action_fn = InprocClient.get_output
975 original_allocate_slots = KVCacheManager.allocate_slots
976 captured: dict[str, Any] = {}
977
978 def temporal_states(model_state, block_tables, kv_cache_config):
979 # Qwen3-Next keeps the temporal (ssm) state as the last Mamba cache.
980 forward_context = (
981 model_state.vllm_config.compilation_config.static_forward_context
982 )
983 mamba_groups = get_mamba_groups(kv_cache_config)
984 for group_ids in mamba_groups.values():
985 for group_id in group_ids:
986 block_table = block_tables[group_id]
987 for layer_name in kv_cache_config.kv_cache_groups[group_id].layer_names:
988 yield forward_context[layer_name].kv_cache[-1], block_table
989
990 def temporal_block(temporal_state, block_table, col):
991 # Resolving the block id for assertions is a deliberate D2H.
992 with gpu_sync_allowed():
993 block_id = int(block_table[0, col].item())
994 return temporal_state[block_id]
995
996 def wrapped_preprocess_state(
997 self: MambaHybridModelState,
998 input_batch: Any,
999 block_tables: tuple[torch.Tensor, ...],
1000 kv_cache_config: KVCacheConfig,
1001 num_computed_tokens: torch.Tensor,
1002 ) -> None:
1003 captured["block_tables"] = block_tables
1004 captured["kv_cache_config"] = kv_cache_config
1005 expected = (
1006 None if cur_step_action is None else cur_step_action.preprocess_copy_idx
1007 )
1008 snapshots = []
1009 if expected is not None and expected != (-1, -1):
1010 for temporal, bt in temporal_states(self, block_tables, kv_cache_config):
1011 snapshots.append(
1012 (temporal, bt, temporal_block(temporal, bt, expected[0]).clone())
1013 )
1014 ret = original_preprocess_state(
1015 self, input_batch, block_tables, kv_cache_config, num_computed_tokens
1016 )
1017 if cur_step_action is not None:
1018 # Reading the GPU-side copy state back to assert on it is a
1019 # deliberate D2H.
1020 with gpu_sync_allowed():
1021 req_idx = int(input_batch.idx_mapping[0].item())
1022 src_col = int(self._mamba_src_col_gpu[req_idx].item())
1023 off = int(self._mamba_src_off_gpu[req_idx].item())
1024 dst = int(self._mamba_state_idx_gpu[req_idx].item())
1025 actual = (-1, -1) if src_col < 0 or src_col == dst else (src_col + off, dst)
1026 assert actual == expected, (
1027 f"V2 align preprocess copy: expected={expected}, "
1028 f"actual={actual}, {cur_step_action=}"
1029 )
1030 # Comparing device tensors for the assertion is a deliberate D2H.
1031 with gpu_sync_allowed():
1032 for temporal, bt, src_state in snapshots:
1033 torch.testing.assert_close(
1034 temporal_block(temporal, bt, expected[1]), src_state
1035 )
1036 return ret
1037
1038 def wrapped_postprocess_state(
1039 self: MambaHybridModelState,
1040 idx_mapping: torch.Tensor,
1041 num_sampled: torch.Tensor | int,
1042 num_computed_tokens: torch.Tensor | None = None,
1043 ) -> None:
1044 action = cur_step_action
1045 block_tables = captured.get("block_tables")
1046 kv_cache_config = captured.get("kv_cache_config")
1047 # The postprocess kernel does not expose its indices, so only the copy
1048 # case is checked, by effect: snapshot the src block, expect dst == src.
1049 if (
1050 action is None
1051 or num_computed_tokens is None
1052 or block_tables is None
1053 or action.postprocess_copy_idx == (-1, -1)
1054 ):
1055 return original_postprocess_state(
1056 self, idx_mapping, num_sampled, num_computed_tokens
1057 )
1058 expected = action.postprocess_copy_idx
1059 snapshots = [
1060 (temporal, bt, temporal_block(temporal, bt, expected[0]).clone())
1061 for temporal, bt in temporal_states(self, block_tables, kv_cache_config)
1062 ]
1063 ret = original_postprocess_state(
1064 self, idx_mapping, num_sampled, num_computed_tokens
1065 )
1066 # Comparing device tensors for the assertion is a deliberate D2H.
1067 with gpu_sync_allowed():
1068 for temporal, bt, src_state in snapshots:
1069 torch.testing.assert_close(
1070 temporal_block(temporal, bt, expected[1]), src_state
1071 )
1072 return ret
1073
1074 def wrapped_execute_model(
1075 self: MRV2GPUModelRunner,
1076 scheduler_output: SchedulerOutput,
1077 *args: Any,
1078 **kwargs: Any,
1079 ):
1080 events.extend(
1081 req.num_computed_tokens for req in scheduler_output.scheduled_new_reqs
1082 )
1083 events.extend(scheduler_output.scheduled_cached_reqs.num_computed_tokens)
1084 if cur_step_action is not None:
1085 num_scheduled_tokens = next(
1086 iter(scheduler_output.num_scheduled_tokens.values())
1087 )
1088 assert num_scheduled_tokens == cur_step_action.num_scheduled_tokens
1089 ret = original_execute_model(self, scheduler_output, *args, **kwargs)
1090 if cur_step_action is not None and self.execute_model_state is not None:
1091 input_batch = self.execute_model_state.input_batch
1092 # Reading positions back to assert on them is a deliberate D2H.
1093 with gpu_sync_allowed():
1094 start_pos = input_batch.positions[input_batch.query_start_loc[0]].item()
1095 assert cur_step_action.num_computed_tokens_start == start_pos
1096 return ret
1097
1098 def fake_sample(
1099 self: MRV2GPUModelRunner,
1100 hidden_states: torch.Tensor,
1101 input_batch: Any,
1102 grammar_output: Any,
1103 ):
1104 if cur_step_action is None:
1105 return original_sample(self, hidden_states, input_batch, grammar_output)
1106
1107 num_reqs = input_batch.num_reqs
1108 sampled_token_ids = torch.ones(
1109 (num_reqs, self.num_speculative_steps + 1),
1110 device=hidden_states.device,
1111 dtype=torch.int64,
1112 )
1113 num_logits = async_tensor_h2d(
1114 input_batch.cu_num_logits_np[1 : num_reqs + 1]
1115 - input_batch.cu_num_logits_np[:num_reqs],
1116 device=hidden_states.device,
1117 )
1118 accepted = torch.full_like(num_logits, num_accepted_tokens)
1119 num_sampled = torch.minimum(accepted, num_logits)
1120 prefill_lens = self.req_states.prefill_len.gpu[input_batch.idx_mapping]
1121 is_chunked_prefill = input_batch.seq_lens[:num_reqs] < prefill_lens
1122 num_sampled = torch.where(is_chunked_prefill, 0, num_sampled)
1123 num_rejected = torch.where(is_chunked_prefill, 0, num_logits - num_sampled)
1124 sampler_output = MRV2SamplerOutput(
1125 sampled_token_ids=sampled_token_ids,
1126 logprobs_tensors=None,
1127 num_nans=None,
1128 num_sampled=num_sampled,
1129 num_rejected=num_rejected,
1130 )
1131 return sampler_output, num_sampled, num_rejected
1132
1133 monkeypatch.setattr(
1134 InprocClient,
1135 "get_output",
1136 get_fake_step_action_fn(original_step_action_fn),
1137 )
1138 monkeypatch.setattr(
1139 KVCacheManager,
1140 "allocate_slots",
1141 get_fake_allocate_slots_fn(original_allocate_slots),
1142 )
1143 monkeypatch.setattr(MRV2GPUModelRunner, "execute_model", wrapped_execute_model)
1144 monkeypatch.setattr(MRV2GPUModelRunner, "sample", fake_sample)
1145 monkeypatch.setattr(
1146 MambaHybridModelState, "preprocess_state", wrapped_preprocess_state
1147 )
1148 monkeypatch.setattr(
1149 MambaHybridModelState, "postprocess_state", wrapped_postprocess_state
1150 )
1151
1152 engine = LLM(
1153 model=MODEL,
1154 load_format="dummy",
1155 enforce_eager=True,
1156 skip_tokenizer_init=True,
1157 enable_prefix_caching=True,
1158 block_size=BLOCK_SIZE,
1159 mamba_cache_mode="align",
1160 speculative_config={
1161 "method": "qwen3_next_mtp",
1162 "num_speculative_tokens": num_speculative_tokens,
1163 },
1164 max_num_batched_tokens=3072,
1165 max_model_len=BLOCK_SIZE * 12,
1166 hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS},
1167 async_scheduling=async_scheduling,
1168 seed=42,
1169 )
1170
1171 try:
1172 tests = get_mamba_prefix_cache_step_configs(async_scheduling)
1173
1174 global step_actions
1175 global cur_step_action_idx
1176 global num_accepted_tokens
1177 for test_name, test_config in tests.items():
1178 num_accepted_tokens = test_config.num_accepted_tokens
1179 cur_step_action_idx = 0
1180 step_actions = test_config.step_actions
1181 sampling_params = SamplingParams(
1182 temperature=0.0,
1183 max_tokens=test_config.num_generated_tokens,
1184 ignore_eos=True,
1185 )
1186 _ = engine.generate(
1187 [TokensPrompt(prompt_token_ids=[1] * test_config.num_prompt_tokens)],
1188 sampling_params=sampling_params,
1189 )
1190 assert cur_step_action_idx == len(test_config.step_actions), test_name
1191 assert (
1192 engine.llm_engine.engine_core.engine_core.scheduler.reset_prefix_cache()
1193 )
1194
1195 step_actions = []
1196 cur_step_action_idx = 0
1197 num_accepted_tokens = 1
1198 prompt = TokensPrompt(prompt_token_ids=[1] * (BLOCK_SIZE * 2))
1199 sampling_params = SamplingParams(
1200 temperature=0.0,
1201 max_tokens=1,
1202 ignore_eos=True,
1203 )
1204 _ = engine.generate([prompt], sampling_params=sampling_params)
1205 first_event_count = len(events)
1206 _ = engine.generate([prompt], sampling_params=sampling_params)
1207 second_events = events[first_event_count:]
1208 prefix_hits = [
1209 num_computed_tokens
1210 for num_computed_tokens in second_events
1211 if num_computed_tokens >= BLOCK_SIZE
1212 ]
1213 assert prefix_hits, (
1214 "Expected the second identical prompt to hit prefix cache, "
1215 f"got events={second_events!r}"
1216 )
1217 assert engine.llm_engine.engine_core.engine_core.scheduler.reset_prefix_cache()
1218 finally:
1219 del engine
1220 torch.accelerator.empty_cache()
1221 cleanup_dist_env_and_memory()
1222
1223
1224 @create_new_process_for_each_test()
1225 def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch):
1226 _run_mamba_prefix_cache_mrv2(monkeypatch, async_scheduling=False)
1227
1228
1229 @create_new_process_for_each_test()
1230 def test_mamba_prefix_cache_mrv2_async(monkeypatch: pytest.MonkeyPatch):
1231 _run_mamba_prefix_cache_mrv2(monkeypatch, async_scheduling=True)