@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5206e5e28c
Fork 0

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 9, 18:15 UTC

[V1][Hybrid] Mamba Prefix Caching with align mode (#30877)

Signed-off-by: huanghaoyan.hhy <huanghaoyan.hhy@alibaba-inc.com> Signed-off-by: Chen Zhang <zhangch99@outlook.com> Co-authored-by: Chen Zhang <zhangch99@outlook.com>

Harry Huang committed Jan 24, 2026 at 01:56 UTC 5206e5e28c25dcc43d3a2373b591a0cdf9792f5d
42 files changed +1775 -129
tests/v1/core/test_single_type_kv_cache_manager.py
+18 -7
@@ -24,7 +24,7 @@ pytestmark = pytest.mark.cpu_test
24 def get_sliding_window_manager(sliding_window_spec, block_pool, enable_caching=True):
25 return SlidingWindowManager(
26 sliding_window_spec,
27 - block_pool,
27 + block_pool=block_pool,
28 enable_caching=enable_caching,
29 kv_cache_group_id=0,
30 )
@@ -35,7 +35,7 @@ def get_chunked_local_attention_manager(
35 ):
36 return ChunkedLocalAttentionManager(
37 chunked_local_attention_spec,
38 - block_pool,
38 + block_pool=block_pool,
39 enable_caching=enable_caching,
40 kv_cache_group_id=0,
41 )
@@ -342,11 +342,15 @@ def test_get_num_blocks_to_allocate():
342 ]
343
344 assert (
345 - manager.get_num_blocks_to_allocate("1", 20 * block_size, cached_blocks_1, 0)
345 + manager.get_num_blocks_to_allocate(
346 + "1", 20 * block_size, cached_blocks_1, 0, 20 * block_size
347 + )
348 == 20
349 )
350 assert (
349 - manager.get_num_blocks_to_allocate("2", 20 * block_size, cached_blocks_2, 0)
351 + manager.get_num_blocks_to_allocate(
352 + "2", 20 * block_size, cached_blocks_2, 0, 20 * block_size
353 + )
354 == 15
355 )
356
@@ -375,6 +379,7 @@ def test_evictable_cached_blocks_not_double_allocated():
379 num_tokens=2 * block_size,
380 new_computed_blocks=[evictable_block],
381 total_computed_tokens=block_size,
382 + num_tokens_main_model=2 * block_size,
383 )
384 # Free capacity check should count evictable cached blocks, but allocation
385 # should only allocate the truly new block.
@@ -386,7 +391,9 @@ def test_evictable_cached_blocks_not_double_allocated():
391 num_local_computed_tokens=block_size,
392 num_external_computed_tokens=0,
393 )
389 - new_blocks = manager.allocate_new_blocks(request_id, num_tokens=4)
394 + new_blocks = manager.allocate_new_blocks(
395 + request_id, num_tokens=4, num_tokens_main_model=4
396 + )
397 assert len(new_blocks) == 1
398 assert len(manager.req_to_blocks[request_id]) == 2
399
@@ -411,10 +418,14 @@ def test_chunked_local_attention_get_num_blocks_to_allocate():
418 ]
419
420 assert (
414 - manager.get_num_blocks_to_allocate("1", 20 * block_size, cached_blocks_1, 0)
421 + manager.get_num_blocks_to_allocate(
422 + "1", 20 * block_size, cached_blocks_1, 0, 20 * block_size
423 + )
424 == 20
425 )
426 assert (
418 - manager.get_num_blocks_to_allocate("2", 20 * block_size, cached_blocks_2, 0)
427 + manager.get_num_blocks_to_allocate(
428 + "2", 20 * block_size, cached_blocks_2, 0, 20 * block_size
429 + )
430 == 15
431 )
tests/v1/e2e/test_mamba_prefix_cache.py new
+764
@@ -0,0 +1,764 @@
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 +from vllm import LLM, SamplingParams, TokensPrompt
15 +from vllm.config import CacheConfig
16 +from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
17 +from vllm.sequence import IntermediateTensors
18 +from vllm.v1.attention.backends.utils import CommonAttentionMetadata
19 +from vllm.v1.core.kv_cache_manager import KVCacheBlocks, KVCacheManager
20 +from vllm.v1.core.sched.output import SchedulerOutput
21 +from vllm.v1.engine.core_client import InprocClient
22 +from vllm.v1.kv_cache_interface import KVCacheConfig
23 +from vllm.v1.outputs import SamplerOutput
24 +from vllm.v1.request import Request
25 +from vllm.v1.sample.metadata import SamplingMetadata
26 +from vllm.v1.spec_decode.metadata import SpecDecodeMetadata
27 +from vllm.v1.worker import mamba_utils
28 +from vllm.v1.worker.gpu_input_batch import CachedRequestState
29 +from vllm.v1.worker.gpu_model_runner import GPUModelRunner
30 +from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch
31 +from vllm.v1.worker.mamba_utils import get_mamba_groups
32 +
33 +
34 +@dataclass
35 +class StepAction:
36 + num_computed_tokens_start: int
37 + num_scheduled_tokens: int
38 + kv_cache_block_ids: list[int] # [] to follow last step
39 + preprocess_copy_idx: tuple[int, int] # -1, -1 for no copy
40 + postprocess_copy_idx: tuple[int, int] # -1, -1 for no copy
41 +
42 +
43 +num_speculative_tokens = 3
44 +
45 +num_accepted_tokens = 1
46 +prompt_token_ids: list[int] = []
47 +MODEL = "Qwen/Qwen3-Next-80B-A3B-Instruct-FP8"
48 +BLOCK_SIZE = 560
49 +NUM_HIDDEN_LAYERS = 1
50 +cur_step_action_idx = 0
51 +cur_step_action: StepAction | None = None
52 +step_actions: list[StepAction] = []
53 +
54 +
55 +def get_fake_sample_fn() -> SamplerOutput:
56 + def fake_sample_fn(
57 + self: GPUModelRunner,
58 + logits: torch.Tensor | None,
59 + spec_decode_metadata: SpecDecodeMetadata | None,
60 + ) -> SamplerOutput:
61 + assert logits is not None
62 + num_computed_tokens_cpu_tensor = self.input_batch.num_computed_tokens_cpu_tensor
63 + num_computed_tokens = num_computed_tokens_cpu_tensor[0].item()
64 + if num_computed_tokens < self.input_batch.num_prompt_tokens[0].item():
65 + first_token_id_index = self.input_batch.num_prompt_tokens[0].item()
66 + else:
67 + first_token_id_index = num_computed_tokens + 1
68 + if spec_decode_metadata is None:
69 + return SamplerOutput(
70 + sampled_token_ids=torch.tensor(
71 + [[prompt_token_ids[first_token_id_index]]],
72 + device="cuda",
73 + dtype=torch.int32,
74 + ),
75 + logprobs_tensors=None,
76 + )
77 + num_sampled_tokens = spec_decode_metadata.cu_num_sampled_tokens[0].item() + 1
78 + accpeted_tokens = prompt_token_ids[
79 + first_token_id_index : first_token_id_index
80 + + min(num_accepted_tokens, logits.shape[0])
81 + ]
82 + sampled_token_ids = accpeted_tokens + [-1] * (
83 + num_sampled_tokens - len(accpeted_tokens)
84 + )
85 + return SamplerOutput(
86 + sampled_token_ids=torch.tensor(
87 + [sampled_token_ids], device="cuda", dtype=torch.int32
88 + ),
89 + logprobs_tensors=None,
90 + )
91 +
92 + return fake_sample_fn
93 +
94 +
95 +def get_fake_propose_draft_token_ids_fn():
96 + def fake_propose_draft_token_ids_fn(
97 + self: GPUModelRunner,
98 + scheduler_output: SchedulerOutput,
99 + sampled_token_ids: torch.Tensor | list[list[int]],
100 + sampling_metadata: SamplingMetadata,
101 + hidden_states: torch.Tensor,
102 + sample_hidden_states: torch.Tensor,
103 + aux_hidden_states: list[torch.Tensor] | None,
104 + spec_decode_metadata: SpecDecodeMetadata | None,
105 + common_attn_metadata: CommonAttentionMetadata,
106 + ) -> list[list[int]]:
107 + num_computed_tokens_cpu_tensor = self.input_batch.num_computed_tokens_cpu_tensor
108 + num_computed_tokens = num_computed_tokens_cpu_tensor[0].item()
109 + if (
110 + self.input_batch.num_tokens_no_spec[0].item()
111 + <= self.input_batch.num_prompt_tokens[0].item()
112 + ):
113 + first_token_id_index = self.input_batch.num_prompt_tokens[0].item()
114 + else:
115 + first_token_id_index = (
116 + num_computed_tokens + 1
117 + ) # bonus token isn't considered as computed
118 + first_token_id_index += self.input_batch.num_accepted_tokens_cpu[0].item()
119 + proposed_draft_token_ids = [
120 + prompt_token_ids[
121 + first_token_id_index : first_token_id_index + num_speculative_tokens
122 + ]
123 + ]
124 + return proposed_draft_token_ids
125 +
126 + return fake_propose_draft_token_ids_fn
127 +
128 +
129 +def get_fake_step_action_fn(original_step_action_fn: Callable):
130 + def fake_get_output(self: InprocClient):
131 + global cur_step_action_idx
132 + global cur_step_action
133 + if cur_step_action_idx < len(step_actions):
134 + cur_step_action = step_actions[cur_step_action_idx]
135 + cur_step_action_idx += 1
136 + else:
137 + cur_step_action = None
138 + print(f"cur_step_action: {cur_step_action_idx=} {cur_step_action=}")
139 + return original_step_action_fn(self)
140 +
141 + return fake_get_output
142 +
143 +
144 +def get_fake_allocate_slots_fn(original_allocate_slots_fn: Callable):
145 + def fake_allocate_slots_fn(
146 + self: KVCacheManager,
147 + request: Request,
148 + num_new_tokens: int,
149 + num_new_computed_tokens: int = 0,
150 + new_computed_blocks: KVCacheBlocks | None = None,
151 + num_lookahead_tokens: int = 0,
152 + num_external_computed_tokens: int = 0,
153 + delay_cache_blocks: bool = False,
154 + num_encoder_tokens: int = 0,
155 + ):
156 + ret = original_allocate_slots_fn(
157 + self,
158 + request,
159 + num_new_tokens,
160 + num_new_computed_tokens,
161 + new_computed_blocks,
162 + num_lookahead_tokens,
163 + num_external_computed_tokens,
164 + delay_cache_blocks,
165 + num_encoder_tokens,
166 + )
167 + if cur_step_action is not None:
168 + cur_block_ids = self.coordinator.single_type_managers[0].req_to_blocks[
169 + request.request_id
170 + ]
171 + not_null_block_flags = [not block.is_null for block in cur_block_ids]
172 + block_ids = [1 if block else 0 for block in not_null_block_flags]
173 + assert block_ids == cur_step_action.kv_cache_block_ids
174 + return ret
175 +
176 + return fake_allocate_slots_fn
177 +
178 +
179 +mamba_kv_cache_dict = {}
180 +
181 +
182 +def get_fake_execute_model_fn(original_execute_model_fn: Callable):
183 + last_num_computed_tokens = 0
184 +
185 + def fake_execute_model_fn(
186 + self: GPUModelRunner,
187 + scheduler_output: SchedulerOutput,
188 + intermediate_tensors: IntermediateTensors | None = None,
189 + ):
190 + if cur_step_action is not None:
191 + num_scheduled_tokens = next(
192 + iter(scheduler_output.num_scheduled_tokens.values())
193 + )
194 + assert num_scheduled_tokens == cur_step_action.num_scheduled_tokens
195 + mamba_group_ids, mamba_spec = get_mamba_groups(self.kv_cache_config)
196 + mamba_group_id = mamba_group_ids[0]
197 + mamba_layer_name = self.kv_cache_config.kv_cache_groups[
198 + mamba_group_id
199 + ].layer_names[0]
200 + nonlocal last_num_computed_tokens
201 + if len(scheduler_output.scheduled_cached_reqs.req_ids) > 0:
202 + num_computed_tokens = (
203 + scheduler_output.scheduled_cached_reqs.num_computed_tokens[0]
204 + )
205 + if (
206 + num_computed_tokens // BLOCK_SIZE
207 + > last_num_computed_tokens // BLOCK_SIZE
208 + ):
209 + # generated a new aligned block in this step
210 + block_idx = num_computed_tokens // mamba_spec.block_size - 1
211 + block_id = (
212 + self.input_batch.block_table.block_tables[mamba_group_id]
213 + .block_table.cpu[0, block_idx]
214 + .item()
215 + )
216 + if block_id != 0:
217 + kv_cache = self.compilation_config.static_forward_context[
218 + mamba_layer_name
219 + ].kv_cache
220 + mamba_kv_cache_dict[
221 + num_computed_tokens - num_computed_tokens % BLOCK_SIZE
222 + ] = (
223 + kv_cache[0][0][block_id].clone(),
224 + kv_cache[0][1][block_id].clone(),
225 + )
226 +
227 + last_num_computed_tokens = num_computed_tokens
228 + else:
229 + last_num_computed_tokens = 0
230 +
231 + ret = original_execute_model_fn(self, scheduler_output, intermediate_tensors)
232 +
233 + if cur_step_action is not None:
234 + assert (
235 + cur_step_action.num_computed_tokens_start
236 + == self.input_batch.num_computed_tokens_cpu[0].item()
237 + )
238 +
239 + return ret
240 +
241 + return fake_execute_model_fn
242 +
243 +
244 +def get_fake_process_mamba_fn(
245 + original_preprocess_mamba_fn: Callable,
246 + original_post_process_mamba_fn: Callable,
247 + original_copy_fn: Callable,
248 +):
249 + copy_info: tuple[list[int], list[int], list[int]] | None = None
250 +
251 + def check_copy_info(
252 + action: tuple[int, int],
253 + kv_cache_config: KVCacheConfig,
254 + forward_context: dict[str, Any],
255 + input_batch: GPUInputBatch,
256 + ):
257 + assert copy_info is not None
258 + if action == (-1, -1):
259 + assert len(copy_info[0]) == len(copy_info[1]) == len(copy_info[2]) == 0
260 + else:
261 + assert len(copy_info[0]) == len(copy_info[1]) == len(copy_info[2]) == 2
262 + mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config)
263 + mamba_group_id = mamba_group_ids[0]
264 + mamba_layer_name = kv_cache_config.kv_cache_groups[
265 + mamba_group_id
266 + ].layer_names[0]
267 + mamba_kv_cache = forward_context[mamba_layer_name].kv_cache[0][-1]
268 + mamba_block_table = input_batch.block_table.block_tables[
269 + mamba_group_id
270 + ].block_table.cpu[0]
271 + expected_temporal_src = mamba_kv_cache[
272 + mamba_block_table[action[0]]
273 + ].data_ptr()
274 + expected_temporal_dest = mamba_kv_cache[
275 + mamba_block_table[action[1]]
276 + ].data_ptr()
277 + # -1 is qwen3-next's temporal. We skip checking conv as it is more complex.
278 + assert copy_info[0][-1] == expected_temporal_src
279 + assert copy_info[1][-1] == expected_temporal_dest
280 +
281 + def fake_preprocess_mamba_fn(
282 + scheduler_output: SchedulerOutput,
283 + kv_cache_config: KVCacheConfig,
284 + cache_config: CacheConfig,
285 + mamba_state_idx: dict[str, int],
286 + input_batch: GPUInputBatch,
287 + requests: dict[str, CachedRequestState],
288 + forward_context: dict[str, Any],
289 + mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
290 + ):
291 + nonlocal copy_info
292 + copy_info = None
293 + ret = original_preprocess_mamba_fn(
294 + scheduler_output,
295 + kv_cache_config,
296 + cache_config,
297 + mamba_state_idx,
298 + input_batch,
299 + requests,
300 + forward_context,
301 + mamba_state_copy_funcs,
302 + )
303 + if cur_step_action is not None:
304 + check_copy_info(
305 + cur_step_action.preprocess_copy_idx,
306 + kv_cache_config,
307 + forward_context,
308 + input_batch,
309 + )
310 + return ret
311 +
312 + def fake_post_process_mamba_fn(
313 + scheduler_output: SchedulerOutput,
314 + kv_cache_config: KVCacheConfig,
315 + input_batch: GPUInputBatch,
316 + requests: dict[str, CachedRequestState],
317 + mamba_state_idx: dict[str, int],
318 + forward_context: dict[str, Any],
319 + mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
320 + ):
321 + nonlocal copy_info
322 + copy_info = None
323 + ret = original_post_process_mamba_fn(
324 + scheduler_output,
325 + kv_cache_config,
326 + input_batch,
327 + requests,
328 + mamba_state_idx,
329 + forward_context,
330 + mamba_state_copy_funcs,
331 + )
332 + if cur_step_action is not None:
333 + check_copy_info(
334 + cur_step_action.postprocess_copy_idx,
335 + kv_cache_config,
336 + forward_context,
337 + input_batch,
338 + )
339 + return ret
340 +
341 + def fake_copy_fn(
342 + src_state_list: list[int],
343 + dest_state_list: list[int],
344 + num_elements_list: list[int],
345 + ):
346 + nonlocal copy_info
347 + assert copy_info is None
348 + copy_info = (src_state_list, dest_state_list, num_elements_list)
349 + return original_copy_fn(
350 + src_state_list,
351 + dest_state_list,
352 + num_elements_list,
353 + )
354 +
355 + return fake_preprocess_mamba_fn, fake_post_process_mamba_fn, fake_copy_fn
356 +
357 +
358 +def run_ref_mamba_state_in_subprocess() -> None:
359 + ctx = mp.get_context("spawn")
360 + proc = ctx.Process(target=_run_ref_mamba_state_worker)
361 + proc.start()
362 + proc.join(timeout=600)
363 + if proc.exitcode != 0:
364 + raise RuntimeError(f"Ref mamba state process exited with code {proc.exitcode}.")
365 +
366 +
367 +def _run_ref_mamba_state_worker():
368 + try:
369 + os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "0"
370 + num_generated_tokens = 8000
371 + num_prompt_tokens = 500
372 + sampling_params = SamplingParams(
373 + temperature=0.0, max_tokens=num_generated_tokens
374 + )
375 + prompt_dataset = datasets.load_dataset("heheda/a_long_article")
376 + full_prompt = prompt_dataset["train"][0]["text"]
377 + fake_execute_model_fn = get_fake_execute_model_fn(GPUModelRunner.execute_model)
378 + GPUModelRunner.execute_model = fake_execute_model_fn
379 + fake_sample_fn = get_fake_sample_fn()
380 + GPUModelRunner._sample = fake_sample_fn
381 + engine = LLM(
382 + model=MODEL,
383 + block_size=BLOCK_SIZE,
384 + hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS},
385 + seed=42,
386 + )
387 + global prompt_token_ids
388 + prompt_token_ids = engine.get_tokenizer().encode(full_prompt)
389 + print(f"Token IDs length: {len(prompt_token_ids)}")
390 +
391 + _outputs = engine.generate(
392 + [TokensPrompt(prompt_token_ids=prompt_token_ids[:num_prompt_tokens])],
393 + sampling_params,
394 + )
395 + # ref_mamba_kv_cache_dict = torch.load("mamba_kv_cache_dict.pth")
396 + # check_mamba_state_equal(ref_mamba_kv_cache_dict, mamba_kv_cache_dict)
397 + # torch.save(mamba_kv_cache_dict, "mamba_kv_cache_dict.pth")
398 + cpu_state_ref = {
399 + key: tuple(tensor.detach().cpu() for tensor in tensors)
400 + for key, tensors in mamba_kv_cache_dict.items()
401 + }
402 + torch.save(cpu_state_ref, "mamba_kv_cache_dict_ref.pth")
403 + mamba_kv_cache_dict.clear()
404 + except Exception:
405 + traceback.print_exc()
406 + raise
407 +
408 +
409 +def check_mamba_state_equal(
410 + mamba_state_ref: dict, mamba_state_new: dict, keys_to_check: list[int]
411 +):
412 + atol = 1e-2
413 + rtol = 1e-2
414 + for key in keys_to_check:
415 + assert key in mamba_state_new
416 + assert key in mamba_state_ref
417 + # mamba state new is a subset of mamba state ref
418 + for i, (ref, new) in enumerate(zip(mamba_state_ref[key], mamba_state_new[key])):
419 + if ref.device != new.device:
420 + new = new.to(ref.device)
421 + new = new[: ref.shape[0]]
422 + if not torch.allclose(ref, new, atol=atol, rtol=rtol):
423 + diff_mask = ~torch.isclose(ref, new, atol=atol, rtol=rtol)
424 + diff_idx = torch.nonzero(diff_mask)
425 + if diff_idx.shape[0] * 100 < ref.numel():
426 + print(
427 + f"[WARNING] found {diff_idx.shape[0] * 100 / ref.numel()}% of the elements are different" # noqa: E501
428 + )
429 + continue
430 + raise ValueError(
431 + f"Mamba state is not equal for key: {key} at index {i}"
432 + )
433 + return True
434 +
435 +
436 +@dataclass
437 +class TestConfig:
438 + num_prompt_tokens: int
439 + num_generated_tokens: int
440 + num_accepted_tokens: int
441 + step_actions: list[StepAction]
442 +
443 +
444 +def apply_patch(monkeypatch: pytest.MonkeyPatch):
445 + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
446 +
447 + fake_sample_fn = get_fake_sample_fn()
448 + monkeypatch.setattr(GPUModelRunner, "_sample", fake_sample_fn)
449 +
450 + fake_propose_draft_token_ids_fn = get_fake_propose_draft_token_ids_fn()
451 + monkeypatch.setattr(
452 + GPUModelRunner, "propose_draft_token_ids", fake_propose_draft_token_ids_fn
453 + )
454 +
455 + fake_execute_model_fn = get_fake_execute_model_fn(GPUModelRunner.execute_model)
456 + monkeypatch.setattr(GPUModelRunner, "execute_model", fake_execute_model_fn)
457 +
458 + fake_step_action_fn = get_fake_step_action_fn(InprocClient.get_output)
459 + monkeypatch.setattr(InprocClient, "get_output", fake_step_action_fn)
460 +
461 + fake_allocate_slots_fn = get_fake_allocate_slots_fn(KVCacheManager.allocate_slots)
462 + monkeypatch.setattr(KVCacheManager, "allocate_slots", fake_allocate_slots_fn)
463 +
464 + fake_preprocess_mamba_fn, fake_post_process_mamba_fn, fake_copy_fn = (
465 + get_fake_process_mamba_fn(
466 + mamba_utils.preprocess_mamba,
467 + mamba_utils.postprocess_mamba,
468 + mamba_utils.do_mamba_copy_block,
469 + )
470 + )
471 + monkeypatch.setattr(mamba_utils, "preprocess_mamba", fake_preprocess_mamba_fn)
472 + monkeypatch.setattr(mamba_utils, "postprocess_mamba", fake_post_process_mamba_fn)
473 + monkeypatch.setattr(mamba_utils, "do_mamba_copy_block", fake_copy_fn)
474 +
475 +
476 +@pytest.mark.skip(
477 + reason="Skipping test_mamba_prefix_cache because it is based on spec "
478 + "decode which is not allowed now."
479 +)
480 +def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
481 + run_ref_mamba_state_in_subprocess()
482 + apply_patch(monkeypatch)
483 + prompt_dataset = datasets.load_dataset("heheda/a_long_article")
484 + full_prompt = prompt_dataset["train"][0]["text"]
485 + tests = {
486 + "accept_1": TestConfig(
487 + num_prompt_tokens=554,
488 + num_generated_tokens=20,
489 + num_accepted_tokens=1,
490 + step_actions=[
491 + StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
492 + StepAction(554, 4, [], (-1, -1), (-1, -1)),
493 + StepAction(555, 4, [], (-1, -1), (-1, -1)),
494 + StepAction(556, 4, [], (-1, -1), (-1, -1)),
495 + StepAction(557, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
496 + StepAction(558, 4, [], (-1, -1), (-1, -1)),
497 + StepAction(559, 4, [], (-1, -1), (1, 0)),
498 + StepAction(560, 4, [], (-1, -1), (-1, -1)),
499 + StepAction(561, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
500 + ],
501 + ),
502 + # test case 2.1: no hit, accept 2 tokens
503 + "accept_2_1": TestConfig(
504 + num_prompt_tokens=554,
505 + num_generated_tokens=20,
506 + num_accepted_tokens=2,
507 + step_actions=[
508 + StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
509 + StepAction(554, 4, [], (-1, -1), (-1, -1)),
510 + StepAction(556, 4, [], (-1, -1), (-1, -1)),
511 + StepAction(558, 4, [1, 1, 1, 1, 1], (1, 1), (2, 0)),
512 + StepAction(560, 4, [], (-1, -1), (-1, -1)),
513 + StepAction(562, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
514 + ],
515 + ),
516 + # test case 2.2: no hit, accept 2 tokens
517 + "accept_2_2": TestConfig(
518 + num_prompt_tokens=555,
519 + num_generated_tokens=20,
520 + num_accepted_tokens=2,
521 + step_actions=[
522 + StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
523 + StepAction(555, 4, [], (-1, -1), (-1, -1)),
524 + StepAction(557, 4, [1, 1, 1, 1, 1], (1, 1), (-1, -1)),
525 + StepAction(559, 4, [], (-1, -1), (1, 0)),
526 + StepAction(561, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
527 + ],
528 + ),
529 + "accept_3_1": TestConfig(
530 + num_prompt_tokens=553,
531 + num_generated_tokens=20,
532 + num_accepted_tokens=3,
533 + step_actions=[
534 + StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)),
535 + StepAction(553, 4, [], (-1, -1), (-1, -1)),
536 + StepAction(556, 4, [], (-1, -1), (-1, -1)),
537 + StepAction(559, 4, [1, 1, 1, 1, 1], (2, 1), (1, 0)),
538 + StepAction(562, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
539 + ],
540 + ),
541 + "accept_3_2": TestConfig(
542 + num_prompt_tokens=554,
543 + num_generated_tokens=20,
544 + num_accepted_tokens=3,
545 + step_actions=[
546 + StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
547 + StepAction(554, 4, [], (-1, -1), (-1, -1)),
548 + StepAction(557, 4, [1, 1, 1, 1, 1], (2, 1), (3, 0)),
549 + StepAction(560, 4, [], (-1, -1), (-1, -1)),
550 + StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
551 + ],
552 + ),
553 + "accept_3_3": TestConfig(
554 + num_prompt_tokens=555,
555 + num_generated_tokens=20,
556 + num_accepted_tokens=3,
557 + step_actions=[
558 + StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
559 + StepAction(555, 4, [], (-1, -1), (-1, -1)),
560 + StepAction(558, 4, [1, 1, 1, 1, 1], (2, 1), (2, 0)),
561 + StepAction(561, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
562 + ],
563 + ),
564 + "accept_4_1": TestConfig(
565 + num_prompt_tokens=553,
566 + num_generated_tokens=20,
567 + num_accepted_tokens=4,
568 + step_actions=[
569 + StepAction(0, 553, [1, 1, 1, 1], (-1, -1), (-1, -1)),
570 + StepAction(553, 4, [], (-1, -1), (-1, -1)),
571 + StepAction(557, 4, [1, 1, 1, 1, 1], (3, 1), (3, 0)),
572 + StepAction(561, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
573 + StepAction(565, 4, [], (-1, -1), (-1, -1)),
574 + ],
575 + ),
576 + "accept_4_2": TestConfig(
577 + num_prompt_tokens=554,
578 + num_generated_tokens=25,
579 + num_accepted_tokens=4,
580 + step_actions=[
581 + StepAction(0, 554, [1, 1, 1, 1], (-1, -1), (-1, -1)),
582 + StepAction(554, 4, [], (-1, -1), (-1, -1)),
583 + StepAction(558, 4, [1, 1, 1, 1, 1], (3, 1), (2, 0)),
584 + StepAction(562, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
585 + StepAction(566, 4, [], (-1, -1), (-1, -1)),
586 + ],
587 + ),
588 + "accept_4_3": TestConfig(
589 + num_prompt_tokens=555,
590 + num_generated_tokens=25,
591 + num_accepted_tokens=4,
592 + step_actions=[
593 + StepAction(0, 555, [1, 1, 1, 1], (-1, -1), (-1, -1)),
594 + StepAction(555, 4, [], (-1, -1), (-1, -1)),
595 + StepAction(559, 4, [1, 1, 1, 1, 1], (3, 1), (1, 0)),
596 + StepAction(563, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
597 + ],
598 + ),
599 + "accept_4_4": TestConfig(
600 + num_prompt_tokens=556,
601 + num_generated_tokens=25,
602 + num_accepted_tokens=4,
603 + step_actions=[
604 + StepAction(0, 556, [1, 1, 1, 1], (-1, -1), (-1, -1)),
605 + StepAction(556, 4, [], (-1, -1), (3, 0)),
606 + StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
607 + StepAction(564, 4, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
608 + ],
609 + ),
610 + "prompt_block_size": TestConfig(
611 + num_prompt_tokens=560,
612 + num_generated_tokens=10,
613 + num_accepted_tokens=4,
614 + step_actions=[
615 + StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
616 + StepAction(560, 4, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
617 + ],
618 + ),
619 + "prompt_2_block_size": TestConfig(
620 + num_prompt_tokens=560 * 2,
621 + num_generated_tokens=10,
622 + num_accepted_tokens=4,
623 + step_actions=[
624 + StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
625 + StepAction(560, 560, [1, 1, 1, 1, 1], (0, 1), (-1, -1)),
626 + StepAction(560 * 2, 4, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
627 + ],
628 + ),
629 + "prompt_2_block_size_10": TestConfig(
630 + num_prompt_tokens=560 * 2 + 10,
631 + num_generated_tokens=10,
632 + num_accepted_tokens=4,
633 + step_actions=[
634 + StepAction(0, 560, [1, 1, 1, 1], (-1, -1), (-1, -1)),
635 + StepAction(560, 570, [1, 0, 1, 1, 1, 1], (0, 2), (-1, -1)),
636 + StepAction(560 * 2 + 10, 4, [0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
637 + ],
638 + ),
639 + "prompt_3_block_size": TestConfig(
640 + num_prompt_tokens=560 * 3,
641 + num_generated_tokens=10,
642 + num_accepted_tokens=4,
643 + step_actions=[
644 + StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
645 + StepAction(560 * 2, 560, [0, 1, 1, 1, 1, 1], (1, 2), (-1, -1)),
646 + StepAction(560 * 3, 4, [0, 0, 1, 1, 1, 1, 1], (2, 3), (-1, -1)),
647 + ],
648 + ),
649 + "prompt_3_block_size_10": TestConfig(
650 + num_prompt_tokens=560 * 3 + 10,
651 + num_generated_tokens=10,
652 + num_accepted_tokens=4,
653 + step_actions=[
654 + StepAction(0, 560 * 2, [0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
655 + StepAction(560 * 2, 570, [0, 1, 0, 1, 1, 1, 1], (1, 3), (-1, -1)),
656 + StepAction(560 * 3 + 10, 4, [0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
657 + ],
658 + ),
659 + "prompt_10_block_size": TestConfig(
660 + num_prompt_tokens=560 * 10,
661 + num_generated_tokens=10,
662 + num_accepted_tokens=4,
663 + step_actions=[
664 + StepAction(0, 560 * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
665 + StepAction(
666 + 560 * 5,
667 + 560 * 4,
668 + [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1],
669 + (4, 8),
670 + (-1, -1),
671 + ),
672 + StepAction(
673 + 560 * 9,
674 + 560,
675 + [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
676 + (8, 9),
677 + (-1, -1),
678 + ),
679 + StepAction(
680 + 560 * 10,
681 + 4,
682 + [0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1],
683 + (9, 10),
684 + (-1, -1),
685 + ),
686 + ],
687 + ),
688 + "prompt_10_block_size_10": TestConfig(
689 + num_prompt_tokens=560 * 10 + 10,
690 + num_generated_tokens=10,
691 + num_accepted_tokens=4,
692 + step_actions=[
693 + StepAction(0, 560 * 5, [0, 0, 0, 0, 1, 1, 1, 1], (-1, -1), (-1, -1)),
694 + StepAction(
695 + 560 * 5,
696 + 560 * 4,
697 + [0, 0, 0, 0, 1, 0, 0, 0, 1, 1, 1, 1],
698 + (4, 8),
699 + (-1, -1),
700 + ),
701 + StepAction(
702 + 560 * 9,
703 + 560 + 10,
704 + [0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 1, 1, 1, 1],
705 + (8, 10),
706 + (-1, -1),
707 + ),
708 + ],
709 + ),
710 + }
711 +
712 + engine = LLM(
713 + model=MODEL,
714 + enable_prefix_caching=True,
715 + block_size=BLOCK_SIZE,
716 + mamba_cache_mode="align",
717 + speculative_config={
718 + "method": "qwen3_next_mtp",
719 + "num_speculative_tokens": num_speculative_tokens,
720 + },
721 + max_num_batched_tokens=3072,
722 + hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS},
723 + seed=42,
724 + )
725 + global prompt_token_ids
726 + prompt_token_ids = engine.get_tokenizer().encode(full_prompt)
727 + print(f"Token IDs length: {len(prompt_token_ids)}")
728 + for test_case_name, test_config in tests.items():
729 + print(f"Running test case: {test_case_name}")
730 + num_generated_tokens = test_config.num_generated_tokens
731 + num_prompt_tokens = test_config.num_prompt_tokens
732 + global num_accepted_tokens
733 + num_accepted_tokens = test_config.num_accepted_tokens
734 + sampling_params = SamplingParams(
735 + temperature=0.0, max_tokens=num_generated_tokens
736 + )
737 + global cur_step_action_idx
738 + cur_step_action_idx = 0
739 + for step_action_prev, step_action_next in zip(
740 + test_config.step_actions[:-1], test_config.step_actions[1:]
741 + ):
742 + if (
743 + step_action_next.kv_cache_block_ids is not None
744 + and len(step_action_next.kv_cache_block_ids) == 0
745 + ):
746 + prev_block_ids = step_action_prev.kv_cache_block_ids
747 + if prev_block_ids is not None:
748 + step_action_next.kv_cache_block_ids = prev_block_ids.copy()
749 + global step_actions
750 + step_actions = test_config.step_actions
751 + _ = engine.generate(
752 + [TokensPrompt(prompt_token_ids=prompt_token_ids[:num_prompt_tokens])],
753 + sampling_params,
754 + )
755 + assert engine.llm_engine.engine_core.engine_core.scheduler.reset_prefix_cache()
756 + print(f"End test case: {test_case_name}")
757 + keys_to_check = [
758 + (action.postprocess_copy_idx[1] + 1) * BLOCK_SIZE
759 + for action in test_config.step_actions
760 + if action.postprocess_copy_idx and action.postprocess_copy_idx[0] != -1
761 + ]
762 + mamba_state_ref = torch.load("mamba_kv_cache_dict_ref.pth")
763 + check_mamba_state_equal(mamba_state_ref, mamba_kv_cache_dict, keys_to_check)
764 + mamba_kv_cache_dict.clear()
vllm/config/cache.py
+10
@@ -31,6 +31,7 @@ CacheDType = Literal[
31 "fp8_ds_mla",
32 ]
33 MambaDType = Literal["auto", "float32", "float16"]
34 +MambaCacheMode = Literal["all", "align", "none"]
35 PrefixCachingHashAlgo = Literal["sha256", "sha256_cbor", "xxhash", "xxhash_cbor"]
36 KVOffloadingBackend = Literal["native", "lmcache"]
37
@@ -123,6 +124,15 @@ class CacheConfig:
124 """The data type to use for the Mamba cache (ssm state only, conv state will
125 still be controlled by mamba_cache_dtype). If set to 'auto', the data type
126 for the ssm state will be determined by mamba_cache_dtype."""
127 + mamba_cache_mode: MambaCacheMode = "none"
128 + """The cache strategy for Mamba layers.
129 + - "none": set when prefix caching is disabled.
130 + - "all": cache the mamba state of all tokens at position i * block_size. This is
131 + the default behavior (for models that support it) when prefix caching is
132 + enabled.
133 + - "align": only cache the mamba state of the last token of each scheduler step and
134 + when the token is at position i * block_size.
135 + """
136
137 # Will be set after profiling.
138 num_gpu_blocks: int | None = field(default=None, init=False)
vllm/config/vllm.py
+11
@@ -999,6 +999,17 @@ class VllmConfig:
999 # Default to enable HMA if not explicitly disabled by user or logic above.
1000 self.scheduler_config.disable_hybrid_kv_cache_manager = False
1001
1002 + if self.cache_config.mamba_cache_mode == "align":
1003 + if self.scheduler_config.long_prefill_token_threshold > 0:
1004 + assert (
1005 + self.scheduler_config.long_prefill_token_threshold
1006 + >= self.cache_config.block_size
1007 + )
1008 + assert not self.scheduler_config.disable_chunked_mm_input, (
1009 + "Chunked MM input is required because we need the flexibility to "
1010 + "schedule a multiple of block_size tokens even if they are in the "
1011 + "middle of a mm input"
1012 + )
1013 if self.compilation_config.debug_dump_path:
1014 self.compilation_config.debug_dump_path = (
1015 self.compilation_config.debug_dump_path.absolute().expanduser()
vllm/engine/arg_utils.py
+6
@@ -60,6 +60,7 @@ from vllm.config.cache import (
60 BlockSize,
61 CacheDType,
62 KVOffloadingBackend,
63 + MambaCacheMode,
64 MambaDType,
65 PrefixCachingHashAlgo,
66 )
@@ -556,6 +557,7 @@ class EngineArgs:
557 mamba_cache_dtype: MambaDType = CacheConfig.mamba_cache_dtype
558 mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype
559 mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size")
560 + mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode
561
562 additional_config: dict[str, Any] = get_field(VllmConfig, "additional_config")
563
@@ -939,6 +941,9 @@ class EngineArgs:
941 cache_group.add_argument(
942 "--mamba-block-size", **cache_kwargs["mamba_block_size"]
943 )
944 + cache_group.add_argument(
945 + "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"]
946 + )
947 cache_group.add_argument(
948 "--kv-offloading-size", **cache_kwargs["kv_offloading_size"]
949 )
@@ -1416,6 +1421,7 @@ class EngineArgs:
1421 mamba_cache_dtype=self.mamba_cache_dtype,
1422 mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype,
1423 mamba_block_size=self.mamba_block_size,
1424 + mamba_cache_mode=self.mamba_cache_mode,
1425 kv_offloading_size=self.kv_offloading_size,
1426 kv_offloading_backend=self.kv_offloading_backend,
1427 )
vllm/model_executor/layers/mamba/abstract.py
+1
@@ -56,6 +56,7 @@ class MambaBase(AttentionLayerBase):
56 block_size=mamba_block_size,
57 page_size_padded=page_size_padded,
58 mamba_type=self.mamba_type,
59 + mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode,
60 num_speculative_blocks=(
61 vllm_config.speculative_config.num_speculative_tokens
62 if vllm_config.speculative_config
vllm/model_executor/layers/mamba/mamba_mixer.py
+3 -3
@@ -255,7 +255,7 @@ class MambaMixer(MambaBase, CustomOp):
255
256 assert self.cache_config is not None
257 mamba_block_size = self.cache_config.mamba_block_size
258 - prefix_caching_enabled = self.cache_config.enable_prefix_caching
258 + is_mamba_cache_all = self.cache_config.mamba_cache_mode == "all"
259
260 if attn_metadata is not None:
261 assert isinstance(attn_metadata, dict)
@@ -304,7 +304,7 @@ class MambaMixer(MambaBase, CustomOp):
304 state_indices_tensor_p = prefill_decode_split.state_indices_tensor_p
305 state_indices_tensor_d = prefill_decode_split.state_indices_tensor_d
306
307 - if prefix_caching_enabled:
307 + if is_mamba_cache_all:
308 block_idx_last_computed_token_d, block_idx_last_computed_token_p = (
309 torch.split(
310 attn_metadata.block_idx_last_computed_token,
@@ -380,7 +380,7 @@ class MambaMixer(MambaBase, CustomOp):
380 ssm_outputs.append(scan_out_p)
381
382 if has_decode:
383 - if prefix_caching_enabled:
383 + if is_mamba_cache_all:
384 state_indices_tensor_d_input = state_indices_tensor_d.gather(
385 1, block_idx_last_computed_token_d.unsqueeze(1)
386 ).squeeze(1)
vllm/model_executor/layers/mamba/mamba_mixer2.py
+6 -6
@@ -570,7 +570,7 @@ class MambaMixer2(MambaBase, CustomOp):
570
571 assert self.cache_config is not None
572 mamba_block_size = self.cache_config.mamba_block_size
573 - prefix_caching_enabled = self.cache_config.enable_prefix_caching
573 + is_mamba_cache_all = self.cache_config.mamba_cache_mode == "all"
574 if attn_metadata is not None:
575 assert isinstance(attn_metadata, dict)
576 attn_metadata = attn_metadata[self.prefix]
@@ -622,7 +622,7 @@ class MambaMixer2(MambaBase, CustomOp):
622 dim=0,
623 )
624
625 - if prefix_caching_enabled:
625 + if is_mamba_cache_all:
626 # If prefix caching is enabled, retrieve the relevant variables
627 # for prefill and decode
628 block_idx_last_computed_token_d, block_idx_last_computed_token_p = (
@@ -701,7 +701,7 @@ class MambaMixer2(MambaBase, CustomOp):
701 initial_states = None
702 if has_initial_states_p is not None and prep_initial_states:
703 kernel_ssm_indices = state_indices_tensor_p
704 - if prefix_caching_enabled:
704 + if is_mamba_cache_all:
705 kernel_ssm_indices = state_indices_tensor_p.gather(
706 1, block_idx_last_computed_token_p.unsqueeze(1)
707 ).squeeze(1)
@@ -729,14 +729,14 @@ class MambaMixer2(MambaBase, CustomOp):
729 cu_chunk_seqlens=cu_chunk_seqlen_p,
730 last_chunk_indices=last_chunk_indices_p,
731 initial_states=initial_states,
732 - return_intermediate_states=prefix_caching_enabled,
732 + return_intermediate_states=is_mamba_cache_all,
733 dt_softplus=True,
734 dt_limit=(0.0, float("inf")),
735 out=preallocated_ssm_out_p.view(num_prefill_tokens, -1, self.head_dim),
736 state_dtype=ssm_state.dtype,
737 )
738
739 - if prefix_caching_enabled:
739 + if is_mamba_cache_all:
740 # The chunk_stride is the number of chunks per mamba block
741 # e.g., if mamba_block_size = 512 and chunk_size = 256,
742 # then chunk_stride = 2
@@ -815,7 +815,7 @@ class MambaMixer2(MambaBase, CustomOp):
815
816 # Process decode requests
817 if has_decode:
818 - if prefix_caching_enabled:
818 + if is_mamba_cache_all:
819 state_indices_tensor_d_input = state_indices_tensor_d.gather(
820 1, block_idx_last_computed_token_d.unsqueeze(1)
821 ).squeeze(1)
vllm/model_executor/layers/mamba/mamba_utils.py
+95
@@ -1,6 +1,10 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from collections.abc import Callable
5 +from dataclasses import dataclass
6 +from typing import TypeAlias
7 +
8 import torch
9
10 from vllm.config.cache import MambaDType
@@ -223,3 +227,94 @@ class MambaStateShapeCalculator:
227 conv_state_k_shape,
228 recurrent_state_shape,
229 )
230 +
231 +
232 +@dataclass
233 +class MambaCopySpec:
234 + """
235 + Data class specifying the memory-copy parameters for Mamba states used for
236 + prefix caching in align mode.
237 +
238 + Attributes:
239 + start_addr (int): Starting address for the memory copy operation.
240 + num_elements (int): Number of elements to copy from the starting address.
241 + """
242 +
243 + start_addr: int
244 + num_elements: int
245 +
246 +
247 +MambaStateCopyFunc: TypeAlias = Callable[
248 + [torch.Tensor, list[int], int, int], MambaCopySpec
249 +]
250 +"""
251 +Type alias for a function that computes a MambaCopySpec for copying state slices.
252 +Parameters:
253 + state: torch.Tensor - the Mamba state tensor (e.g., conv or temporal states).
254 + block_ids: list[int] - the list of block indices for the state to copy.
255 + cur_block_idx: int - current block index within `block_ids` to copy from.
256 + num_accepted_tokens: int - number of accepted tokens used to compute the copy offset.
257 + Range: 1 .. 1 + num_speculative_tokens (inclusive).
258 +"""
259 +
260 +
261 +def get_conv_copy_spec(
262 + state: torch.Tensor,
263 + block_ids: list[int],
264 + cur_block_idx: int,
265 + num_accepted_tokens: int,
266 +) -> MambaCopySpec:
267 + """Return a MambaCopySpec for copying a convolutional state slice."""
268 + src_block_id = block_ids[cur_block_idx]
269 + src_state = state[src_block_id, num_accepted_tokens - 1 :]
270 + return MambaCopySpec(
271 + start_addr=src_state.data_ptr(), num_elements=src_state.numel()
272 + )
273 +
274 +
275 +def get_temporal_copy_spec(
276 + state: torch.Tensor,
277 + block_ids: list[int],
278 + cur_block_idx: int,
279 + num_accepted_tokens: int,
280 +) -> MambaCopySpec:
281 + """Return a MambaCopySpec for copying a temporal state slice."""
282 + src_block_id = block_ids[cur_block_idx + num_accepted_tokens - 1]
283 + src_state = state[src_block_id]
284 + return MambaCopySpec(
285 + start_addr=src_state.data_ptr(), num_elements=src_state.numel()
286 + )
287 +
288 +
289 +get_full_copy_spec = get_temporal_copy_spec
290 +
291 +
292 +class MambaStateCopyFuncCalculator:
293 + @classmethod
294 + def linear_attention_state_copy_func(cls):
295 + return (get_temporal_copy_spec,)
296 +
297 + @classmethod
298 + def mamba1_state_copy_func(cls):
299 + return (get_conv_copy_spec, get_temporal_copy_spec)
300 +
301 + @classmethod
302 + def mamba2_state_copy_func(cls):
303 + return get_conv_copy_spec, get_temporal_copy_spec
304 +
305 + @classmethod
306 + def short_conv_state_copy_func(cls):
307 + return (get_conv_copy_spec,)
308 +
309 + @classmethod
310 + def gated_delta_net_state_copy_func(cls):
311 + return (get_conv_copy_spec, get_temporal_copy_spec)
312 +
313 + @classmethod
314 + def kda_state_copy_func(cls):
315 + return (
316 + get_conv_copy_spec,
317 + get_conv_copy_spec,
318 + get_conv_copy_spec,
319 + get_temporal_copy_spec,
320 + )
vllm/model_executor/models/bamba.py
+6
@@ -24,6 +24,8 @@ from vllm.model_executor.layers.linear import (
24 from vllm.model_executor.layers.logits_processor import LogitsProcessor
25 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
26 from vllm.model_executor.layers.mamba.mamba_utils import (
27 + MambaStateCopyFunc,
28 + MambaStateCopyFuncCalculator,
29 MambaStateDtypeCalculator,
30 MambaStateShapeCalculator,
31 )
@@ -455,6 +457,10 @@ class BambaForCausalLM(
457 conv_kernel=hf_config.mamba_d_conv,
458 )
459
460 + @classmethod
461 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
462 + return MambaStateCopyFuncCalculator.mamba2_state_copy_func()
463 +
464 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
465 config = vllm_config.model_config.hf_config
466 self.vllm_config = vllm_config
vllm/model_executor/models/config.py
+55 -20
@@ -330,26 +330,54 @@ class MambaModelConfig(VerifyAndUpdateConfig):
330 cache_config = vllm_config.cache_config
331
332 if cache_config.enable_prefix_caching:
333 - if model_config.supports_mamba_prefix_caching:
334 - logger.info(
335 - "Warning: Prefix caching is currently enabled. "
336 - "Its support for Mamba layers is experimental. "
337 - "Please report any issues you may observe."
333 + if cache_config.mamba_cache_mode == "none":
334 + cache_config.mamba_cache_mode = (
335 + "all" if model_config.supports_mamba_prefix_caching else "align"
336 )
339 - # By default, mamba block size will be set to max_model_len (see
340 - # below). When enabling prefix caching, we align mamba block size
341 - # to the block size as the basic granularity for prefix caching.
342 - if cache_config.mamba_block_size is None:
343 - cache_config.mamba_block_size = cache_config.block_size
344 - else:
345 - logger.info(
346 - "Hybrid or mamba-based model detected without "
347 - "support for prefix caching: disabling."
337 + logger.warning(
338 + "Mamba cache mode is set to '%s' for %s by default "
339 + "when prefix caching is enabled",
340 + cache_config.mamba_cache_mode,
341 + model_config.architecture,
342 )
349 - cache_config.enable_prefix_caching = False
350 -
351 - if cache_config.mamba_block_size is None:
352 - cache_config.mamba_block_size = model_config.max_model_len
343 + if (
344 + cache_config.mamba_cache_mode == "all"
345 + and not model_config.supports_mamba_prefix_caching
346 + ):
347 + cache_config.mamba_cache_mode = "align"
348 + logger.warning(
349 + "Hybrid or mamba-based model detected without support "
350 + "for prefix caching with Mamba cache 'all' mode: "
351 + "falling back to 'align' mode."
352 + )
353 + if cache_config.mamba_cache_mode == "align":
354 + assert vllm_config.scheduler_config.enable_chunked_prefill, (
355 + "Chunked prefill is required for mamba cache mode 'align'."
356 + )
357 + assert not vllm_config.speculative_config, (
358 + "Mamba cache mode 'align' is currently not compatible "
359 + "with speculative decoding."
360 + )
361 + logger.info(
362 + "Warning: Prefix caching in Mamba cache '%s' "
363 + "mode is currently enabled. "
364 + "Its support for Mamba layers is experimental. "
365 + "Please report any issues you may observe.",
366 + cache_config.mamba_cache_mode,
367 + )
368 + # By default, mamba block size will be set to max_model_len (see
369 + # below). When enabling prefix caching, we align mamba block size
370 + # to the block size as the basic granularity for prefix caching.
371 + if cache_config.mamba_block_size is None:
372 + cache_config.mamba_block_size = cache_config.block_size
373 + else:
374 + if cache_config.mamba_cache_mode != "none":
375 + cache_config.mamba_cache_mode = "none"
376 + logger.warning(
377 + "Mamba cache mode is set to 'none' when prefix caching is disabled"
378 + )
379 + if cache_config.mamba_block_size is None:
380 + cache_config.mamba_block_size = model_config.max_model_len
381
382
383 class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
@@ -426,7 +454,7 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
454 mamba_page_size = MambaSpec(
455 shapes=model_cls.get_mamba_state_shape_from_config(vllm_config),
456 dtypes=model_cls.get_mamba_state_dtype_from_config(vllm_config),
429 - block_size=model_config.max_model_len,
457 + block_size=-1, # block_size doesn't matter for mamba page size
458 ).page_size_bytes
459
460 # Model may be marked as is_hybrid
@@ -435,7 +463,7 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
463 if mamba_page_size == 0:
464 return
465
438 - if cache_config.enable_prefix_caching:
466 + if cache_config.mamba_cache_mode == "all":
467 # With prefix caching, select attention block size to
468 # optimize for mamba kernel performance
469
@@ -479,6 +507,13 @@ class HybridAttentionMambaModelConfig(VerifyAndUpdateConfig):
507 attn_block_size,
508 )
509
510 + # By default, mamba block size will be set to max_model_len.
511 + # When enabling prefix caching and using align mamba cache
512 + # mode, we align mamba block size to the block size as the
513 + # basic granularity for prefix caching.
514 + if cache_config.mamba_cache_mode == "align":
515 + cache_config.mamba_block_size = cache_config.block_size
516 +
517 # compute new attention page size
518 attn_page_size = cache_config.block_size * attn_page_size_1_token
519
vllm/model_executor/models/falcon_h1.py
+6
@@ -24,6 +24,8 @@ from vllm.model_executor.layers.linear import (
24 from vllm.model_executor.layers.logits_processor import LogitsProcessor
25 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
26 from vllm.model_executor.layers.mamba.mamba_utils import (
27 + MambaStateCopyFunc,
28 + MambaStateCopyFuncCalculator,
29 MambaStateDtypeCalculator,
30 MambaStateShapeCalculator,
31 )
@@ -551,6 +553,10 @@ class FalconH1ForCausalLM(
553 conv_kernel=hf_config.mamba_d_conv,
554 )
555
556 + @classmethod
557 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
558 + return MambaStateCopyFuncCalculator.mamba2_state_copy_func()
559 +
560 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
561 config = vllm_config.model_config.hf_config
562 self.vllm_config = vllm_config
vllm/model_executor/models/granitemoehybrid.py
+6
@@ -19,6 +19,8 @@ from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLine
19 from vllm.model_executor.layers.logits_processor import LogitsProcessor
20 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
21 from vllm.model_executor.layers.mamba.mamba_utils import (
22 + MambaStateCopyFunc,
23 + MambaStateCopyFuncCalculator,
24 MambaStateDtypeCalculator,
25 MambaStateShapeCalculator,
26 )
@@ -641,6 +643,10 @@ class GraniteMoeHybridForCausalLM(
643 conv_kernel=hf_config.mamba_d_conv,
644 )
645
646 + @classmethod
647 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
648 + return MambaStateCopyFuncCalculator.mamba2_state_copy_func()
649 +
650 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
651 super().__init__()
652
vllm/model_executor/models/interfaces.py
+14
@@ -24,6 +24,7 @@ from vllm.config import ModelConfig, SpeechToTextConfig
24 from vllm.inputs import TokensPrompt
25 from vllm.inputs.data import PromptType
26 from vllm.logger import init_logger
27 +from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc
28 from vllm.model_executor.layers.quantization import QuantizationConfig
29 from vllm.utils.collection_utils import common_prefix
30 from vllm.utils.func_utils import supports_kw
@@ -776,6 +777,19 @@ class IsHybrid(Protocol):
777 """
778 ...
779
780 + @classmethod
781 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, ...]:
782 + """Calculate copy-function callables for each Mamba state.
783 +
784 + Returns:
785 + A tuple of MambaStateCopyFunc callables that correspond, in order,
786 + to the Mamba states produced by the model. Each callable accepts
787 + (state, block_ids, cur_block_idx, num_accepted_tokens) and returns
788 + a MambaCopySpec describing the memory-copy parameters for prefix
789 + caching in align mode.
790 + """
791 + ...
792 +
793
794 @overload
795 def is_hybrid(model: object) -> TypeIs[IsHybrid]: ...
vllm/model_executor/models/jamba.py
+6
@@ -24,6 +24,8 @@ from vllm.model_executor.layers.linear import (
24 from vllm.model_executor.layers.logits_processor import LogitsProcessor
25 from vllm.model_executor.layers.mamba.mamba_mixer import MambaMixer
26 from vllm.model_executor.layers.mamba.mamba_utils import (
27 + MambaStateCopyFunc,
28 + MambaStateCopyFuncCalculator,
29 MambaStateDtypeCalculator,
30 MambaStateShapeCalculator,
31 )
@@ -558,6 +560,10 @@ class JambaForCausalLM(
560 conv_kernel=hf_config.mamba_d_conv,
561 )
562
563 + @classmethod
564 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
565 + return MambaStateCopyFuncCalculator.mamba1_state_copy_func()
566 +
567 def compute_logits(
568 self,
569 hidden_states: torch.Tensor,
vllm/model_executor/models/kimi_linear.py
+10
@@ -26,6 +26,8 @@ from vllm.model_executor.layers.linear import (
26 )
27 from vllm.model_executor.layers.logits_processor import LogitsProcessor
28 from vllm.model_executor.layers.mamba.mamba_utils import (
29 + MambaStateCopyFunc,
30 + MambaStateCopyFuncCalculator,
31 MambaStateDtypeCalculator,
32 MambaStateShapeCalculator,
33 )
@@ -544,6 +546,14 @@ class KimiLinearForCausalLM(
546 num_spec=num_spec,
547 )
548
549 + @classmethod
550 + def get_mamba_state_copy_func(
551 + cls,
552 + ) -> tuple[
553 + MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc, MambaStateCopyFunc
554 + ]:
555 + return MambaStateCopyFuncCalculator.kda_state_copy_func()
556 +
557 def compute_logits(
558 self,
559 hidden_states: torch.Tensor,
vllm/model_executor/models/lfm2.py
+11 -4
@@ -20,6 +20,8 @@ from vllm.model_executor.layers.linear import (
20 )
21 from vllm.model_executor.layers.logits_processor import LogitsProcessor
22 from vllm.model_executor.layers.mamba.mamba_utils import (
23 + MambaStateCopyFunc,
24 + MambaStateCopyFuncCalculator,
25 MambaStateDtypeCalculator,
26 MambaStateShapeCalculator,
27 )
@@ -459,14 +461,19 @@ class Lfm2ForCausalLM(
461 conv_kernel=hf_config.conv_L_cache,
462 )
463
464 + @classmethod
465 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc]:
466 + return MambaStateCopyFuncCalculator.short_conv_state_copy_func()
467 +
468 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
469 config = vllm_config.model_config.hf_config
470 quant_config = vllm_config.quant_config
471 cache_config = vllm_config.cache_config
466 -
467 - assert not cache_config.enable_prefix_caching, (
468 - "Lfm2 currently does not support prefix caching"
469 - )
472 + if cache_config.mamba_cache_mode == "all":
473 + raise NotImplementedError(
474 + "Lfm2 currently does not support 'all' prefix caching, "
475 + "please use '--mamba-cache-mode=align' instead"
476 + )
477
478 super().__init__()
479 self.config = config
vllm/model_executor/models/lfm2_moe.py
+6
@@ -25,6 +25,8 @@ from vllm.model_executor.layers.linear import (
25 )
26 from vllm.model_executor.layers.logits_processor import LogitsProcessor
27 from vllm.model_executor.layers.mamba.mamba_utils import (
28 + MambaStateCopyFunc,
29 + MambaStateCopyFuncCalculator,
30 MambaStateDtypeCalculator,
31 MambaStateShapeCalculator,
32 )
@@ -640,6 +642,10 @@ class Lfm2MoeForCausalLM(
642 conv_kernel=hf_config.conv_L_cache,
643 )
644
645 + @classmethod
646 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc]:
647 + return MambaStateCopyFuncCalculator.short_conv_state_copy_func()
648 +
649 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
650 config = vllm_config.model_config.hf_config
651 quant_config = vllm_config.quant_config
vllm/model_executor/models/mamba.py
+6
@@ -16,6 +16,8 @@ from vllm.model_executor.layers.layernorm import RMSNorm
16 from vllm.model_executor.layers.logits_processor import LogitsProcessor
17 from vllm.model_executor.layers.mamba.mamba_mixer import MambaMixer
18 from vllm.model_executor.layers.mamba.mamba_utils import (
19 + MambaStateCopyFunc,
20 + MambaStateCopyFuncCalculator,
21 MambaStateDtypeCalculator,
22 MambaStateShapeCalculator,
23 )
@@ -261,6 +263,10 @@ class MambaForCausalLM(
263 conv_kernel=hf_config.conv_kernel,
264 )
265
266 + @classmethod
267 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
268 + return MambaStateCopyFuncCalculator.mamba1_state_copy_func()
269 +
270 def copy_inputs_before_cuda_graphs(self, input_buffers, **kwargs):
271 return self.mamba_cache.copy_inputs_before_cuda_graphs(input_buffers, **kwargs)
272
vllm/model_executor/models/mamba2.py
+6
@@ -15,6 +15,8 @@ from vllm.model_executor.layers.layernorm import RMSNorm
15 from vllm.model_executor.layers.logits_processor import LogitsProcessor
16 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
17 from vllm.model_executor.layers.mamba.mamba_utils import (
18 + MambaStateCopyFunc,
19 + MambaStateCopyFuncCalculator,
20 MambaStateDtypeCalculator,
21 MambaStateShapeCalculator,
22 )
@@ -228,6 +230,10 @@ class Mamba2ForCausalLM(
230 conv_kernel=hf_config.conv_kernel,
231 )
232
233 + @classmethod
234 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
235 + return MambaStateCopyFuncCalculator.mamba2_state_copy_func()
236 +
237 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
238 config = vllm_config.model_config.hf_config
239
vllm/model_executor/models/minimax_text_01.py
+6
@@ -35,6 +35,8 @@ from vllm.model_executor.layers.linear import (
35 from vllm.model_executor.layers.logits_processor import LogitsProcessor
36 from vllm.model_executor.layers.mamba.linear_attn import MiniMaxText01LinearAttention
37 from vllm.model_executor.layers.mamba.mamba_utils import (
38 + MambaStateCopyFunc,
39 + MambaStateCopyFuncCalculator,
40 MambaStateDtypeCalculator,
41 MambaStateShapeCalculator,
42 )
@@ -1006,3 +1008,7 @@ class MiniMaxText01ForCausalLM(nn.Module, HasInnerState, IsHybrid):
1008 tp_size=parallel_config.tensor_parallel_size,
1009 head_dim=hf_config.head_dim,
1010 )
1011 +
1012 + @classmethod
1013 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc]:
1014 + return MambaStateCopyFuncCalculator.linear_attention_state_copy_func()
vllm/model_executor/models/nano_nemotron_vl.py
+4
@@ -2128,3 +2128,7 @@ class NemotronH_Nano_VL_V2(
2128 temp_vllm_config = copy.deepcopy(vllm_config)
2129 temp_vllm_config.model_config.hf_config = text_config
2130 return NemotronHForCausalLM.get_mamba_state_dtype_from_config(temp_vllm_config)
2131 +
2132 + @classmethod
2133 + def get_mamba_state_copy_func(cls):
2134 + return NemotronHForCausalLM.get_mamba_state_copy_func()
vllm/model_executor/models/nemotron_h.py
+6
@@ -45,6 +45,8 @@ from vllm.model_executor.layers.linear import (
45 from vllm.model_executor.layers.logits_processor import LogitsProcessor
46 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
47 from vllm.model_executor.layers.mamba.mamba_utils import (
48 + MambaStateCopyFunc,
49 + MambaStateCopyFuncCalculator,
50 MambaStateDtypeCalculator,
51 MambaStateShapeCalculator,
52 )
@@ -809,6 +811,10 @@ class NemotronHForCausalLM(
811 conv_kernel=hf_config.conv_kernel,
812 )
813
814 + @classmethod
815 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
816 + return MambaStateCopyFuncCalculator.mamba2_state_copy_func()
817 +
818 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
819 config = vllm_config.model_config.hf_config
820 self.vllm_config = vllm_config
vllm/model_executor/models/plamo2.py
+6
@@ -27,6 +27,8 @@ from vllm.model_executor.layers.linear import (
27 from vllm.model_executor.layers.logits_processor import LogitsProcessor
28 from vllm.model_executor.layers.mamba.abstract import MambaBase
29 from vllm.model_executor.layers.mamba.mamba_utils import (
30 + MambaStateCopyFunc,
31 + MambaStateCopyFuncCalculator,
32 MambaStateDtypeCalculator,
33 MambaStateShapeCalculator,
34 )
@@ -899,6 +901,10 @@ class Plamo2ForCausalLM(
901 conv_kernel=hf_config.mamba_d_conv,
902 )
903
904 + @classmethod
905 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
906 + return MambaStateCopyFuncCalculator.mamba2_state_copy_func()
907 +
908 def compute_logits(
909 self,
910 hidden_states: torch.Tensor,
vllm/model_executor/models/qwen3_next.py
+11 -3
@@ -48,6 +48,8 @@ from vllm.model_executor.layers.logits_processor import LogitsProcessor
48 from vllm.model_executor.layers.mamba.abstract import MambaBase
49 from vllm.model_executor.layers.mamba.mamba_mixer2 import mamba_v2_sharded_weight_loader
50 from vllm.model_executor.layers.mamba.mamba_utils import (
51 + MambaStateCopyFunc,
52 + MambaStateCopyFuncCalculator,
53 MambaStateDtypeCalculator,
54 MambaStateShapeCalculator,
55 )
@@ -1205,9 +1207,11 @@ class Qwen3NextForCausalLM(
1207 cache_config = vllm_config.cache_config
1208
1209 scheduler_config = vllm_config.scheduler_config
1208 - assert not cache_config.enable_prefix_caching, (
1209 - "Qwen3Next currently does not support prefix caching"
1210 - )
1210 + if cache_config.mamba_cache_mode == "all":
1211 + raise NotImplementedError(
1212 + "Qwen3Next currently does not support 'all' prefix caching, "
1213 + "please use '--mamba-cache-mode=align' instead"
1214 + )
1215 self.quant_config = vllm_config.quant_config
1216
1217 super().__init__()
@@ -1278,6 +1282,10 @@ class Qwen3NextForCausalLM(
1282 num_spec,
1283 )
1284
1285 + @classmethod
1286 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
1287 + return MambaStateCopyFuncCalculator.gated_delta_net_state_copy_func()
1288 +
1289 def compute_logits(
1290 self,
1291 hidden_states: torch.Tensor,
vllm/model_executor/models/qwen3_next_mtp.py
+5 -3
@@ -234,9 +234,11 @@ class Qwen3NextMTP(nn.Module, QwenNextMixtureOfExperts):
234 config = vllm_config.model_config.hf_config
235 self.vllm_config = vllm_config
236 cache_config = vllm_config.cache_config
237 - assert not cache_config.enable_prefix_caching, (
238 - "Qwen3NextMTP currently does not support prefix caching"
239 - )
237 + if cache_config.mamba_cache_mode == "all":
238 + raise NotImplementedError(
239 + "Qwen3NextMTP currently does not support 'all' prefix caching, "
240 + "please use '--mamba-cache-mode=align' instead"
241 + )
242
243 self.quant_config = vllm_config.quant_config
244
vllm/model_executor/models/zamba2.py
+6
@@ -32,6 +32,8 @@ from vllm.model_executor.layers.linear import (
32 from vllm.model_executor.layers.logits_processor import LogitsProcessor
33 from vllm.model_executor.layers.mamba.mamba_mixer2 import MambaMixer2
34 from vllm.model_executor.layers.mamba.mamba_utils import (
35 + MambaStateCopyFunc,
36 + MambaStateCopyFuncCalculator,
37 MambaStateDtypeCalculator,
38 MambaStateShapeCalculator,
39 )
@@ -891,6 +893,10 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsMambaPrefixC
893 conv_kernel=hf_config.mamba_d_conv,
894 )
895
896 + @classmethod
897 + def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc, MambaStateCopyFunc]:
898 + return MambaStateCopyFuncCalculator.mamba2_state_copy_func()
899 +
900 def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
901 """Initialize the Zamba2 model for causal language modeling.
902
vllm/v1/attention/backends/gdn_attn.py
+11 -4
@@ -16,6 +16,7 @@ from vllm.v1.attention.backend import (
16 from vllm.v1.attention.backends.utils import (
17 PAD_SLOT_ID,
18 compute_causal_conv1d_metadata,
19 + mamba_get_block_table_tensor,
20 split_decodes_and_prefills,
21 )
22 from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
@@ -158,6 +159,12 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
159 query_start_loc_cpu = m.query_start_loc_cpu
160 context_lens_tensor = m.compute_num_computed_tokens()
161 nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
162 + block_table_tensor = mamba_get_block_table_tensor(
163 + m.block_table_tensor,
164 + m.seq_lens,
165 + self.kv_cache_spec,
166 + self.vllm_config.cache_config.mamba_cache_mode,
167 + )
168
169 spec_sequence_masks_cpu: torch.Tensor | None = None
170 if (
@@ -189,7 +196,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
196 spec_token_indx = None
197 non_spec_token_indx = None
198 spec_state_indices_tensor = None
192 - non_spec_state_indices_tensor = m.block_table_tensor[:, 0]
199 + non_spec_state_indices_tensor = block_table_tensor[:, 0]
200 spec_query_start_loc = None
201 non_spec_query_start_loc = query_start_loc
202 non_spec_query_start_loc_cpu = query_start_loc_cpu
@@ -221,7 +228,7 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
228 non_spec_token_indx = torch.empty(
229 0, dtype=torch.int32, device=query_start_loc.device
230 )
224 - spec_state_indices_tensor = m.block_table_tensor[:, : self.num_spec + 1]
231 + spec_state_indices_tensor = block_table_tensor[:, : self.num_spec + 1]
232 non_spec_state_indices_tensor = None
233 spec_query_start_loc = query_start_loc
234 non_spec_query_start_loc = None
@@ -235,10 +242,10 @@ class GDNAttentionMetadataBuilder(AttentionMetadataBuilder[GDNAttentionMetadata]
242 non_spec_token_indx = index[:num_non_spec_tokens]
243 spec_token_indx = index[num_non_spec_tokens:]
244
238 - spec_state_indices_tensor = m.block_table_tensor[
245 + spec_state_indices_tensor = block_table_tensor[
246 spec_sequence_masks, : self.num_spec + 1
247 ]
241 - non_spec_state_indices_tensor = m.block_table_tensor[
248 + non_spec_state_indices_tensor = block_table_tensor[
249 ~spec_sequence_masks, 0
250 ]
251
vllm/v1/attention/backends/linear_attn.py
+10 -2
@@ -11,7 +11,10 @@ from vllm.v1.attention.backend import (
11 AttentionMetadataBuilder,
12 CommonAttentionMetadata,
13 )
14 -from vllm.v1.attention.backends.utils import split_decodes_and_prefills
14 +from vllm.v1.attention.backends.utils import (
15 + mamba_get_block_table_tensor,
16 + split_decodes_and_prefills,
17 +)
18 from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
19
20
@@ -61,7 +64,12 @@ class LinearAttentionMetadataBuilder(AttentionMetadataBuilder[LinearAttentionMet
64 query_start_loc = common_attn_metadata.query_start_loc
65 seq_lens = common_attn_metadata.seq_lens
66
64 - state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
67 + state_indices_tensor = mamba_get_block_table_tensor(
68 + common_attn_metadata.block_table_tensor,
69 + common_attn_metadata.seq_lens,
70 + self.kv_cache_spec,
71 + self.vllm_config.cache_config.mamba_cache_mode,
72 + )[:, 0]
73
74 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
75 split_decodes_and_prefills(
vllm/v1/attention/backends/mamba_attn.py
+27 -8
@@ -18,6 +18,7 @@ from vllm.v1.attention.backend import (
18 from vllm.v1.attention.backends.utils import (
19 PAD_SLOT_ID,
20 compute_causal_conv1d_metadata,
21 + mamba_get_block_table_tensor,
22 split_decodes_and_prefills,
23 )
24 from vllm.v1.kv_cache_interface import AttentionSpec, MambaSpec
@@ -41,11 +42,15 @@ class BaseMambaAttentionMetadata:
42
43 state_indices_tensor: torch.Tensor
44
44 - # The following tensors are only used for prefix caching and are None if disabled
45 + # The following tensors are only used for prefix caching in all mode and
46 + # are None if disabled
47 block_idx_last_scheduled_token: torch.Tensor | None
48 block_idx_first_scheduled_token_p: torch.Tensor | None
49 block_idx_last_computed_token: torch.Tensor | None
50
51 + # The following tensor is only used for prefix caching in align mode
52 + seq_lens: torch.Tensor
53 +
54 # The following attributes are for triton implementation of causal_conv1d
55 nums_dict: dict | None = None
56 batch_ptr: torch.Tensor | None = None
@@ -78,7 +83,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
83 self.compilation_config.max_cudagraph_capture_size,
84 )
85
81 - if self.vllm_config.cache_config.enable_prefix_caching:
86 + if self.vllm_config.cache_config.mamba_cache_mode == "all":
87 self.state_indices_tensor = torch.empty(
88 (
89 self.decode_cudagraph_max_bs,
@@ -198,7 +203,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
203 # for causal_conv1d
204 nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
205
201 - if self.vllm_config.cache_config.enable_prefix_caching:
206 + if self.vllm_config.cache_config.mamba_cache_mode == "all":
207 num_computed_tokens = common_attn_metadata.compute_num_computed_tokens()
208
209 # Return a tensor of shape (#requests, #max blocks)
@@ -214,7 +219,12 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
219 )
220 else:
221 # Always return just a single block per each request:
217 - state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
222 + state_indices_tensor = mamba_get_block_table_tensor(
223 + common_attn_metadata.block_table_tensor,
224 + common_attn_metadata.seq_lens,
225 + self.kv_cache_spec,
226 + self.vllm_config.cache_config.mamba_cache_mode,
227 + )[:, 0]
228
229 if num_prefills > 0:
230 if num_computed_tokens is None:
@@ -239,7 +249,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
249 )
250 )
251
242 - if self.vllm_config.cache_config.enable_prefix_caching:
252 + if self.vllm_config.cache_config.mamba_cache_mode == "all":
253 assert num_computed_tokens is not None
254 num_computed_tokens_p = num_computed_tokens[
255 num_reqs - num_prefills : num_reqs
@@ -258,7 +268,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
268 state_indices_tensor = self.state_indices_tensor[:num_decode_tokens]
269 state_indices_tensor[num_decodes:] = PAD_SLOT_ID
270
261 - if self.vllm_config.cache_config.enable_prefix_caching:
271 + if self.vllm_config.cache_config.mamba_cache_mode == "all":
272 self.block_idx_last_scheduled_token[:num_decodes].copy_(
273 block_idx_last_scheduled_token, non_blocking=True
274 )
@@ -286,6 +296,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
296 block_idx_last_computed_token=block_idx_last_computed_token,
297 num_computed_tokens_p=num_computed_tokens_p,
298 num_reqs=num_reqs,
299 + seq_lens=common_attn_metadata.seq_lens,
300 nums_dict=nums_dict,
301 batch_ptr=batch_ptr,
302 token_chunk_offset_ptr=token_chunk_offset_ptr,
@@ -298,8 +309,16 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
309 slot_mapping: torch.Tensor,
310 ) -> M:
311 new_metadata = copy.copy(metadata)
301 - prefix_caching = self.vllm_config.cache_config.enable_prefix_caching
302 - state_indices_t = blk_table if prefix_caching else blk_table[:, 0]
312 + state_indices_t = mamba_get_block_table_tensor(
313 + blk_table,
314 + metadata.seq_lens,
315 + self.kv_cache_spec,
316 + self.vllm_config.cache_config.mamba_cache_mode,
317 + )
318 + if self.vllm_config.cache_config.mamba_cache_mode in ("none", "align"):
319 + # Only needs the block that saves the running state
320 + state_indices_t = state_indices_t[:, 0]
321 +
322 num_reqs = blk_table.shape[0]
323
324 # For CUDA graphs, copy to persistent buffer
vllm/v1/attention/backends/utils.py
+38
@@ -17,6 +17,7 @@ from typing_extensions import runtime_checkable
17
18 from vllm.config import VllmConfig, get_layers_from_vllm_config
19 from vllm.utils.math_utils import cdiv
20 +from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec
21
22 if TYPE_CHECKING:
23 from vllm.v1.core.sched.output import SchedulerOutput
@@ -854,3 +855,40 @@ def extend_all_queries_by_1(
855 slot_mapping=new_slot_mapping,
856 )
857 return new_cad
858 +
859 +
860 +def mamba_get_block_table_tensor(
861 + block_table: torch.Tensor,
862 + seq_lens: torch.Tensor,
863 + kv_cache_spec: KVCacheSpec,
864 + mamba_cache_mode: str,
865 +) -> torch.Tensor:
866 + """
867 + Get the block table tensor for mamba kernels from the input
868 + common_attn_metadata.block_table_tensor given different mamba cache modes.
869 +
870 + - "all": input (#requests, cdiv(max_model_len, block_size));
871 + output (#requests, cdiv(max_model_len, block_size)).
872 +
873 + - "none": input (#requests, 1 + num_speculative_blocks);
874 + output (#requests, 1 + num_speculative_blocks).
875 +
876 + - "align": input (#requests, cdiv(max_model_len, block_size));
877 + output (#requests, 1 + num_speculative_blocks), which are the last
878 + 1 + num_speculative_blocks of each request.
879 + """
880 + if mamba_cache_mode in ("all", "none"):
881 + return block_table
882 + else:
883 + assert isinstance(kv_cache_spec, MambaSpec)
884 + # NOTE: For 0-length requests in CUDA graph, use a start_index of 0
885 + # to handle the invalid block table.
886 + start_indices = torch.clamp(
887 + (seq_lens - 1) // kv_cache_spec.block_size,
888 + min=0,
889 + )
890 + offsets = torch.arange(
891 + 1 + kv_cache_spec.num_speculative_blocks, device=block_table.device
892 + )
893 + indices_to_gather = start_indices.unsqueeze(1) + offsets
894 + return torch.gather(block_table, 1, indices_to_gather)
vllm/v1/core/block_pool.py
+2 -1
@@ -255,7 +255,8 @@ class BlockPool:
255 )
256 for i, blk in enumerate(new_full_blocks):
257 # Some blocks may be null blocks when enabling sparse attention like
258 - # sliding window attention. We skip null blocks here.
258 + # sliding window attention, or Mamba models with prefix-caching in
259 + # align mode. We skip null blocks here.
260 if blk.is_null:
261 continue
262 assert blk.block_hash is None
vllm/v1/core/kv_cache_coordinator.py
+11 -1
@@ -75,6 +75,7 @@ class KVCacheCoordinator(ABC):
75 new_computed_blocks: tuple[Sequence[KVCacheBlock], ...],
76 num_encoder_tokens: int,
77 total_computed_tokens: int,
78 + num_tokens_main_model: int,
79 ) -> int:
80 """
81 Get the number of blocks needed to be allocated for the request.
@@ -88,6 +89,9 @@ class KVCacheCoordinator(ABC):
89 num_encoder_tokens: The number of encoder tokens for allocating
90 blocks for cross-attention.
91 total_computed_tokens: Include both local and external tokens.
92 + num_tokens_main_model: The number of tokens for the main model (aka target
93 + model in spec decode). w/o spec decode, it is num_tokens;
94 + with spec decode, it is num_tokens - num_lookahead_tokens.
95
96 Returns:
97 The number of blocks to allocate.
@@ -98,7 +102,7 @@ class KVCacheCoordinator(ABC):
102 # For cross-attention, we issue a single static allocation
103 # of blocks based on the number of encoder input tokens.
104 num_blocks_to_allocate += manager.get_num_blocks_to_allocate(
101 - request_id, num_encoder_tokens, [], 0
105 + request_id, num_encoder_tokens, [], 0, num_encoder_tokens
106 )
107 else:
108 num_blocks_to_allocate += manager.get_num_blocks_to_allocate(
@@ -106,6 +110,7 @@ class KVCacheCoordinator(ABC):
110 num_tokens,
111 new_computed_blocks[i],
112 total_computed_tokens,
113 + num_tokens_main_model,
114 )
115 return num_blocks_to_allocate
116
@@ -139,6 +144,7 @@ class KVCacheCoordinator(ABC):
144 self,
145 request_id: str,
146 num_tokens: int,
147 + num_tokens_main_model: int,
148 num_encoder_tokens: int = 0,
149 ) -> tuple[list[KVCacheBlock], ...]:
150 """
@@ -149,6 +155,9 @@ class KVCacheCoordinator(ABC):
155 request_id: The request ID.
156 num_tokens: The total number of tokens that need a slot (including
157 tokens that are already allocated).
158 + num_tokens_main_model: The number of tokens for the main model (aka target
159 + model in spec decode). w/o spec decode, it is num_tokens;
160 + with spec decode, it is num_tokens - num_lookahead_tokens.
161 num_encoder_tokens: The number of encoder tokens for allocating
162 blocks for cross-attention.
163
@@ -161,6 +170,7 @@ class KVCacheCoordinator(ABC):
170 num_encoder_tokens
171 if isinstance(manager, CrossAttentionManager)
172 else num_tokens,
173 + num_tokens_main_model,
174 )
175 for manager in self.single_type_managers
176 )
vllm/v1/core/kv_cache_manager.py
+7 -2
@@ -307,8 +307,9 @@ class KVCacheManager:
307 num_local_computed_tokens + num_external_computed_tokens,
308 self.max_model_len,
309 )
310 + num_tokens_main_model = total_computed_tokens + num_new_tokens
311 num_tokens_need_slot = min(
311 - total_computed_tokens + num_new_tokens + num_lookahead_tokens,
312 + num_tokens_main_model + num_lookahead_tokens,
313 self.max_model_len,
314 )
315
@@ -329,6 +330,7 @@ class KVCacheManager:
330 num_encoder_tokens=num_encoder_tokens,
331 total_computed_tokens=num_local_computed_tokens
332 + num_external_computed_tokens,
333 + num_tokens_main_model=num_tokens_main_model,
334 )
335
336 if num_blocks_to_allocate > self.block_pool.get_num_free_blocks():
@@ -349,7 +351,10 @@ class KVCacheManager:
351 )
352
353 new_blocks = self.coordinator.allocate_new_blocks(
352 - request.request_id, num_tokens_need_slot, num_encoder_tokens
354 + request.request_id,
355 + num_tokens_need_slot,
356 + num_tokens_main_model,
357 + num_encoder_tokens,
358 )
359
360 # P/D: delay caching blocks if we have to recv from
vllm/v1/core/sched/scheduler.py
+76 -1
@@ -47,7 +47,7 @@ from vllm.v1.core.sched.output import (
47 from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue
48 from vllm.v1.core.sched.utils import check_stop, remove_all
49 from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs
50 -from vllm.v1.kv_cache_interface import KVCacheConfig
50 +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
51 from vllm.v1.metrics.perf import ModelMetrics, PerfStats
52 from vllm.v1.metrics.stats import (
53 PrefixCacheStats,
@@ -226,6 +226,17 @@ class Scheduler(SchedulerInterface):
226 )
227 self.use_pp = self.parallel_config.pipeline_parallel_size > 1
228 self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER
229 +
230 + def has_mamba_layers(kv_cache_config: KVCacheConfig) -> bool:
231 + return any(
232 + isinstance(group_spec.kv_cache_spec, MambaSpec)
233 + for group_spec in kv_cache_config.kv_cache_groups
234 + )
235 +
236 + self.has_mamba_layers = has_mamba_layers(kv_cache_config)
237 + self.need_mamba_block_aligned_split = (
238 + self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align"
239 + )
240 self.perf_metrics: ModelMetrics | None = None
241 if self.log_stats and vllm_config.observability_config.enable_mfu_metrics:
242 self.perf_metrics = ModelMetrics(vllm_config)
@@ -250,6 +261,53 @@ class Scheduler(SchedulerInterface):
261 vllm_config=self.vllm_config,
262 )
263
264 + def _mamba_block_aligned_split(
265 + self,
266 + request: Request,
267 + num_new_tokens: int,
268 + num_new_local_computed_tokens: int = 0,
269 + num_external_computed_tokens: int = 0,
270 + ) -> int:
271 + assert num_external_computed_tokens == 0, (
272 + "External KV connector is not verified yet"
273 + )
274 + # TODO: need check for resume requests
275 + if request.num_output_tokens == 0: # prefill
276 + # To enable block-aligned caching of the Mamba state, `num_new_tokens`
277 + # must be a multiple of `block_size`.
278 + # As an exception, if `num_new_tokens` is less than `block_size`, the
279 + # state is simply not cached, requiring no special handling.
280 + # Additionally, when Eagle mode is enabled, FullAttn prunes the last
281 + # matching block. To prevent this from causing a Mamba cache miss, the
282 + # last chunk must be larger than `block_size`.
283 + block_size = self.cache_config.block_size
284 + last_cache_position = (
285 + request.num_prompt_tokens - request.num_prompt_tokens % block_size
286 + )
287 + # eagle prune
288 + if self.use_eagle:
289 + last_cache_position = max(last_cache_position - block_size, 0)
290 + num_computed_tokens = (
291 + request.num_computed_tokens
292 + + num_new_local_computed_tokens
293 + + num_external_computed_tokens
294 + )
295 + num_computed_tokens_after_sched = num_computed_tokens + num_new_tokens
296 + if num_computed_tokens_after_sched < last_cache_position:
297 + # align to block_size
298 + num_new_tokens = num_new_tokens // block_size * block_size
299 + elif (
300 + num_computed_tokens
301 + < last_cache_position
302 + < num_computed_tokens_after_sched
303 + ):
304 + # force to cache the last chunk
305 + num_new_tokens = last_cache_position - num_computed_tokens
306 + else:
307 + # prefill the last few tokens
308 + pass
309 + return num_new_tokens
310 +
311 def schedule(self) -> SchedulerOutput:
312 # NOTE(woosuk) on the scheduling algorithm:
313 # There's no "decoding phase" nor "prefill phase" in the scheduler.
@@ -340,6 +398,11 @@ class Scheduler(SchedulerInterface):
398 shift_computed_tokens=1 if self.use_eagle else 0,
399 )
400
401 + if self.need_mamba_block_aligned_split:
402 + num_new_tokens = self._mamba_block_aligned_split(
403 + request, num_new_tokens
404 + )
405 +
406 if num_new_tokens == 0:
407 # The request cannot be scheduled because one of the following
408 # reasons:
@@ -350,6 +413,8 @@ class Scheduler(SchedulerInterface):
413 # its max_total_tokens or max_model_len.
414 # 2. The encoder budget is exhausted.
415 # 3. The encoder cache is exhausted.
416 + # 4. Insufficient budget for a block-aligned chunk in hybrid
417 + # models with mamba cache mode \"align\".
418 # NOTE(woosuk): Here, by doing `continue` instead of `break`,
419 # we do not strictly follow the FCFS scheduling policy and
420 # allow the lower-priority requests to be scheduled.
@@ -608,6 +673,16 @@ class Scheduler(SchedulerInterface):
673 # The request cannot be scheduled.
674 break
675
676 + if self.need_mamba_block_aligned_split:
677 + num_new_tokens = self._mamba_block_aligned_split(
678 + request,
679 + num_new_tokens,
680 + num_new_local_computed_tokens,
681 + num_external_computed_tokens,
682 + )
683 + if num_new_tokens == 0:
684 + break
685 +
686 # Handles an edge case when P/D Disaggregation
687 # is used with Spec Decoding where an
688 # extra block gets allocated which
vllm/v1/core/single_type_kv_cache_manager.py
+171 -32
@@ -66,12 +66,17 @@ class SingleTypeKVCacheManager(ABC):
66 self.kv_cache_group_id = kv_cache_group_id
67 self._null_block = block_pool.null_block
68
69 + @classmethod
70 + def _get_num_evictable_blocks(cls, blocks: Sequence[KVCacheBlock]):
71 + return sum(blk.ref_cnt == 0 and not blk.is_null for blk in blocks)
72 +
73 def get_num_blocks_to_allocate(
74 self,
75 request_id: str,
76 num_tokens: int,
77 new_computed_blocks: Sequence[KVCacheBlock],
78 total_computed_tokens: int,
79 + num_tokens_main_model: int,
80 ) -> int:
81 """
82 Get the number of blocks needed to be allocated for the request.
@@ -84,6 +89,9 @@ class SingleTypeKVCacheManager(ABC):
89 prefix caching.
90 total_computed_tokens: Include both local and external computed
91 tokens.
92 + num_tokens_main_model: The number of tokens for the main model (aka target
93 + model in spec decode). w/o spec decode, it is num_tokens;
94 + with spec decode, it is num_tokens - num_lookahead_tokens.
95
96 Returns:
97 The number of blocks to allocate.
@@ -121,9 +129,8 @@ class SingleTypeKVCacheManager(ABC):
129 # If a computed block is an eviction candidate (in the free queue and
130 # ref_cnt == 0), it will be removed from the free queue when touched by
131 # the allocated request, so we must count it in the free-capacity check.
124 - num_evictable_blocks = sum(
125 - blk.ref_cnt == 0 and not blk.is_null
126 - for blk in new_computed_blocks[num_skipped_new_computed_blocks:]
132 + num_evictable_blocks = self._get_num_evictable_blocks(
133 + new_computed_blocks[num_skipped_new_computed_blocks:]
134 )
135 return num_new_blocks + num_evictable_blocks
136
@@ -199,7 +206,7 @@ class SingleTypeKVCacheManager(ABC):
206 req_blocks.extend(allocated_blocks)
207
208 def allocate_new_blocks(
202 - self, request_id: str, num_tokens: int
209 + self, request_id: str, num_tokens: int, num_tokens_main_model: int
210 ) -> list[KVCacheBlock]:
211 """
212 Allocate new blocks for the request to give it at least `num_tokens`
@@ -209,7 +216,9 @@ class SingleTypeKVCacheManager(ABC):
216 request_id: The request ID.
217 num_tokens: The total number of tokens that need a slot (including
218 tokens that are already allocated).
212 -
219 + num_tokens_main_model: The number of tokens for the main model (aka target
220 + model in spec decode). w/o spec decode, it is num_tokens;
221 + with spec decode, it is num_tokens - num_lookahead_tokens.
222 Returns:
223 The new allocated blocks.
224 """
@@ -450,12 +459,9 @@ class FullAttentionManager(SingleTypeKVCacheManager):
459
460
461 class SlidingWindowManager(SingleTypeKVCacheManager):
453 - def __init__(
454 - self, kv_cache_spec: SlidingWindowSpec, block_pool: BlockPool, **kwargs
455 - ) -> None:
456 - super().__init__(kv_cache_spec, block_pool, **kwargs)
462 + def __init__(self, kv_cache_spec: SlidingWindowSpec, **kwargs) -> None:
463 + super().__init__(kv_cache_spec, **kwargs)
464 self.sliding_window = kv_cache_spec.sliding_window
458 - self._null_block = block_pool.null_block
465
466 @classmethod
467 def find_longest_cache_hit(
@@ -586,12 +592,9 @@ class SlidingWindowManager(SingleTypeKVCacheManager):
592
593
594 class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
589 - def __init__(
590 - self, kv_cache_spec: ChunkedLocalAttentionSpec, block_pool: BlockPool, **kwargs
591 - ) -> None:
592 - super().__init__(kv_cache_spec, block_pool, **kwargs)
595 + def __init__(self, kv_cache_spec: ChunkedLocalAttentionSpec, **kwargs) -> None:
596 + super().__init__(kv_cache_spec, **kwargs)
597 self.attention_chunk_size = kv_cache_spec.attention_chunk_size
594 - self._null_block = block_pool.null_block
598
599 @classmethod
600 def find_longest_cache_hit(
@@ -739,6 +742,17 @@ class ChunkedLocalAttentionManager(SingleTypeKVCacheManager):
742
743
744 class MambaManager(SingleTypeKVCacheManager):
745 + def __init__(self, kv_cache_spec: MambaSpec, **kwargs) -> None:
746 + super().__init__(kv_cache_spec, **kwargs)
747 + self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode
748 + self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks
749 + if self.mamba_cache_mode == "align":
750 + # Mapping from request ID to the index of the block
751 + # allocated in the previous step
752 + self.last_state_block_idx: dict[str, int] = {}
753 + # The set of the requests that have been allocated blocks
754 + self._allocated_block_reqs: set[str] = set()
755 +
756 @classmethod
757 def find_longest_cache_hit(
758 cls,
@@ -787,6 +801,28 @@ class MambaManager(SingleTypeKVCacheManager):
801
802 return computed_blocks
803
804 + def remove_skipped_blocks(self, request_id: str, num_computed_tokens: int) -> None:
805 + assert isinstance(self.kv_cache_spec, MambaSpec)
806 + super().remove_skipped_blocks(request_id, num_computed_tokens)
807 + if self.mamba_cache_mode == "align":
808 + # `last_state_block_idx` refers to the block index allocated two steps ago.
809 + # The block allocated in the previous step is used to copy Mamba states
810 + # into the block allocated in the current step; the earlier block is
811 + # no longer needed and should be freed here.
812 + last_state_block_idx = self.last_state_block_idx.get(request_id)
813 + # Blocks allocated during prefill may be non-contiguous. Use
814 + # `last_state_block_idx` to free the appropriate block and replace it
815 + # with a null block.
816 + if (
817 + last_state_block_idx is not None
818 + and last_state_block_idx
819 + < cdiv(num_computed_tokens, self.block_size) - 1
820 + ):
821 + blocks = self.req_to_blocks[request_id]
822 + if blocks[last_state_block_idx] != self._null_block:
823 + self.block_pool.free_blocks([blocks[last_state_block_idx]])
824 + blocks[last_state_block_idx] = self._null_block
825 +
826 def get_num_common_prefix_blocks(self, running_request_id: str) -> int:
827 """
828 cascade attention is not supported by mamba
@@ -799,31 +835,134 @@ class MambaManager(SingleTypeKVCacheManager):
835 num_tokens: int,
836 new_computed_blocks: Sequence[KVCacheBlock],
837 total_computed_tokens: int,
838 + num_tokens_main_model: int,
839 ) -> int:
803 - # Allocate extra `num_speculative_blocks` blocks for
804 - # speculative decoding (MTP/EAGLE) with linear attention.
840 assert isinstance(self.kv_cache_spec, MambaSpec)
806 - if self.kv_cache_spec.num_speculative_blocks > 0:
807 - num_tokens += (
808 - self.kv_cache_spec.block_size
809 - * self.kv_cache_spec.num_speculative_blocks
841 + if self.mamba_cache_mode != "align":
842 + # Allocate extra `num_speculative_blocks` blocks for
843 + # speculative decoding (MTP/EAGLE) with linear attention.
844 + if self.num_speculative_blocks > 0:
845 + num_tokens += (
846 + self.kv_cache_spec.block_size * self.num_speculative_blocks
847 + )
848 + return super().get_num_blocks_to_allocate(
849 + request_id,
850 + num_tokens,
851 + new_computed_blocks,
852 + total_computed_tokens,
853 + num_tokens_main_model,
854 )
811 - return super().get_num_blocks_to_allocate(
812 - request_id, num_tokens, new_computed_blocks, total_computed_tokens
813 - )
855 + else:
856 + # We don't allocate blocks for lookahead tokens in align mode, because if
857 + # x * block_size tokens are scheduled, num_tokens is
858 + # x * block_size + num_lookahead_tokens and breaks the alignment.
859 + # We can ignore lookahead tokens because current draft models don't have
860 + # mamba layers.
861 + num_tokens = num_tokens_main_model
862 + num_required_blocks = (
863 + cdiv(num_tokens, self.block_size) + self.num_speculative_blocks
864 + )
865 + num_new_blocks = (
866 + num_required_blocks
867 + - len(new_computed_blocks)
868 + - len(self.req_to_blocks[request_id])
869 + )
870 + if num_new_blocks > 0:
871 + if request_id in self._allocated_block_reqs:
872 + # Old request. Needs at most 1 more blocks as we can reuse the
873 + # speculative blocks in previous step.
874 + num_new_blocks = 1
875 + else:
876 + # First prefill. Allocate 1 block for running state and the
877 + # speculative blocks.
878 + num_new_blocks = 1 + self.num_speculative_blocks
879 +
880 + num_evictable_computed_blocks = self._get_num_evictable_blocks(
881 + new_computed_blocks
882 + )
883 + return num_new_blocks + num_evictable_computed_blocks
884
885 def allocate_new_blocks(
816 - self, request_id: str, num_tokens: int
886 + self, request_id: str, num_tokens: int, num_tokens_main_model: int
887 ) -> list[KVCacheBlock]:
818 - # Allocate extra `num_speculative_blocks` blocks for
819 - # speculative decoding (MTP/EAGLE) with linear attention.
888 assert isinstance(self.kv_cache_spec, MambaSpec)
821 - if self.kv_cache_spec.num_speculative_blocks > 0:
822 - num_tokens += (
823 - self.kv_cache_spec.block_size
824 - * self.kv_cache_spec.num_speculative_blocks
889 + if self.mamba_cache_mode != "align":
890 + # Allocate extra `num_speculative_blocks` blocks for
891 + # speculative decoding (MTP/EAGLE) with linear attention.
892 + if self.num_speculative_blocks > 0:
893 + num_tokens += self.block_size * self.num_speculative_blocks
894 + return super().allocate_new_blocks(
895 + request_id, num_tokens, num_tokens_main_model
896 )
826 - return super().allocate_new_blocks(request_id, num_tokens)
897 + else:
898 + # We don't allocate blocks for lookahead tokens in align mode, because if
899 + # x * block_size tokens are scheduled, num_tokens is
900 + # x * block_size + num_lookahead_tokens and breaks the alignment.
901 + # We can ignore lookahead tokens because current draft models don't have
902 + # mamba layers.
903 + num_tokens = num_tokens_main_model
904 + req_blocks: list[KVCacheBlock] = self.req_to_blocks[request_id]
905 + num_required_blocks = (
906 + cdiv(num_tokens, self.block_size) + self.num_speculative_blocks
907 + )
908 + if num_required_blocks == len(req_blocks):
909 + return []
910 + else:
911 + assert num_required_blocks > len(req_blocks), (
912 + "num_required_blocks "
913 + f"{num_required_blocks} < len(req_blocks) {len(req_blocks)}"
914 + )
915 + prev_block_len = len(req_blocks)
916 + blocks_allocated = request_id in self._allocated_block_reqs
917 + # Record the last state block
918 + if blocks_allocated:
919 + # We always save the running state at the last
920 + # (1 + num_speculative_blocks) block
921 + self.last_state_block_idx[request_id] = (
922 + prev_block_len - 1 - self.num_speculative_blocks
923 + )
924 + elif prev_block_len > 0:
925 + # When a new request hits the prefix cache, the last block
926 + # saves the hit state.
927 + self.last_state_block_idx[request_id] = prev_block_len - 1
928 +
929 + num_skipped_blocks = (
930 + num_required_blocks - self.num_speculative_blocks - 1
931 + )
932 + # null blocks
933 + if prev_block_len < num_skipped_blocks:
934 + req_blocks.extend(
935 + [
936 + self._null_block
937 + for _ in range(prev_block_len, num_skipped_blocks)
938 + ]
939 + )
940 +
941 + if blocks_allocated:
942 + # reuse previous speculative blocks in this step
943 + for block_idx in range(
944 + prev_block_len - self.num_speculative_blocks, prev_block_len
945 + ):
946 + if block_idx < num_skipped_blocks:
947 + req_blocks.append(req_blocks[block_idx])
948 + req_blocks[block_idx] = self._null_block
949 + else:
950 + break
951 + num_new_blocks = num_required_blocks - len(req_blocks)
952 + if blocks_allocated:
953 + assert num_new_blocks <= 1
954 + else:
955 + assert num_new_blocks <= self.num_speculative_blocks + 1
956 + new_blocks = self.block_pool.get_new_blocks(num_new_blocks)
957 + req_blocks.extend(new_blocks)
958 + self._allocated_block_reqs.add(request_id)
959 + return req_blocks[prev_block_len:]
960 +
961 + def free(self, request_id: str) -> None:
962 + if self.mamba_cache_mode == "align":
963 + self._allocated_block_reqs.discard(request_id)
964 + self.last_state_block_idx.pop(request_id, None)
965 + super().free(request_id)
966
967 def get_num_skipped_tokens(self, num_computed_tokens: int) -> int:
968 """
vllm/v1/kv_cache_interface.py
+8 -2
@@ -276,6 +276,7 @@ class MambaSpec(KVCacheSpec):
276 dtypes: tuple[torch.dtype]
277 page_size_padded: int | None = None
278 mamba_type: str = "mamba2"
279 + mamba_cache_mode: str = "none"
280 num_speculative_blocks: int = 0
281
282 @property
@@ -290,8 +291,13 @@ class MambaSpec(KVCacheSpec):
291 return page_size
292
293 def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
293 - max_model_len = vllm_config.model_config.max_model_len
294 - return cdiv(max_model_len, self.block_size) * self.page_size_bytes
294 + if vllm_config.cache_config.mamba_cache_mode == "all":
295 + max_model_len = vllm_config.model_config.max_model_len
296 + return cdiv(max_model_len, self.block_size) * self.page_size_bytes
297 + elif vllm_config.cache_config.mamba_cache_mode == "align":
298 + return self.page_size_bytes * (2 + self.num_speculative_blocks)
299 + else:
300 + return self.page_size_bytes * (1 + self.num_speculative_blocks)
301
302
303 @dataclass(frozen=True)
vllm/v1/worker/block_table.py
+22 -23
@@ -8,6 +8,7 @@ from vllm.distributed import get_dcp_group, get_pcp_group
8 from vllm.logger import init_logger
9 from vllm.utils.math_utils import cdiv
10 from vllm.v1.utils import CpuGpuBuffer
11 +from vllm.v1.worker.cp_utils import get_total_cp_world_size
12
13 logger = init_logger(__name__)
14
@@ -261,47 +262,45 @@ class MultiGroupBlockTable:
262 device: torch.device,
263 block_sizes: list[int],
264 kernel_block_sizes: list[int],
264 - num_speculative_tokens: int = 0,
265 + max_num_blocks: list[int] | None = None,
266 cp_kv_cache_interleave_size: int = 1,
267 ) -> None:
267 - # Note(hc): each dcp rank only store
268 - # (max_model_len//dcp_world_size) tokens in kvcache,
269 - # so the block_size which used for calc max_num_blocks_per_req
270 - # must be multiplied by dcp_world_size.
271 - try:
272 - pcp_world_size = get_pcp_group().world_size
273 - except AssertionError:
274 - # PCP might not be initialized in testing
275 - pcp_world_size = 1
276 - try:
277 - dcp_world_size = get_dcp_group().world_size
278 - except AssertionError:
279 - # DCP might not be initialized in testing
280 - dcp_world_size = 1
281 -
268 if len(kernel_block_sizes) != len(block_sizes):
269 raise ValueError(
270 f"kernel_block_sizes length ({len(kernel_block_sizes)}) "
271 f"must match block_sizes length ({len(block_sizes)})"
272 )
287 -
288 - total_cp_world_size = dcp_world_size * pcp_world_size
273 + if max_num_blocks is None:
274 + # Note(hc): each dcp rank only store
275 + # (max_model_len//dcp_world_size) tokens in kvcache,
276 + # so the block_size which used for calc max_num_blocks_per_req
277 + # must be multiplied by dcp_world_size.
278 + total_cp_world_size = get_total_cp_world_size()
279 + max_num_blocks = [
280 + cdiv(max_model_len, block_size * total_cp_world_size)
281 + for block_size in block_sizes
282 + ]
283 +
284 + if len(max_num_blocks) != len(block_sizes):
285 + raise ValueError(
286 + f"max_num_blocks length ({len(max_num_blocks)}) "
287 + f"must match block_sizes length ({len(block_sizes)})"
288 + )
289
290 self.block_tables = [
291 BlockTable(
292 block_size,
293 max_num_reqs,
294 - max(
295 - cdiv(max_model_len, block_size * total_cp_world_size),
296 - 1 + num_speculative_tokens,
297 - ),
294 + max_num_blocks_per_req,
295 max_num_batched_tokens,
296 pin_memory,
297 device,
298 kernel_block_size,
299 cp_kv_cache_interleave_size,
300 )
304 - for block_size, kernel_block_size in zip(block_sizes, kernel_block_sizes)
301 + for block_size, kernel_block_size, max_num_blocks_per_req in zip(
302 + block_sizes, kernel_block_sizes, max_num_blocks
303 + )
304 ]
305
306 def append_row(self, block_ids: tuple[list[int], ...], row_idx: int) -> None:
vllm/v1/worker/cp_utils.py
+15
@@ -3,6 +3,7 @@
3 from typing import TYPE_CHECKING, Any, cast
4
5 from vllm.config import VllmConfig, get_layers_from_vllm_config
6 +from vllm.distributed import get_dcp_group, get_pcp_group
7
8 if TYPE_CHECKING:
9 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
@@ -40,3 +41,17 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None:
41 f"but the impl {layer_impl.__class__.__name__} "
42 "does not support PCP."
43 )
44 +
45 +
46 +def get_total_cp_world_size():
47 + try:
48 + pcp_world_size = get_pcp_group().world_size
49 + except AssertionError:
50 + # PCP might not be initialized in testing
51 + pcp_world_size = 1
52 + try:
53 + dcp_world_size = get_dcp_group().world_size
54 + except AssertionError:
55 + # DCP might not be initialized in testing
56 + dcp_world_size = 1
57 + return dcp_world_size * pcp_world_size
vllm/v1/worker/gpu_input_batch.py
+2 -2
@@ -89,11 +89,11 @@ class InputBatch:
89 vocab_size: int,
90 block_sizes: list[int], # The block_size of each kv cache group
91 kernel_block_sizes: list[int],
92 + max_num_blocks_per_req: list[int] | None = None,
93 logitsprocs: LogitsProcessors | None = None,
94 logitsprocs_need_output_token_ids: bool = False,
95 is_spec_decode: bool = False,
96 is_pooling_model: bool = False,
96 - num_speculative_tokens: int = 0,
97 cp_kv_cache_interleave_size: int = 1,
98 ):
99 self.is_pooling_model = is_pooling_model
@@ -146,7 +146,7 @@ class InputBatch:
146 device=device,
147 block_sizes=block_sizes,
148 kernel_block_sizes=kernel_block_sizes,
149 - num_speculative_tokens=num_speculative_tokens,
149 + max_num_blocks=max_num_blocks_per_req,
150 cp_kv_cache_interleave_size=cp_kv_cache_interleave_size,
151 )
152
vllm/v1/worker/gpu_model_runner.py
+53 -5
@@ -152,7 +152,11 @@ from vllm.v1.spec_decode.ngram_proposer import NgramProposer
152 from vllm.v1.spec_decode.suffix_decoding import SuffixDecodingProposer
153 from vllm.v1.structured_output.utils import apply_grammar_bitmask
154 from vllm.v1.utils import CpuGpuBuffer, record_function_or_nullcontext
155 -from vllm.v1.worker.cp_utils import check_attention_cp_compatibility
155 +from vllm.v1.worker import mamba_utils
156 +from vllm.v1.worker.cp_utils import (
157 + check_attention_cp_compatibility,
158 + get_total_cp_world_size,
159 +)
160 from vllm.v1.worker.dp_utils import coordinate_batch_across_dp
161 from vllm.v1.worker.ec_connector_model_runner_mixin import ECConnectorModelRunnerMixin
162 from vllm.v1.worker.gpu_input_batch import CachedRequestState, InputBatch
@@ -688,6 +692,7 @@ class GPUModelRunner(
692 # Ephemeral state transferred between execute_model() and sample_tokens().
693 self.execute_model_state: ExecuteModelState | None = None
694 self.kv_connector_output: KVConnectorOutput | None = None
695 + self.mamba_state_idx: dict[str, int] = {}
696 self.layerwise_nvtx_hooks_registered = False
697
698 def update_max_model_len(self, max_model_len: int) -> None:
@@ -1075,7 +1080,7 @@ class GPUModelRunner(
1080 self.input_batch.refresh_metadata()
1081
1082 def _update_states_after_model_execute(
1078 - self, output_token_ids: torch.Tensor
1083 + self, output_token_ids: torch.Tensor, scheduler_output: "SchedulerOutput"
1084 ) -> None:
1085 """Update the cached states after model execution.
1086
@@ -1111,6 +1116,16 @@ class GPUModelRunner(
1116 )
1117 for i, num_tokens in enumerate(num_accepted_tokens):
1118 self.input_batch.num_accepted_tokens_cpu[i] = num_tokens
1119 + if self.cache_config.mamba_cache_mode == "align":
1120 + mamba_utils.postprocess_mamba(
1121 + scheduler_output,
1122 + self.kv_cache_config,
1123 + self.input_batch,
1124 + self.requests,
1125 + self.mamba_state_idx,
1126 + self.compilation_config.static_forward_context,
1127 + self.model.get_mamba_state_copy_func(),
1128 + )
1129
1130 def _init_mrope_positions(self, req_state: CachedRequestState):
1131 model = self.get_model()
@@ -2751,7 +2766,6 @@ class GPUModelRunner(
2766 logits,
2767 sampling_metadata,
2768 )
2754 - self._update_states_after_model_execute(sampler_output.sampled_token_ids)
2769 return sampler_output
2770
2771 def _bookkeeping_sync(
@@ -3237,6 +3251,18 @@ class GPUModelRunner(
3251
3252 pad_attn = cudagraph_mode == CUDAGraphMode.FULL
3253
3254 + if self.cache_config.mamba_cache_mode == "align":
3255 + mamba_utils.preprocess_mamba(
3256 + scheduler_output,
3257 + self.kv_cache_config,
3258 + self.cache_config,
3259 + self.mamba_state_idx,
3260 + self.input_batch,
3261 + self.requests,
3262 + self.compilation_config.static_forward_context,
3263 + self.model.get_mamba_state_copy_func(),
3264 + )
3265 +
3266 use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0
3267 ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices
3268
@@ -3423,6 +3449,10 @@ class GPUModelRunner(
3449 with record_function_or_nullcontext("gpu_model_runner: sample"):
3450 sampler_output = self._sample(logits, spec_decode_metadata)
3451
3452 + self._update_states_after_model_execute(
3453 + sampler_output.sampled_token_ids, scheduler_output
3454 + )
3455 +
3456 self._draft_token_ids = None
3457 self._draft_token_req_ids = None
3458 self.input_batch.prev_sampled_token_ids = None
@@ -5322,6 +5352,24 @@ class GPUModelRunner(
5352 for kv_cache_group in kv_cache_config.kv_cache_groups
5353 if not isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec)
5354 ]
5355 + max_num_blocks = []
5356 + max_model_len = max(self.max_model_len, self.max_encoder_len)
5357 + for i, kv_cache_group in enumerate(kv_cache_config.kv_cache_groups):
5358 + if isinstance(kv_cache_group.kv_cache_spec, EncoderOnlyAttentionSpec):
5359 + continue
5360 + max_num_blocks_per_req = cdiv(
5361 + max_model_len, block_sizes[i] * get_total_cp_world_size()
5362 + )
5363 + if isinstance(kv_cache_group.kv_cache_spec, MambaSpec):
5364 + mamba_blocks_per_req = (
5365 + max_num_blocks_per_req
5366 + if self.cache_config.enable_prefix_caching
5367 + else 1
5368 + ) + kv_cache_group.kv_cache_spec.num_speculative_blocks
5369 + max_num_blocks_per_req = max(
5370 + max_num_blocks_per_req, mamba_blocks_per_req
5371 + )
5372 + max_num_blocks.append(max_num_blocks_per_req)
5373
5374 if block_sizes != [self.cache_config.block_size] or kernel_block_sizes != [
5375 self.cache_config.block_size
@@ -5333,18 +5381,18 @@ class GPUModelRunner(
5381 )
5382 self.input_batch = InputBatch(
5383 max_num_reqs=self.max_num_reqs,
5336 - max_model_len=max(self.max_model_len, self.max_encoder_len),
5384 + max_model_len=max_model_len,
5385 max_num_batched_tokens=self.max_num_tokens,
5386 device=self.device,
5387 pin_memory=self.pin_memory,
5388 vocab_size=self.model_config.get_vocab_size(),
5389 block_sizes=block_sizes,
5390 kernel_block_sizes=kernel_block_sizes,
5391 + max_num_blocks_per_req=max_num_blocks,
5392 is_spec_decode=bool(self.vllm_config.speculative_config),
5393 logitsprocs=self.input_batch.logitsprocs,
5394 logitsprocs_need_output_token_ids=self.input_batch.logitsprocs_need_output_token_ids,
5395 is_pooling_model=self.is_pooling_model,
5347 - num_speculative_tokens=self.num_spec_tokens,
5396 )
5397
5398 def _allocate_kv_cache_tensors(
vllm/v1/worker/mamba_utils.py new
+232
@@ -0,0 +1,232 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import itertools
4 +from typing import Any
5 +
6 +import torch
7 +import triton
8 +import triton.language as tl
9 +
10 +from vllm.config import CacheConfig
11 +from vllm.model_executor.layers.mamba.mamba_utils import (
12 + MambaStateCopyFunc,
13 +)
14 +from vllm.v1.core.sched.output import SchedulerOutput
15 +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
16 +from vllm.v1.worker.gpu_input_batch import CachedRequestState
17 +from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch
18 +
19 +
20 +@triton.jit
21 +def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr):
22 + pid = tl.program_id(0)
23 +
24 + src_ptr = tl.load(src_ptrs + pid)
25 + dst_ptr = tl.load(dst_ptrs + pid)
26 + size = tl.load(sizes + pid)
27 +
28 + offsets = tl.arange(0, BLOCK_SIZE)
29 + for i in range(0, size, BLOCK_SIZE):
30 + mask = (i + offsets) < size
31 +
32 + curr_src_ptr = (src_ptr + i + offsets).to(tl.pointer_type(tl.uint8))
33 + curr_dst_ptr = (dst_ptr + i + offsets).to(tl.pointer_type(tl.uint8))
34 +
35 + data = tl.load(curr_src_ptr, mask=mask)
36 + tl.store(curr_dst_ptr, data, mask=mask)
37 +
38 +
39 +def batch_memcpy(src_ptrs, dst_ptrs, sizes):
40 + batch = src_ptrs.shape[0]
41 + assert dst_ptrs.shape[0] == batch
42 + assert sizes.shape[0] == batch
43 +
44 + grid = (batch,)
45 + BLOCK_SIZE = 1024
46 + batch_memcpy_kernel[grid](src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=BLOCK_SIZE)
47 +
48 +
49 +def get_mamba_groups(kv_cache_config: KVCacheConfig) -> tuple[list[int], MambaSpec]:
50 + mamba_group_ids: list[int] = []
51 + mamba_specs: list[MambaSpec] = []
52 + for i in range(len(kv_cache_config.kv_cache_groups)):
53 + kv_cache_spec = kv_cache_config.kv_cache_groups[i].kv_cache_spec
54 + if isinstance(kv_cache_spec, MambaSpec):
55 + mamba_group_ids.append(i)
56 + mamba_specs.append(kv_cache_spec)
57 + assert len(mamba_group_ids) > 0, "no mamba layers in the model"
58 + assert all(mamba_specs[0] == spec for spec in mamba_specs)
59 + return mamba_group_ids, mamba_specs[0]
60 +
61 +
62 +def collect_mamba_copy_meta(
63 + src_state_list: list[int],
64 + dest_state_list: list[int],
65 + num_elements_list: list[int],
66 + kv_cache_config: KVCacheConfig,
67 + mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
68 + mamba_group_ids: list[int],
69 + src_block_idx: int,
70 + dest_block_idx: int,
71 + accept_token_bias: int,
72 + req_state: CachedRequestState,
73 + forward_context: dict[str, Any],
74 +):
75 + if src_block_idx == dest_block_idx and accept_token_bias == 0:
76 + return
77 +
78 + for mamba_group_id in mamba_group_ids:
79 + block_ids = req_state.block_ids[mamba_group_id]
80 + dest_block_id = block_ids[dest_block_idx]
81 + layer_names = kv_cache_config.kv_cache_groups[mamba_group_id].layer_names
82 + for layer_name in layer_names:
83 + attention = forward_context[layer_name]
84 + kv_caches: list[torch.Tensor] = attention.kv_cache[0]
85 + for state, state_copy_func in zip(kv_caches, mamba_state_copy_funcs):
86 + copy_spec = state_copy_func(
87 + state, block_ids, src_block_idx, accept_token_bias + 1
88 + )
89 +
90 + src_state_list.append(copy_spec.start_addr)
91 + dest_state_list.append(state[dest_block_id].data_ptr())
92 + num_elements_list.append(copy_spec.num_elements * state.element_size())
93 +
94 +
95 +def do_mamba_copy_block(
96 + src_state_list: list[int],
97 + dest_state_list: list[int],
98 + num_elements_list: list[int],
99 +):
100 + if len(src_state_list) == 0:
101 + return
102 + assert len(src_state_list) == len(dest_state_list)
103 + assert len(src_state_list) == len(num_elements_list)
104 + src_state_ptrs = torch.tensor(src_state_list, device="cuda", dtype=torch.int64)
105 + dst_state_ptrs = torch.tensor(dest_state_list, device="cuda", dtype=torch.int64)
106 + num_elements = torch.tensor(num_elements_list, device="cuda", dtype=torch.int32)
107 +
108 + batch_memcpy(src_state_ptrs, dst_state_ptrs, num_elements)
109 +
110 +
111 +def preprocess_mamba(
112 + scheduler_output: SchedulerOutput,
113 + kv_cache_config: KVCacheConfig,
114 + cache_config: CacheConfig,
115 + mamba_state_idx: dict[str, int],
116 + input_batch: GPUInputBatch,
117 + requests: dict[str, CachedRequestState],
118 + forward_context: dict[str, Any],
119 + mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
120 +):
121 + """
122 + Copy the mamba state of previous step to the last
123 + (1 + num_speculative_blocks) block.
124 + """
125 + mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config)
126 + num_speculative_blocks = mamba_spec.num_speculative_blocks
127 + # TODO(Chen): we need to optimize this function a lot
128 + assert cache_config.enable_prefix_caching
129 + block_size = mamba_spec.block_size
130 + finished_req_ids = scheduler_output.finished_req_ids
131 + preempted_req_ids = scheduler_output.preempted_req_ids or set()
132 + for req_id in itertools.chain(finished_req_ids, preempted_req_ids):
133 + mamba_state_idx.pop(req_id, None)
134 +
135 + src_state_list: list[int] = []
136 + dest_state_list: list[int] = []
137 + num_elements_list: list[int] = []
138 + for i, req_id in enumerate(input_batch.req_ids):
139 + req_state = requests[req_id]
140 + prev_state_idx = mamba_state_idx.get(req_id)
141 + if prev_state_idx is None:
142 + # new / resumed request, no previous state
143 + # if num_computed_tokens is 0, prev_state_idx will be -1
144 + prev_state_idx = (req_state.num_computed_tokens - 1) // block_size
145 +
146 + num_blocks = len(req_state.block_ids[mamba_group_ids[0]])
147 +
148 + # We always save the current running state at the last
149 + # (1 + num_speculative_blocks) block.
150 + # A corner case worth mention here: assume we have block_size = 4 and
151 + # num_speculative_tokens = 2. The request is [A, B, C] and contains 2 draft
152 + # tokens [draft 1, draft 2]. Then we will have:
153 + # Block 0: [A, B, C, draft 1]
154 + # Block 1: [draft 2, TOFILL, TOFILL, TOFILL]
155 + # Block 2: speculative block
156 + # Block 3: speculative block
157 + # And use block 1 to save the running state.
158 + curr_state_idx = num_blocks - 1 - num_speculative_blocks
159 + mamba_state_idx[req_id] = curr_state_idx
160 + if prev_state_idx != -1 and prev_state_idx != curr_state_idx:
161 + collect_mamba_copy_meta(
162 + src_state_list,
163 + dest_state_list,
164 + num_elements_list,
165 + kv_cache_config,
166 + mamba_state_copy_funcs,
167 + mamba_group_ids,
168 + prev_state_idx,
169 + curr_state_idx,
170 + input_batch.num_accepted_tokens_cpu[i] - 1,
171 + req_state,
172 + forward_context,
173 + )
174 + input_batch.num_accepted_tokens_cpu[i] = 1
175 + do_mamba_copy_block(src_state_list, dest_state_list, num_elements_list)
176 +
177 +
178 +def postprocess_mamba(
179 + scheduler_output: SchedulerOutput,
180 + kv_cache_config: KVCacheConfig,
181 + input_batch: GPUInputBatch,
182 + requests: dict[str, CachedRequestState],
183 + mamba_state_idx: dict[str, int],
184 + forward_context: dict[str, Any],
185 + mamba_state_copy_funcs: tuple[MambaStateCopyFunc, ...],
186 +):
187 + """
188 + If a blocks is converted from partial block to full block in this step, copy the
189 + state from the block for running state to the new full block.
190 + """
191 + num_scheduled_tokens_dict = scheduler_output.num_scheduled_tokens
192 + scheduled_spec_decode_tokens_dict = scheduler_output.scheduled_spec_decode_tokens
193 + num_accepted_tokens_cpu = input_batch.num_accepted_tokens_cpu
194 + # NOTE: can be optimized as this function always returns the same result
195 + mamba_group_ids, mamba_spec = get_mamba_groups(kv_cache_config)
196 + src_state_list: list[int] = []
197 + dest_state_list: list[int] = []
198 + num_elements_list: list[int] = []
199 + for i, req_id in enumerate(input_batch.req_ids):
200 + req_state = requests[req_id]
201 + num_computed_tokens = req_state.num_computed_tokens
202 + num_draft_tokens = len(scheduled_spec_decode_tokens_dict.get(req_id, []))
203 + num_scheduled_tokens = num_scheduled_tokens_dict[req_id]
204 + num_accepted_tokens = num_accepted_tokens_cpu[i]
205 + num_tokens_running_state = (
206 + num_computed_tokens + num_scheduled_tokens - num_draft_tokens
207 + )
208 + new_num_computed_tokens = num_tokens_running_state + num_accepted_tokens - 1
209 + aligned_new_computed_tokens = (
210 + new_num_computed_tokens // mamba_spec.block_size * mamba_spec.block_size
211 + )
212 + # TODO: how to ensure all blocks that cache_blocks called are cached here?
213 + if aligned_new_computed_tokens >= num_tokens_running_state:
214 + accept_token_bias = aligned_new_computed_tokens - num_tokens_running_state
215 + src_block_idx = mamba_state_idx[req_id]
216 + dest_block_idx = aligned_new_computed_tokens // mamba_spec.block_size - 1
217 + collect_mamba_copy_meta(
218 + src_state_list,
219 + dest_state_list,
220 + num_elements_list,
221 + kv_cache_config,
222 + mamba_state_copy_funcs,
223 + mamba_group_ids,
224 + src_block_idx,
225 + dest_block_idx,
226 + accept_token_bias,
227 + req_state,
228 + forward_context,
229 + )
230 + if src_block_idx == dest_block_idx:
231 + num_accepted_tokens_cpu[i] = 1
232 + do_mamba_copy_block(src_state_list, dest_state_list, num_elements_list)