A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Mamba CPU Offloading (#44599)
Signed-off-by: varun sundar rabindranath <vsundarr@redhat.com> Co-authored-by: varun sundar rabindranath <vsundarr@redhat.com>
Varun Sundar Rabindranath committed
Jun 12, 2026 at 00:07 UTC
b927004c44e20c8cb86918d500adb431b1661607
2 files changed
+114
-1
tests/v1/kv_connector/unit/test_offloading_connector.py
+88
@@ -554,3 +554,91 @@ def test_fs_tiering_offloading(tmp_path) -> None:
554
finally:
555
subscriber.close()
556
del llm
557
+
558
+
559
+@pytest.mark.skipif(
560
+ not current_platform.is_cuda(),
561
+ reason="HMA mamba-align CPU offload test is CUDA-only",
562
+)
563
+@pytest.mark.parametrize(
564
+ "model,block_size,tp_size",
565
+ [
566
+ # ("Qwen/Qwen3.6-35B-A3B", 1056, 2),
567
+ # ("tiiuae/falcon-mamba-7b", 16, 1),
568
+ ("state-spaces/mamba-1.4b-hf", 16, 1)
569
+ ],
570
+)
571
+def test_mamba_align_cpu_offload(model: str, block_size: int, tp_size: int):
572
+ kv_transfer_config = KVTransferConfig(
573
+ kv_connector="OffloadingConnector",
574
+ kv_role="kv_both",
575
+ kv_connector_extra_config={
576
+ "cpu_bytes_to_use": 4 << 30,
577
+ "block_size": block_size,
578
+ },
579
+ )
580
+ llm = LLM(
581
+ model=model,
582
+ max_model_len=block_size * 10,
583
+ gpu_memory_utilization=0.85,
584
+ tensor_parallel_size=tp_size,
585
+ kv_transfer_config=kv_transfer_config,
586
+ language_model_only=True,
587
+ enable_prefix_caching=True,
588
+ mamba_cache_mode="align",
589
+ disable_hybrid_kv_cache_manager=False,
590
+ )
591
+
592
+ _PROMPT_SIZE: int = block_size * 2
593
+ _PROMPT_TEXT = "Hi. Give me a set of trivia questions and their answers "
594
+
595
+ # build prompt ids to match prompt_size
596
+ tokenizer = llm.get_tokenizer()
597
+ raw_ids: list[int] = tokenizer.encode(_PROMPT_TEXT)
598
+ while len(raw_ids) < _PROMPT_SIZE:
599
+ raw_ids = tokenizer.encode("....") + raw_ids
600
+ initial_ids: list[int] = raw_ids[:_PROMPT_SIZE]
601
+
602
+ sampling_params = SamplingParams(max_tokens=128, temperature=0, ignore_eos=True)
603
+
604
+ failures: list[str] = []
605
+
606
+ def _get_output_str(outputs):
607
+ return outputs[0].outputs[0].text
608
+
609
+ def _verify(llm, prompt, label: str):
610
+ cold_outputs = llm.generate([prompt], sampling_params, use_tqdm=False)
611
+ _wait_for_prefix_cache_reset(llm)
612
+ cpu_outputs = llm.generate([prompt], sampling_params, use_tqdm=False)
613
+
614
+ cold_text = _get_output_str(cold_outputs)
615
+ cpu_text = _get_output_str(cpu_outputs)
616
+ print(f"{label} : cold outputs\n{cold_text}")
617
+ print(f"{label} : cpu outputs\n{cpu_text}")
618
+
619
+ if cold_text != cpu_text:
620
+ failures.append(
621
+ f"{label}: mismatch\n cold: {cold_text!r}\n cpu: {cpu_text!r}"
622
+ )
623
+
624
+ try:
625
+ # Mamba has only a single state. The CPU cache stores are triggered
626
+ # at offload block boundaries. When the prompt is exactly at the boundary,
627
+ # The CPU offload should not load the cached block.
628
+ # This is because we'd use that state to recompute the last token. This
629
+ # does not work for mamba as there is only one KV value and that is for
630
+ # for the token at the boundary.
631
+ # This is fine for other attention types as we have all the necessary
632
+ # token KV values in the hit blocks.
633
+ prompt = TokensPrompt(prompt_token_ids=initial_ids)
634
+ _verify(llm, prompt, "block-boundary-prompt")
635
+
636
+ # Test for prompt token ids at non-block boundaries.
637
+ # Reuse is okay for this case.
638
+ prompt = TokensPrompt(prompt_token_ids=[0] + initial_ids)
639
+ _verify(llm, prompt, "block-mid-prompt")
640
+
641
+ assert not failures, "\n\n".join(failures)
642
+
643
+ finally:
644
+ del llm
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+26
-1
@@ -19,7 +19,7 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
19
_TransferMetricName,
20
)
21
from vllm.logger import init_logger
22
-from vllm.utils.math_utils import cdiv
22
+from vllm.utils.math_utils import cdiv, round_down
23
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
24
from vllm.v1.core.sched.output import SchedulerOutput
25
from vllm.v1.kv_cache_interface import (
@@ -94,6 +94,24 @@ def get_sliding_window_size_in_blocks(
94
return None
95
96
97
+def resolve_mamba_align_size(spec: "OffloadingSpec") -> int | None:
98
+ """Scan all KV cache groups in *spec* and return the single mamba alignment
99
+ size, or None if no group requires mamba alignment.
100
+
101
+ For MambaSpec groups in "align" cache mode the hit window must be rounded
102
+ down to a multiple of the offloaded block size. Asserts that all such
103
+ groups agree on the same value.
104
+ """
105
+ mamba_align_size: int | None = None
106
+ for idx, gpu_block_size in enumerate(spec.gpu_block_size):
107
+ kv_spec = spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec
108
+ if isinstance(kv_spec, MambaSpec) and kv_spec.mamba_cache_mode == "align":
109
+ offload_block_size = gpu_block_size * spec.block_size_factor
110
+ assert mamba_align_size is None or mamba_align_size == offload_block_size
111
+ mamba_align_size = offload_block_size
112
+ return mamba_align_size
113
+
114
+
115
class SchedulerOffloadConfig(NamedTuple):
116
kv_group_configs: tuple[GroupOffloadConfig, ...]
117
block_size_factor: int
@@ -290,6 +308,7 @@ class OffloadingConnectorScheduler:
308
# used by _lookup
309
self._sliding_window_groups: tuple[int, ...] = tuple(sliding_window_groups)
310
self._lookup_groups = tuple(full_attention_groups) + self._sliding_window_groups
311
+ self._mamba_align_size: int | None = resolve_mamba_align_size(spec)
312
313
self._req_status: dict[ReqId, RequestOffloadState] = {}
314
self._current_batch_load_jobs: dict[int, TransferJob] = {}
@@ -408,6 +427,12 @@ class OffloadingConnectorScheduler:
427
# for sliding window attention, we must reduce by 1 to make sure
428
# we still have a hit after reduction
429
max_hit_size_tokens -= 1
430
+ if self._mamba_align_size is not None:
431
+ # Constrain hit-window to the mamba block size.
432
+ max_hit_size_tokens = round_down(
433
+ max_hit_size_tokens, self._mamba_align_size
434
+ )
435
+
436
num_hit_tokens: int = 0
437
defer_lookup = False
438
lookup_groups = self._lookup_groups