From 2cc2e68d1c5652a5bc7ae10422697308e2193359 Mon Sep 17 00:00:00 2001 From: Jiangyun Zhu Date: Tue, 13 Oct 2026 15:58:59 +0800 Subject: [PATCH] [Core][Mamba] Remove leftovers of mamba_cache_mode="all" (#60043) Signed-off-by: zjy0516 Signed-off-by: Jiangyun Zhu Co-authored-by: Claude Opus 5.5 Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/disaggregated.yaml | 2 +- docs/design/hybrid_kv_cache_manager.md | 2 +- docs/features/automatic_prefix_caching.md | 7 +- .../test_kimi_linear_context_parallel.py | 1 - .../configs/Qwen3.8-Flash-Next-FP8-rocm.yaml | 2 +- tests/kernels/mamba/test_mamba_mixer2.py | 10 +- .../test_mamba_eagle_resume_checkpoint.py | 4 +- tests/v1/core/test_contiguous_kv_packing.py | 9 +- tests/v1/core/test_prefix_caching.py | 90 +++------- .../v1/e2e/general/test_mamba_prefix_cache.py | 2 - tests/v1/e2e/test_hybrid_chunked_prefill.py | 2 - tests/v1/e2e/test_replayssm_decode.py | 4 - .../run_mamba_prefix_cache_test.sh | 2 - .../unit/test_offloading_connector.py | 1 - .../unit/test_remote_prefill_lifecycle.py | 3 + tests/v1/simple_kv_offload/test_scheduler.py | 74 ++++++-- vllm/config/cache.py | 28 ++- vllm/config/model.py | 4 - vllm/config/vllm.py | 10 +- vllm/engine/arg_utils.py | 18 +- .../layers/mamba/mamba_mixer2.py | 6 +- .../layers/mamba/ops/causal_conv1d.py | 164 +----------------- vllm/model_executor/models/config.py | 2 +- vllm/model_executor/models/falcon_h1.py | 2 - .../model_executor/models/granitemoehybrid.py | 2 - vllm/model_executor/models/interfaces.py | 25 +-- vllm/model_executor/models/jamba.py | 2 - vllm/model_executor/models/nemotron_h.py | 2 - vllm/model_executor/models/registry.py | 3 - vllm/model_executor/models/zamba2.py | 8 +- vllm/v1/attention/backends/mamba_attn.py | 9 +- vllm/v1/core/single_type_kv_cache_manager.py | 3 + 33 files changed, 161 insertions(+), 344 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index d815ee9268..7f0a351524 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -3523,7 +3523,7 @@ steps: - export MODEL_NAMES='moonshotai/Kimi-Linear-48B-A3B-Instruct' - export PREFILL_BLOCK_SIZE='2048' - export DECODE_BLOCK_SIZE='2048' - - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768' + - export VLLM_SERVE_EXTRA_ARGS='--trust-remote-code,--max-model-len 32768' - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh diff --git a/.buildkite/test_areas/disaggregated.yaml b/.buildkite/test_areas/disaggregated.yaml index b9bd59583b..4821fac0a2 100644 --- a/.buildkite/test_areas/disaggregated.yaml +++ b/.buildkite/test_areas/disaggregated.yaml @@ -426,7 +426,7 @@ steps: # Keep TP=2 prefill and TP=1 decode above both hybrid page-size minima. PREFILL_BLOCK_SIZE: "2048" DECODE_BLOCK_SIZE: "2048" - VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--mamba-cache-mode align,--max-model-len 32768" + VLLM_SERVE_EXTRA_ARGS: "--trust-remote-code,--max-model-len 32768" commands: - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh - bash v1/kv_connector/nixl_integration/run_accuracy_test.sh diff --git a/docs/design/hybrid_kv_cache_manager.md b/docs/design/hybrid_kv_cache_manager.md index 9891c82f97..3984000c7d 100644 --- a/docs/design/hybrid_kv_cache_manager.md +++ b/docs/design/hybrid_kv_cache_manager.md @@ -204,7 +204,7 @@ The second question is the cache eviction policy. For now, we use one LRU queue ### Case 3: mamba models -The prefix caching support of the mamba model is work in progress. Once implemented, models with mamba layer + full attention layer can be supported via the full attention + X algorithm in case 2. +Models with mamba layers + full attention layers support prefix caching via the full attention + X algorithm in case 2, with the Mamba state cached at block boundaries. See [Automatic Prefix Caching](../features/automatic_prefix_caching.md). ## Implementation diff --git a/docs/features/automatic_prefix_caching.md b/docs/features/automatic_prefix_caching.md index cfc8c19acb..bfd87e7a2c 100644 --- a/docs/features/automatic_prefix_caching.md +++ b/docs/features/automatic_prefix_caching.md @@ -15,23 +15,22 @@ Set `enable_prefix_caching=True` in vLLM engine to enable APC. Here is an exampl ## Hybrid Mamba models -Under `--mamba-cache-mode align`, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-shared-prefix-checkpoint` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it. +With prefix caching enabled, Mamba state is stored only on the Mamba block grid, so a prefix-cache hit can resume only at a block boundary. `--enable-mamba-shared-prefix-checkpoint` also stores a checkpoint at the shared-prefix junction, the point where an earlier request with the same prefix stopped. Requests whose shared prefix ends inside a block can then reuse it. This helps when many requests share a long system prompt and then diverge. It is off by default, and takes effect only when all of the following hold: -- `--mamba-cache-mode align` +- prefix caching enabled - EAGLE/MTP speculative decoding on the Mamba group - `--prefix-match-unit` smaller than the Mamba block size - the model does not use multi-module MTP ```bash vllm serve \ - --mamba-cache-mode align \ --prefix-match-unit 64 \ --enable-mamba-shared-prefix-checkpoint ``` -`--prefix-match-unit` is required. It sets the granularity at which prefix-cache keys are computed. When unset it defaults to the greatest common divisor of the prefix-cacheable KV cache group block sizes. Under `align` that is the block size itself, so no sub-block boundary exists and the flag has no effect. +`--prefix-match-unit` is required. It sets the granularity at which prefix-cache keys are computed. When unset it defaults to the greatest common divisor of the prefix-cacheable KV cache group block sizes. For hybrid Mamba models that is the block size itself, so no sub-block boundary exists and the flag has no effect. Choose a value that divides the block size of every prefix-cacheable KV cache group, and that is a multiple of the per-state compression ratio for models that use one, such as sparse MLA. vLLM validates both at startup and names the offending sizes in the error. Read the served block size from the startup log. 64 is a reasonable starting point. diff --git a/tests/distributed/test_kimi_linear_context_parallel.py b/tests/distributed/test_kimi_linear_context_parallel.py index 06348074bb..8bddf6e609 100644 --- a/tests/distributed/test_kimi_linear_context_parallel.py +++ b/tests/distributed/test_kimi_linear_context_parallel.py @@ -180,7 +180,6 @@ def _run_k3_partial_prefix_reuse( block_size=block_size, enable_chunked_prefill=True, enable_prefix_caching=True, - mamba_cache_mode="align", ) as runner: assert ( runner.llm.llm_engine.vllm_config.compilation_config.cudagraph_mode diff --git a/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml b/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml index db43869eb4..c35614cd5f 100644 --- a/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml +++ b/tests/evals/qwen4_exp/configs/Qwen3.8-Flash-Next-FP8-rocm.yaml @@ -23,4 +23,4 @@ generation_config: extra_body: chat_template_kwargs: enable_thinking: true -server_args: "--tensor-parallel-size 4 --mamba-cache-mode align" +server_args: "--tensor-parallel-size 4" diff --git a/tests/kernels/mamba/test_mamba_mixer2.py b/tests/kernels/mamba/test_mamba_mixer2.py index 973e7885c6..db2274c3c9 100644 --- a/tests/kernels/mamba/test_mamba_mixer2.py +++ b/tests/kernels/mamba/test_mamba_mixer2.py @@ -96,15 +96,17 @@ def mixer2_gated_norm_tensor_parallel( gate_states = torch.randn(batch_size, seq_len, hidden_size) # create gated-norm with TP - mixer = Mixer2RMSNormGated( - full_hidden_size=hidden_size, - full_n_groups=n_groups, - ) + with ensure_current_vllm_config(): + mixer = Mixer2RMSNormGated( + full_hidden_size=hidden_size, + full_n_groups=n_groups, + ) mixer.weight.weight_loader(mixer.weight, weight) # load # create gated-norm without TP to compute reference # - utilize mock patching to disable TP when with ( + ensure_current_vllm_config(), unittest.mock.patch( "vllm.model_executor.layers.mamba.mamba_mixer2." "get_tensor_model_parallel_world_size", diff --git a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py index 324aef425b..54ae43e8e2 100644 --- a/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py +++ b/tests/v1/core/prefix_cache/test_mamba_eagle_resume_checkpoint.py @@ -38,9 +38,7 @@ def _manager( num_prefill_lookahead=0, ): init_none_hash(sha256) - config = _make_hybrid_kv_cache_config( - block_size, num_blocks, ["full", "mamba_align"] - ) + config = _make_hybrid_kv_cache_config(block_size, num_blocks, ["full", "mamba"]) if eagle_group is not None: groups = list(config.kv_cache_groups) groups[eagle_group] = replace(groups[eagle_group], is_eagle_group=True) diff --git a/tests/v1/core/test_contiguous_kv_packing.py b/tests/v1/core/test_contiguous_kv_packing.py index 64cd550540..f71fe160d8 100644 --- a/tests/v1/core/test_contiguous_kv_packing.py +++ b/tests/v1/core/test_contiguous_kv_packing.py @@ -278,7 +278,14 @@ class TestCSALinearGrouping: neither divides the hit length nor bounds the (empty) block list.""" config = _shared_layout_config() config.cache_config.enable_prefix_caching = True - groups = get_kv_cache_groups(config, _make_csa_linear_specs()) + config.cache_config.mamba_cache_mode = "align" + specs = { + name: replace(spec, mamba_cache_mode="align") + if isinstance(spec, MambaSpec) + else spec + for name, spec in _make_csa_linear_specs().items() + } + groups = get_kv_cache_groups(config, specs) kv_cache_config = get_kv_cache_config_from_groups( config, groups, available_memory=BYTES_PER_BLOCK * 64 ) diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 836286d127..52455c6083 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1617,6 +1617,7 @@ def make_kv_cache_config_hybrid_model( block_size=block_size, shapes=((1, 1),), dtypes=(torch.float32,), + mamba_cache_mode="align", ) return KVCacheConfig( @@ -1652,6 +1653,7 @@ def make_kv_cache_config_three_types( block_size=block_size, shapes=((1, 1),), dtypes=(torch.float32,), + mamba_cache_mode="align", ) elif third_spec_type == "sliding_window": third_spec = SlidingWindowSpec( @@ -1732,6 +1734,7 @@ def test_prefix_cache_hit_uses_per_group_dcp_geometry(draft_sharded): block_size=block_size, shapes=(1, 1), dtypes=(torch.float32,), + mamba_cache_mode="align", ), ), ], @@ -1757,10 +1760,13 @@ def test_prefix_cache_hit_uses_per_group_dcp_geometry(draft_sharded): req0 = make_request("0", common_token_ids + [99] * 7, block_size, hash_fn) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req0) assert num_computed_tokens == 0 - blocks = manager.allocate_slots( - req0, len(req0.prompt_token_ids), 0, computed_blocks - ) - assert blocks is not None + # Align-mode Mamba keeps a state only where a chunk ends, so end the first + # chunk on the shared prefix. + for num_new_tokens in (len(common_token_ids), 7): + blocks = manager.allocate_slots(req0, num_new_tokens, 0, computed_blocks) + assert blocks is not None + req0.num_computed_tokens += num_new_tokens + computed_blocks = None req1 = make_request("1", common_token_ids + [100] * 5, block_size, hash_fn) computed_blocks, num_computed_tokens, _ = manager.get_computed_blocks(req1) @@ -2285,7 +2291,7 @@ def _make_hybrid_kv_cache_config( - "full": FullAttentionSpec - "sliding_window": SlidingWindowSpec with window=2*block_size - "sliding_window_large": SlidingWindowSpec with window=4*block_size - - "mamba": MambaSpec + - "mamba": MambaSpec in prefix-caching ("align") mode """ spec_map = { @@ -2313,11 +2319,6 @@ def _make_hybrid_kv_cache_config( block_size=block_size, shapes=((1, 1),), dtypes=(torch.float32,), - ), - "mamba_align": lambda: MambaSpec( - block_size=block_size, - shapes=((1, 1),), - dtypes=(torch.float32,), mamba_cache_mode="align", ), } @@ -2427,14 +2428,16 @@ def test_prefill_hybrid_model_combinations(spec_types: list[str]): assert not computed_blocks.blocks[0] # No cache hit initially assert num_computed_tokens == 0 - blocks = manager.allocate_slots( - req0, 55, len(computed_blocks.blocks[0]) * block_size, computed_blocks - ) - assert blocks is not None - # Should have blocks for all groups - assert len(blocks.get_block_ids()) == num_groups - - manager.new_step_starts() + # Align-mode Mamba keeps a state only where a chunk ends, so end the first + # chunk on the common prefix. + for num_new_tokens in (len(common_token_ids), len(unique_token_ids)): + blocks = manager.allocate_slots(req0, num_new_tokens, 0, computed_blocks) + assert blocks is not None + # Should have blocks for all groups + assert len(blocks.get_block_ids()) == num_groups + req0.num_computed_tokens += num_new_tokens + computed_blocks = None + manager.new_step_starts() # Second request: should hit cached blocks for common prefix req1 = make_request("1", common_token_ids + [4] * 5, block_size, hash_fn) @@ -2549,7 +2552,7 @@ def test_prefill_hybrid_model_mamba_align(): num_blocks = 30 kv_cache_config = _make_hybrid_kv_cache_config( - block_size, num_blocks, ["full", "mamba_align"] + block_size, num_blocks, ["full", "mamba"] ) manager = make_kv_cache_manager( kv_cache_config, @@ -2585,7 +2588,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection(): """ block_size = 16 manager = make_kv_cache_manager( - _make_hybrid_kv_cache_config(block_size, 30, ["full", "mamba_align"]), + _make_hybrid_kv_cache_config(block_size, 30, ["full", "mamba"]), max_model_len=8192, enable_caching=True, hash_block_size=block_size, @@ -2660,7 +2663,7 @@ def test_hybrid_model_mamba_align_with_dynamic_draft_tokens(): num_blocks = 30 kv_cache_config = _make_hybrid_kv_cache_config( - block_size, num_blocks, ["full", "mamba_align"] + block_size, num_blocks, ["full", "mamba"] ) manager = KVCacheManager( kv_cache_config, @@ -5896,49 +5899,6 @@ def test_mamba_reachable_block_mask_pins_shared_prefix(): assert retained(0, None) == {14} -def test_mamba_shared_prefix_survives_zero_retention(): - """Manager-level check of the full wiring: a pinned shared-prefix boundary - (``Request.shared_prefix_boundary``, set by the scheduler on Marconi-style - detection) keeps its Mamba state block cached under - ``prefix_cache_retention_interval=0``, which otherwise retains only the - end-of-prompt replay boundary. Without this, a shared prefix (junction - before ``num_prompt``) would be recomputed by every sharing request.""" - block_size = 16 - - # 16-block (256-token) prompt; replay boundary is block 240 // 16 - 1 = 14. - token_ids = [i for i in range(16) for _ in range(block_size)] - - def cached_mamba_blocks(shared_prefix_boundary): - # Fresh manager per scenario so cached blocks don't leak between runs. - manager = make_kv_cache_manager( - _make_hybrid_kv_cache_config(block_size, 100, ["full", "mamba"]), - max_model_len=8192, - enable_caching=True, - hash_block_size=block_size, - retention_interval=0, - ) - req = make_request("r", token_ids, block_size, sha256) - req.shared_prefix_boundary = shared_prefix_boundary - computed_blocks, num_computed, _ = manager.get_computed_blocks(req) - blocks = manager.allocate_slots( - req, len(token_ids), num_computed, computed_blocks - ) - assert blocks is not None - pool = manager.block_pool - return { - i - for i in range(16) - if pool.get_cached_block(req.block_hashes[i], kv_cache_group_ids=[1]) - is not None - } - - # Without a pinned boundary, retention=0 keeps only the replay boundary (14). - assert cached_mamba_blocks(0) == {14} - # Pinning the shared prefix at token 96 (state block 5) retains it too, so a - # later request sharing that prefix can hit the Mamba state. - assert cached_mamba_blocks(96) == {5, 14} - - def test_mamba_shared_prefix_reuse_under_zero_retention(): """Full cross-request Marconi flow: a partial shared prefix cached by the detecting request must stay reusable by a later request under @@ -5949,7 +5909,7 @@ def test_mamba_shared_prefix_reuse_under_zero_retention(): def last_req_hit(retention, pin): manager = make_kv_cache_manager( - _make_hybrid_kv_cache_config(block_size, 200, ["full", "mamba_align"]), + _make_hybrid_kv_cache_config(block_size, 200, ["full", "mamba"]), max_model_len=8192, enable_caching=True, hash_block_size=block_size, diff --git a/tests/v1/e2e/general/test_mamba_prefix_cache.py b/tests/v1/e2e/general/test_mamba_prefix_cache.py index 50f3473708..6c208e19fb 100644 --- a/tests/v1/e2e/general/test_mamba_prefix_cache.py +++ b/tests/v1/e2e/general/test_mamba_prefix_cache.py @@ -898,7 +898,6 @@ def _run_mamba_prefix_cache_mrv1( load_format="dummy", enable_prefix_caching=True, block_size=BLOCK_SIZE, - mamba_cache_mode="align", speculative_config={ "method": "qwen3_next_mtp", "num_speculative_tokens": num_speculative_tokens, @@ -1158,7 +1157,6 @@ def _run_mamba_prefix_cache_mrv2( skip_tokenizer_init=True, enable_prefix_caching=True, block_size=BLOCK_SIZE, - mamba_cache_mode="align", speculative_config={ "method": "qwen3_next_mtp", "num_speculative_tokens": num_speculative_tokens, diff --git a/tests/v1/e2e/test_hybrid_chunked_prefill.py b/tests/v1/e2e/test_hybrid_chunked_prefill.py index 66c8d5b9b7..c63a4ebc78 100644 --- a/tests/v1/e2e/test_hybrid_chunked_prefill.py +++ b/tests/v1/e2e/test_hybrid_chunked_prefill.py @@ -88,7 +88,6 @@ def test_mtp_speculative_mixed_batch_short_prefill( trust_remote_code=True, enable_chunked_prefill=True, enable_prefix_caching=enable_prefix_caching, - mamba_cache_mode="align" if enable_prefix_caching else "none", ) as llm: sampling_params = SamplingParams( temperature=0.0, @@ -161,7 +160,6 @@ def test_qwen_mtp_mamba_prefix_cache_hit_is_bounded(vllm_runner, monkeypatch): enforce_eager=True, enable_chunked_prefill=True, enable_prefix_caching=True, - mamba_cache_mode="align", speculative_config={"method": "mtp", "num_speculative_tokens": 2}, ) as runner: llm = runner.get_llm() diff --git a/tests/v1/e2e/test_replayssm_decode.py b/tests/v1/e2e/test_replayssm_decode.py index 89797db4cd..f75151414a 100644 --- a/tests/v1/e2e/test_replayssm_decode.py +++ b/tests/v1/e2e/test_replayssm_decode.py @@ -60,7 +60,6 @@ def _check_replayssm_parity( max_model_len=1024, trust_remote_code=True, enable_prefix_caching=False, - mamba_cache_mode="none", tensor_parallel_size=tensor_parallel_size, mamba_backend=mamba_backend, ) @@ -130,7 +129,6 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline( max_model_len=1024, trust_remote_code=True, enable_prefix_caching=False, - mamba_cache_mode="none", speculative_config={ "method": "ngram", "num_speculative_tokens": 3, @@ -182,7 +180,6 @@ def test_replayssm_flashinfer_mtp(vllm_runner, monkeypatch, use_v2_model_runner) max_model_len=1024, trust_remote_code=True, enable_prefix_caching=False, - mamba_cache_mode="none", mamba_backend="flashinfer", tensor_parallel_size=2, disable_log_stats=False, @@ -250,7 +247,6 @@ def _check_replayssm_prefix_caching_parity( trust_remote_code=True, enable_prefix_caching=True, enable_chunked_prefill=True, - mamba_cache_mode="align", disable_log_stats=False, # required for llm.get_metrics() tensor_parallel_size=tensor_parallel_size, ) diff --git a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh index c6f79f3721..6615a5c103 100755 --- a/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_mamba_prefix_cache_test.sh @@ -60,7 +60,6 @@ vllm serve "$MODEL" \ --block-size 128 \ --trust-remote-code \ --enable-prefix-caching \ - --mamba-cache-mode align \ --kv-transfer-config "$KV_CONFIG" \ "${EXTRA_ARGS[@]}" & @@ -78,7 +77,6 @@ vllm serve "$MODEL" \ --block-size 128 \ --trust-remote-code \ --enable-prefix-caching \ - --mamba-cache-mode align \ --kv-transfer-config "$KV_CONFIG" \ "${EXTRA_ARGS[@]}" & diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index 76d2197469..53e7378a7b 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/tests/v1/kv_connector/unit/test_offloading_connector.py @@ -583,7 +583,6 @@ def test_mamba_cpu_offload_boundary(model: str, block_size: int, tp_size: int): kv_transfer_config=kv_transfer_config, language_model_only=True, enable_prefix_caching=True, - mamba_cache_mode="align", # Use lossless state storage so this exact-equality regression isolates # offload boundary selection from low-precision Mamba checkpointing. mamba_cache_dtype="float32", diff --git a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py index b72de811c0..74de3c3a8c 100644 --- a/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py +++ b/tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py @@ -632,6 +632,8 @@ def test_p_side_chunked_prefill_mamba(mock_platform): block_size=BLOCK_SIZE, ) vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False + # mamba_cache_mode "none" only runs without prefix caching. + vllm_config.cache_config.enable_prefix_caching = False kv_cache_config = make_kv_cache_config( block_size=BLOCK_SIZE, @@ -855,6 +857,7 @@ def test_async_load_does_not_double_reserve_mamba_spec_blocks( disable_hybrid_kv_cache_manager=False, ) vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode + vllm_config.cache_config.enable_prefix_caching = mamba_cache_mode == "align" kv_cache_config = make_kv_cache_config( block_size, mamba_enabled=True, diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 9f818c866f..527f7bee05 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -49,6 +49,7 @@ from vllm.v1.kv_cache_interface import ( FullAttentionSpec, KVCacheConfig, KVCacheGroupSpec, + KVCacheSpec, KVCacheTensor, MambaSpec, SlidingWindowSpec, @@ -2355,8 +2356,6 @@ def _make_hybrid_attention_mamba_scheduler( enable_kv_cache_events: bool = False, ) -> SchedulerFixture: """Build a scheduler for one attention group plus one Mamba group.""" - scheduler_block_size = scheduler_block_size or block_size - hash_block_size = hash_block_size or block_size attention_spec = FullAttentionSpec( block_size=attention_block_size, num_kv_heads=NUM_KV_HEADS, @@ -2369,9 +2368,38 @@ def _make_hybrid_attention_mamba_scheduler( dtypes=(torch.float32,), mamba_cache_mode=mamba_cache_mode, ) + return _make_two_group_scheduler( + attention_spec, + ("mamba", mamba_spec), + num_cpu_blocks=num_cpu_blocks, + num_gpu_blocks=num_gpu_blocks, + scheduler_block_size=scheduler_block_size or block_size, + hash_block_size=hash_block_size or block_size, + dcp_world_size=dcp_world_size, + lazy=lazy, + mamba_cache_mode=mamba_cache_mode, + enable_kv_cache_events=enable_kv_cache_events, + ) + + +def _make_two_group_scheduler( + attention_spec: FullAttentionSpec, + second_group: tuple[str, KVCacheSpec], + *, + num_cpu_blocks: int, + num_gpu_blocks: int, + scheduler_block_size: int, + hash_block_size: int, + dcp_world_size: int, + lazy: bool = False, + mamba_cache_mode: MambaCacheMode = "align", + enable_kv_cache_events: bool = False, +) -> SchedulerFixture: + """Build a scheduler for one attention group plus one named second group.""" + second_name, second_spec = second_group groups = [ KVCacheGroupSpec(["attention"], attention_spec), - KVCacheGroupSpec(["mamba"], mamba_spec), + KVCacheGroupSpec([second_name], second_spec), ] tensors = [ KVCacheTensor( @@ -2380,7 +2408,7 @@ def _make_hybrid_attention_mamba_scheduler( layer_stride=spec.page_size_bytes * num_gpu_blocks, block_stride=spec.page_size_bytes, ) - for group, spec in zip(groups, (attention_spec, mamba_spec)) + for group, spec in zip(groups, (attention_spec, second_spec)) ] kv_cache_config = KVCacheConfig( num_blocks=num_gpu_blocks, @@ -2428,34 +2456,48 @@ def test_hybrid_store_uses_resolved_group_block_sizes() -> None: silently offloads only a fraction of them. """ attention_block_size = BLOCK_SIZE - mamba_block_size = 4 * BLOCK_SIZE + replicated_block_size = 4 * BLOCK_SIZE dcp_world_size = 2 - fix = _make_hybrid_attention_mamba_scheduler( + fix = _make_two_group_scheduler( + FullAttentionSpec( + block_size=attention_block_size, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + ), + ( + "replicated_attention", + FullAttentionSpec( + block_size=replicated_block_size, + num_kv_heads=NUM_KV_HEADS, + head_size=HEAD_SIZE, + dtype=DTYPE, + dcp_sharded=False, + ), + ), num_cpu_blocks=32, num_gpu_blocks=32, - attention_block_size=attention_block_size, - block_size=mamba_block_size, + scheduler_block_size=replicated_block_size, # Every prefix-cacheable group's resolved block size must be a multiple # of the hash block, so it cannot exceed the smallest of them. hash_block_size=BLOCK_SIZE, dcp_world_size=dcp_world_size, - mamba_cache_mode="none", ) sched = fix.scheduler gpu_pool = fix.gpu_block_pool - attention_size, mamba_size = sched.group_block_sizes + attention_size, replicated_size = sched.group_block_sizes assert attention_size == attention_block_size * dcp_world_size - assert mamba_size == mamba_block_size + assert replicated_size == replicated_block_size confirmed = 2 * sched.block_size req = _make_cp_request(num_blocks=8, virtual_block_size=BLOCK_SIZE) attention_blocks = _allocate_cp_gpu_blocks( gpu_pool, req, confirmed // attention_size, attention_size, group_id=0 ) - mamba_blocks = _allocate_cp_gpu_blocks( - gpu_pool, req, confirmed // mamba_size, mamba_size, group_id=1 + replicated_blocks = _allocate_cp_gpu_blocks( + gpu_pool, req, confirmed // replicated_size, replicated_size, group_id=1 ) - kv_blocks = KVCacheBlocks(blocks=(attention_blocks, mamba_blocks)) + kv_blocks = KVCacheBlocks(blocks=(attention_blocks, replicated_blocks)) sched.update_state_after_alloc(req, kv_blocks, num_external_tokens=0) sched.build_connector_meta( make_scheduler_output( @@ -2476,11 +2518,11 @@ def test_hybrid_store_uses_resolved_group_block_sizes() -> None: # count. assert sched._reqs_to_store[req.request_id].num_stored_blocks == [ confirmed // attention_size, - confirmed // mamba_size, + confirmed // replicated_size, ] assert set(meta.store_gpu_blocks) == { *(b.block_id for b in attention_blocks), - *(b.block_id for b in mamba_blocks), + *(b.block_id for b in replicated_blocks), } diff --git a/vllm/config/cache.py b/vllm/config/cache.py index 5eb26726c7..9bb7166215 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -176,8 +176,7 @@ class CacheConfig: page.""" mamba_block_size: int | None = Field(default=None, gt=0) """Size of a contiguous cache block in number of tokens for mamba cache. - Can be set only when prefix caching is enabled. - Value must be a multiple of 8 to align with causal_conv1d kernel.""" + Can be set only when prefix caching is enabled.""" mamba_cache_dtype: MambaDType = "auto" """The data type to use for the Mamba cache (both the conv as well as the ssm state). If set to 'auto', the data type will be inferred from the model @@ -187,19 +186,19 @@ class CacheConfig: still be controlled by mamba_cache_dtype). If set to 'auto', the data type for the ssm state will be determined by mamba_cache_dtype.""" mamba_cache_mode: MambaCacheMode = "none" - """The cache strategy for Mamba layers: + """The cache strategy for Mamba layers, derived from `enable_prefix_caching`: - - "none": set when prefix caching is disabled. - - "align": only cache the mamba state of the last token of each scheduler step and - when the token is at position i * block_size. This is the default when prefix - caching is enabled. + - "none": prefix caching is disabled. + - "align": prefix caching is enabled. Only cache the mamba state of the last + token of each scheduler step and when the token is at position + i * block_size. """ enable_mamba_shared_prefix_checkpoint: bool = False """Also register a Mamba "align" checkpoint at the shared-prefix junction -- where an EAGLE/MTP sibling was observed to resume -- instead of only at the - prompt tail. Off by default; only takes effect with `mamba_cache_mode` - "align", EAGLE on the Mamba group, and a prefix match unit smaller than the - Mamba block size.""" + prompt tail. Off by default; only takes effect with prefix caching enabled, + EAGLE on the Mamba group, and a prefix match unit smaller than the Mamba + block size.""" replayssm_buffer_len: int = Field(default=16, gt=0) """ReplaySSM logical history length B for Mamba2. Triton uses B physical rows and FlashInfer uses B+T, where T is the target verification length. @@ -207,11 +206,10 @@ class CacheConfig: use_replayssm: bool = False """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip the per-step full-state store, writing the checkpoint back only on flush. - Requires mamba_cache_mode 'none' or 'align' (prefix caching) and the Triton - or FlashInfer mamba backend. Mamba2 speculative decode requires FlashInfer - and mamba_cache_mode 'none'. In align mode flushes are most efficient when - mamba_block_size is a multiple of replayssm_buffer_len, but this is not - required.""" + Requires the Triton or FlashInfer mamba backend; FlashInfer does not support + prefix caching yet. Mamba2 speculative decode requires FlashInfer. With + prefix caching, flushes are most efficient when mamba_block_size is a + multiple of replayssm_buffer_len, but this is not required.""" use_kda_recoverssm: bool = field(default=False, init=False) """Whether Kimi-K3 KDA uses RecoverSSM speculative decode.""" diff --git a/vllm/config/model.py b/vllm/config/model.py index 87759f4606..c4016e82bc 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -2038,10 +2038,6 @@ class ModelConfig: def has_inner_state(self): return self._model_info.has_inner_state - @property - def supports_mamba_prefix_caching(self) -> bool: - return self._model_info.supports_mamba_prefix_caching - @property def supports_replayssm(self) -> bool: return self._model_info.supports_replayssm diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index f27c8418ab..d6e834861a 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -3360,7 +3360,7 @@ class VllmConfig: and not self.use_v2_model_runner ): raise ValueError( - "RecoverSSM with align mode requires VLLM_USE_V2_MODEL_RUNNER=1" + "RecoverSSM with prefix caching requires VLLM_USE_V2_MODEL_RUNNER=1" ) if self.parallel_config.pipeline_parallel_size > 1: raise ValueError( @@ -3371,8 +3371,8 @@ class VllmConfig: elif use_mamba_replayssm_spec: if self.cache_config.mamba_cache_mode != "none": raise ValueError( - "FlashInfer ReplaySSM speculative decoding requires " - "--mamba-cache-mode none" + "FlashInfer ReplaySSM speculative decoding does not support " + "prefix caching; use --no-enable-prefix-caching" ) query_len = 1 + self.num_speculative_tokens if self.cache_config.replayssm_buffer_len < query_len: @@ -3390,8 +3390,8 @@ class VllmConfig: elif self.mamba_config.backend == MambaBackendEnum.FLASHINFER: if self.cache_config.mamba_cache_mode == "align": raise ValueError( - "FlashInfer ReplaySSM does not support " - "--mamba-cache-mode align yet; use none" + "FlashInfer ReplaySSM does not support prefix caching yet; " + "use --no-enable-prefix-caching" ) elif self.mamba_config.backend != MambaBackendEnum.TRITON: raise ValueError( diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 86199e5f00..c9aa433162 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -453,6 +453,12 @@ _LOG_CONFIG_FILE_DEPRECATION_MESSAGE = ( ) +_MAMBA_CACHE_MODE_DEPRECATION_MESSAGE = ( + "--mamba-cache-mode is deprecated and will be removed in v0.33.0. " + "It has no effect: Mamba prefix caching follows --enable-prefix-caching." +) + + _DISABLE_CASCADE_ATTN_DEPRECATION_MESSAGE = ( "--disable-cascade-attn is deprecated and has no effect since cascade " "attention has been removed. It will be removed in v0.33.0." @@ -798,7 +804,7 @@ class EngineArgs: mamba_ssm_cache_dtype: MambaDType = CacheConfig.mamba_ssm_cache_dtype mamba_block_size: int | None = get_field(CacheConfig, "mamba_block_size") prefix_match_unit: int | None = get_field(CacheConfig, "prefix_match_unit") - mamba_cache_mode: MambaCacheMode = CacheConfig.mamba_cache_mode + mamba_cache_mode: MambaCacheMode | None = None enable_mamba_shared_prefix_checkpoint: bool = ( CacheConfig.enable_mamba_shared_prefix_checkpoint ) @@ -1403,7 +1409,12 @@ class EngineArgs: "--prefix-match-unit", **cache_kwargs["prefix_match_unit"] ) cache_group.add_argument( - "--mamba-cache-mode", **cache_kwargs["mamba_cache_mode"] + "--mamba-cache-mode", + **{ + **cache_kwargs["mamba_cache_mode"], + "default": None, + "help": _MAMBA_CACHE_MODE_DEPRECATION_MESSAGE, + }, ) cache_group.add_argument( "--enable-mamba-shared-prefix-checkpoint", @@ -2263,6 +2274,8 @@ class EngineArgs: assert self.enable_prefix_caching is not None, ( "enable_prefix_caching must be set by this point" ) + if self.mamba_cache_mode is not None: + logger.warning_once(_MAMBA_CACHE_MODE_DEPRECATION_MESSAGE) cache_config = CacheConfig( block_size=self.block_size, # type: ignore[arg-type] @@ -2282,7 +2295,6 @@ class EngineArgs: mamba_ssm_cache_dtype=self.mamba_ssm_cache_dtype, mamba_block_size=self.mamba_block_size, prefix_match_unit=self.prefix_match_unit, - mamba_cache_mode=self.mamba_cache_mode, enable_mamba_shared_prefix_checkpoint=( self.enable_mamba_shared_prefix_checkpoint ), diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index bfa1103bb3..14c483122a 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -812,10 +812,8 @@ class MambaMixer2(MambaBase, PluggableLayer): # - It will read the initial states for every sequence, # that has "has_initial_states_p" == True, # from "cache_indices", using "state_indices_tensor_p". - # - It updates the "conv_state" cache in positions pointed - # to by "state_indices_tensor_p". - # In particular, it will always write the state at the - # sequence end. + # - It writes the state at the sequence end to the "conv_state" + # cache in positions pointed to by "state_indices_tensor_p". x = hidden_states_B_C_p.transpose( 0, 1 ) # this is the form that causal-conv see diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index e8254e4f34..2094df902a 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -20,17 +20,11 @@ def _causal_conv1d_fwd_kernel( # continuous batching w_ptr, # (dim, width) bias_ptr, initial_states_ptr, # conv_states_ptr - cache_indices_ptr, # (batch, n_blocks + padding) The second dimension contains - # the block indices relevant for each sequence - # plus potential 0-padding at the beginning and at the end + cache_indices_ptr, # (batch,) has_initial_states_ptr, query_start_loc_ptr, batch_ptr, token_chunk_offset_ptr, - block_idx_first_scheduled_token, # (batch,) - block_idx_last_scheduled_token, # (batch,) - initial_state_idx, # (batch,) - num_computed_tokens, # (batch,) o_ptr, # (dim, seqlen) - actually pointing to x_ptr # Matrix dimensions dim: tl.constexpr, @@ -46,7 +40,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching stride_cache_indices: tl.constexpr, stride_o_dim: tl.constexpr, stride_o_token: tl.int64, - stride_block_m: tl.constexpr, # Stride block to align divided by BLOCK_M # others pad_slot_id: tl.constexpr, null_block_id: tl.constexpr, @@ -54,7 +47,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching HAS_BIAS: tl.constexpr, KERNEL_WIDTH: tl.constexpr, SILU_ACTIVATION: tl.constexpr, - IS_APC_ENABLED: tl.constexpr, HAS_NULL_BLOCK: tl.constexpr, NP2_STATELEN: tl.constexpr, BLOCK_M: tl.constexpr, @@ -93,40 +85,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching # find the actual sequence length seqlen = sequence_end_index - sequence_start_index - B_size: tl.constexpr = stride_block_m * BLOCK_M - - if IS_APC_ENABLED: - # Handle the case if prefix caching is enabled. - # In particular, if prefix caching is enabled, the program write additional cache states to "cache_indices_ptr" - - # Get the length of the completed sequence so far and compute the offset. - current_first_index = tl.load(block_idx_first_scheduled_token + idx_seq) - current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) - sequence_completed_index = tl.load(num_computed_tokens + idx_seq) - - # Compute the offset where the first stride_block_m-aligned first full block is - # Value in "token-space" - sequence_completed_offset_token = sequence_completed_index % B_size - seq_completed_offset = B_size - sequence_completed_offset_token - seq_end_offset = (seqlen - seq_completed_offset) % B_size - last_full_block_token_index = sequence_end_index - seq_end_offset - # If the sequence without the sequence_offset_index is stride_cache_chunk-aligned, then the last full chunk is the second-to-last one - if seq_end_offset == 0: - last_full_block_token_index = last_full_block_token_index - B_size - - # Get the number of blocks to be filled for the current sequence - # If n_block_to_fill = 0, then only the state at the sequence end is stored - n_block_to_fill = current_last_index - current_first_index - - # Get the index of the init block - conv_state_init_index = tl.load(initial_state_idx + idx_seq) - else: - n_block_to_fill = 0 - current_last_index = 0 - conv_state_init_index = 0 - current_first_index = 0 - last_full_block_token_index = 0 - token_offset = BLOCK_M * chunk_offset segment_len = min(BLOCK_M, seqlen - token_offset) @@ -137,7 +95,7 @@ def _causal_conv1d_fwd_kernel( # continuous batching # cache_idx conv_states_input_coord = tl.load( - conv_state_indices_ptr + idx_seq * stride_cache_indices + conv_state_init_index + conv_state_indices_ptr + idx_seq * stride_cache_indices ).to(tl.int64) if HAS_NULL_BLOCK: # noqa @@ -224,20 +182,10 @@ def _causal_conv1d_fwd_kernel( # continuous batching loaded_x = tl.load(x_ptrs, mask_x, 0.0) idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] - # Compute the offset where the last block should be written in the conv_states - conv_states_output_coord = tl.load( - conv_state_indices_ptr - + idx_seq * stride_cache_indices - + current_last_index - ).to(tl.int64) - conv_states_ptrs_target = ( - conv_states_ptr - + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq - + (idx_feats * stride_conv_state_dim) - )[None, :] + ( # [BLOCK_N,] - idx_tokens_conv * stride_conv_state_tok - )[:, None] + conv_states_base[None, :] + + (idx_tokens_conv * stride_conv_state_tok)[:, None] + ) # [BLOCK_M, BLOCK_N] mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] tl.debug_barrier() # NOTE: use this due to bug in Triton compiler @@ -346,52 +294,6 @@ def _causal_conv1d_fwd_kernel( # continuous batching conv_states_ptrs = prior_tokens - 3 * stride_x_token # [BLOCK_N] col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier=".ca") - # Store intermediate states aligned with stride_block_m - # The additional states are cached starting from the last stride_block_m. - # For example: - # If n_block_to_fill = 0, then only the state at the sequence end is cached and the process below is not involved. - # If n_block_to_fill > 0, then the states at the sequence end and at the n_block_to_fill-last - # stride_block_m are cached. - # For example chunk_offset = n_block_to_fill stores the state at last_full_block - if (chunk_offset - 1) < n_block_to_fill: - # Store the states at the chunk boundaries from the start of the sequence - idx_tokens_last = ( - last_full_block_token_index - - (n_block_to_fill - chunk_offset) * B_size - - state_len - ) + tl.arange(0, NP2_STATELEN) # [BLOCK_M] - x_ptrs = ( - x_ptr - + (idx_tokens_last * stride_x_token)[:, None] - + (idx_feats * stride_x_dim)[None, :] - ) # [BLOCK_M,BLOCK_N,] - - mask_x = (idx_tokens_last >= 0)[:, None] & (idx_feats < dim)[ - None, : - ] # token-index # token-index # feature-index - loaded_x = tl.load(x_ptrs, mask_x, 0.0) - idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M] - - # cache_idx - conv_states_output_coord = tl.load( - conv_state_indices_ptr - + idx_seq * stride_cache_indices - + current_first_index - + (chunk_offset - 1) - ).to(tl.int64) - - conv_states_ptrs_target = ( - conv_states_ptr - + (conv_states_output_coord * stride_conv_state_seq) # Offset from seq - + (idx_feats * stride_conv_state_dim) - )[None, :] + ( # [BLOCK_N,] - idx_tokens_conv * stride_conv_state_tok - )[:, None] - - mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats < dim)[None, :] - tl.debug_barrier() # NOTE: use this due to bug in Triton compiler - tl.store(conv_states_ptrs_target, loaded_x, mask) - if HAS_BIAS: bias = bias_ptr + idx_feats mask_bias = idx_feats < dim @@ -489,11 +391,6 @@ def causal_conv1d_fn( activation: str | None = "silu", pad_slot_id: int = PAD_SLOT_ID, null_block_id: int = NULL_BLOCK_ID, - block_idx_first_scheduled_token: torch.Tensor | None = None, - block_idx_last_scheduled_token: torch.Tensor | None = None, - initial_state_idx: torch.Tensor | None = None, - num_computed_tokens: torch.Tensor | None = None, - block_size_to_align=0, metadata=None, validate_data=False, ): @@ -536,16 +433,6 @@ def causal_conv1d_fn( for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id] in this case, the kernel will not process entries at indices 0 and 3 - block_idx_first_scheduled_token: (batch,), dtype int32 - The pointer into cache_indices, where the first cache block to be filled is located. - block_idx_last_scheduled_token: (batch,), dtype int32 - The pointer into cache_indices, where the last cache block to be filled is located. - initial_state_idx: (batch,), dtype int32 - The pointer into cache_indices, where the cache block containing the initial state is located. - num_computed_tokens: (batch,), dtype int32 - The number of tokens already completed for each sequence - block_size_to_align: int - The block size to align the cached states to out: same shape as `x` """ if isinstance(activation, bool) and activation: @@ -631,12 +518,6 @@ def causal_conv1d_fn( assert weight.stride(1) == 1 assert (dim, width) == weight.shape assert is_channel_last, "Need to run in channel-last layout" - if block_size_to_align is not None and block_size_to_align > 0: - assert (block_size_to_align % BLOCK_M) == 0, ( - "The mamba block size needs to be divisible by the BLOCK_M" - ) - else: - block_size_to_align = BLOCK_M if metadata is None: @@ -720,10 +601,6 @@ def causal_conv1d_fn( query_start_loc, batch_ptr, token_chunk_offset_ptr, - block_idx_first_scheduled_token, - block_idx_last_scheduled_token, - initial_state_idx, - num_computed_tokens, out, # Matrix dimensions dim, @@ -739,7 +616,6 @@ def causal_conv1d_fn( stride_cache_indices, stride_o_dim, stride_o_token, - block_size_to_align // BLOCK_M, # others pad_slot_id, null_block_id, @@ -747,7 +623,6 @@ def causal_conv1d_fn( HAS_BIAS=bias is not None, KERNEL_WIDTH=width, SILU_ACTIVATION=activation in ["silu", "swish"], - IS_APC_ENABLED=block_idx_last_scheduled_token is not None, HAS_NULL_BLOCK=null_block_id is not None, NP2_STATELEN=np2_statelen, # launch_cooperative_grid=True @@ -769,8 +644,6 @@ def _causal_conv1d_update_kernel( conv_state_indices_ptr, num_accepted_tokens_ptr, query_start_loc_ptr, # (batch + 1) - block_idx_last_scheduled_token, # (batch,) - initial_state_idx, # (batch,) o_ptr, # (batch, dim, seqlen) # Matrix dimensions batch: int, @@ -798,7 +671,6 @@ def _causal_conv1d_update_kernel( KERNEL_WIDTH: tl.constexpr, SILU_ACTIVATION: tl.constexpr, IS_VARLEN: tl.constexpr, - IS_APC_ENABLED: tl.constexpr, IS_SPEC_DECODING: tl.constexpr, NP2_STATELEN: tl.constexpr, HAS_NULL_BLOCK: tl.constexpr, @@ -818,17 +690,9 @@ def _causal_conv1d_update_kernel( # [BLOCK_N,] elements along the feature-dimension (channel) idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N) - if IS_APC_ENABLED: - # Get the state from the initial_state_idx - conv_state_init = tl.load(initial_state_idx + idx_seq) - current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq) - else: - conv_state_init = 0 - current_last_index = 0 - # cache_idx conv_states_input_coord = tl.load( - conv_state_indices_ptr + idx_seq * stride_state_indices + conv_state_init + conv_state_indices_ptr + idx_seq * stride_state_indices ).to(tl.int64) if HAS_NULL_BLOCK: # noqa @@ -941,14 +805,9 @@ def _causal_conv1d_update_kernel( new_conv_state = tl.where(mask, conv_state, loaded_x) - # Get the state from the initial_state_idx - # cache_idx - conv_states_offset = tl.load( - conv_state_indices_ptr + idx_seq * stride_state_indices + current_last_index - ).to(tl.int64) conv_state_ptrs_target = ( conv_state_ptr - + (conv_states_offset * stride_conv_state_seq) # Offset from seq + + (conv_states_input_coord * stride_conv_state_seq) # Offset from seq + (idx_feats * stride_conv_state_dim) )[None, :] + ( # [BLOCK_N,] idx_tokens * stride_conv_state_tok @@ -1104,8 +963,6 @@ def causal_conv1d_update( query_start_loc: torch.Tensor | None = None, max_query_len: int = -1, null_block_id: int = NULL_BLOCK_ID, - block_idx_last_scheduled_token: torch.Tensor | None = None, - initial_state_idx: torch.Tensor | None = None, validate_data=False, out: torch.Tensor | None = None, ): @@ -1123,10 +980,6 @@ def causal_conv1d_update( If not None, the conv_state is a larger tensor along the batch dim, and we are selecting the batch coords specified by conv_state_indices. Useful for a continuous batching scenario. - block_idx_last_scheduled_token: (batch,), dtype int32 - The pointer into conv_state_indices, where the last cache block to be filled is located. - initial_state_idx: (batch,), dtype int32 - The pointer into conv_state_indices, where the cache block containing the initial state is located. num_accepted_tokens: (batch,), dtype int32 If not None, it indicates the number of accepted tokens for each sequence in the batch. @@ -1237,8 +1090,6 @@ def causal_conv1d_update( conv_state_indices, num_accepted_tokens, query_start_loc, - block_idx_last_scheduled_token, - initial_state_idx, out, # Matrix dimensions batch, @@ -1266,7 +1117,6 @@ def causal_conv1d_update( KERNEL_WIDTH=width, SILU_ACTIVATION=activation in ["silu", "swish"], IS_VARLEN=query_start_loc is not None, - IS_APC_ENABLED=block_idx_last_scheduled_token is not None, IS_SPEC_DECODING=num_accepted_tokens is not None, NP2_STATELEN=np2_statelen, HAS_NULL_BLOCK=null_block_id is not None, diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index d64e51bfa8..69ca643a3a 100644 --- a/vllm/model_executor/models/config.py +++ b/vllm/model_executor/models/config.py @@ -692,7 +692,7 @@ class MambaModelConfig(VerifyAndUpdateConfig): model_config.architecture, ) assert vllm_config.scheduler_config.enable_chunked_prefill, ( - "Chunked prefill is required for mamba cache mode 'align'." + "Chunked prefill is required for Mamba prefix caching." ) # By default, mamba block size will be set to max_model_len (see # below). When enabling prefix caching, we align mamba block size diff --git a/vllm/model_executor/models/falcon_h1.py b/vllm/model_executor/models/falcon_h1.py index 203a3f3ec3..7aff819dfc 100644 --- a/vllm/model_executor/models/falcon_h1.py +++ b/vllm/model_executor/models/falcon_h1.py @@ -42,7 +42,6 @@ from .interfaces import ( HasInnerState, IsHybrid, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, ) from .utils import ( @@ -499,7 +498,6 @@ class FalconH1ForCausalLM( SupportsLoRA, SupportsPP, IsHybrid, - SupportsMambaPrefixCaching, ): packed_modules_mapping = { "qkv_proj": ["q_proj", "k_proj", "v_proj"], diff --git a/vllm/model_executor/models/granitemoehybrid.py b/vllm/model_executor/models/granitemoehybrid.py index 979d272c13..040f7453fb 100644 --- a/vllm/model_executor/models/granitemoehybrid.py +++ b/vllm/model_executor/models/granitemoehybrid.py @@ -42,7 +42,6 @@ from .interfaces import ( HasInnerState, IsHybrid, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, SupportsQuant, ) @@ -587,7 +586,6 @@ class GraniteMoeHybridForCausalLM( SupportsPP, IsHybrid, SupportsQuant, - SupportsMambaPrefixCaching, ): packed_modules_mapping: dict[str, list[str]] = { "qkv_proj": [ diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 4dc67ff34d..a780cbe23c 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -1134,13 +1134,10 @@ def has_noops( @runtime_checkable class SupportsMambaPrefixCaching(Protocol): - """The interface for models whose mamba layers support prefix caching. - - This is currently experimental. + """The interface for attention-free Mamba models that support prefix + caching. Hybrid models get the same methods from `IsHybrid`. """ - supports_mamba_prefix_caching: ClassVar[Literal[True]] = True - @classmethod def get_mamba_state_copy_func(cls) -> tuple["MambaStateCopyFunc", ...]: """Return copy functions for the model's Mamba states.""" @@ -1156,24 +1153,6 @@ class SupportsMambaPrefixCaching(Protocol): return {mamba_type: copy_funcs for mamba_type in mamba_types} -@overload -def supports_mamba_prefix_caching( - model: type[object], -) -> TypeIs[type[SupportsMambaPrefixCaching]]: ... - - -@overload -def supports_mamba_prefix_caching( - model: object, -) -> TypeIs[SupportsMambaPrefixCaching]: ... - - -def supports_mamba_prefix_caching( - model: type[object] | object, -) -> TypeIs[type[SupportsMambaPrefixCaching]] | TypeIs[SupportsMambaPrefixCaching]: - return getattr(model, "supports_mamba_prefix_caching", False) - - @runtime_checkable class SupportsReplaySSM(Protocol): """The interface for models whose recurrent layers support ReplaySSM diff --git a/vllm/model_executor/models/jamba.py b/vllm/model_executor/models/jamba.py index fe9349c30b..315c821d6a 100644 --- a/vllm/model_executor/models/jamba.py +++ b/vllm/model_executor/models/jamba.py @@ -44,7 +44,6 @@ from .interfaces import ( HasInnerState, IsHybrid, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, ) from .utils import ( @@ -380,7 +379,6 @@ class JambaForCausalLM( SupportsLoRA, SupportsPP, IsHybrid, - SupportsMambaPrefixCaching, ): hf_to_vllm_mapper = WeightsMapper( orig_to_new_substr={".self_attn.": ".", ".A_log": ".A"}, diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index ae653950ac..050a13431d 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -74,7 +74,6 @@ from vllm.model_executor.models.interfaces import ( SupportsEagle, SupportsEagle3, SupportsLoRA, - SupportsMambaPrefixCaching, SupportsPP, SupportsQuant, SupportsReplaySSM, @@ -725,7 +724,6 @@ class NemotronHForCausalLM( IsHybrid, SupportsQuant, MixtureOfExperts, - SupportsMambaPrefixCaching, SupportsReplaySSM, ): # Relevant only if self.has_moe is True diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index cd0557ac45..09e5025c4e 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -50,7 +50,6 @@ from .interfaces import ( is_hybrid, requires_raw_input_tokens, supports_hc_sp, - supports_mamba_prefix_caching, supports_multimodal, supports_multimodal_encoder_tp_data, supports_multimodal_raw_input_only, @@ -881,7 +880,6 @@ class _ModelInfo: is_attention_free: bool is_hybrid: bool has_noops: bool - supports_mamba_prefix_caching: bool supports_replayssm: bool supports_transcription: bool supports_transcription_only: bool @@ -911,7 +909,6 @@ class _ModelInfo: has_inner_state=has_inner_state(model), is_attention_free=is_attention_free(model), is_hybrid=is_hybrid(model), - supports_mamba_prefix_caching=supports_mamba_prefix_caching(model), supports_replayssm=supports_replayssm(model), supports_transcription=supports_transcription(model), supports_transcription_only=( diff --git a/vllm/model_executor/models/zamba2.py b/vllm/model_executor/models/zamba2.py index 2e90da86ea..35e9728187 100644 --- a/vllm/model_executor/models/zamba2.py +++ b/vllm/model_executor/models/zamba2.py @@ -45,7 +45,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import ( ) from vllm.sequence import IntermediateTensors -from .interfaces import HasInnerState, IsHybrid, SupportsMambaPrefixCaching +from .interfaces import HasInnerState, IsHybrid from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix @@ -833,7 +833,7 @@ class Zamba2Model(nn.Module): return hidden_states -class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsMambaPrefixCaching): +class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid): """Zamba2 model with causal language modeling head. This class wraps the core Zamba2 model and adds: @@ -911,10 +911,6 @@ class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsMambaPrefixC LoRA and scheduler settings prefix: Optional prefix for parameter names - Raises: - AssertionError: If prefix caching is enabled - (not supported by Mamba) - """ config = vllm_config.model_config.hf_config diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index 3e7ae905a3..9820e8660e 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -245,13 +245,10 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): The code below carefully constructs the chunks such that: 1. Chunks contain tokens from a *single* sequence only. - 2. For every sequence, we are guaranteed that we can - retrieve the mamba state *every* chunk_size tokens. + 2. Chunk boundaries fall on multiples of chunk_size in each + sequence's absolute token positions, so the chunking does not + depend on how chunked prefill split the sequence. Constraint (1) dramatically simplifies the mamba kernels. - Constraint (2) dramatically simplifies the implementation - of prefix caching for mamba (wip). We need to take care - of the interaction with chunked prefill in order to - satisfy constraint (2). `checkpoint_offsets_p[i]` forces an extra chunk end that many tokens into row i's query (0 = none), because the SSD scan only materializes diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 2e4fe55b96..28a7a1106f 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1426,6 +1426,9 @@ class MambaManager(SingleTypeKVCacheManager): ) -> None: super().__init__(kv_cache_spec, block_pool, **kwargs) self.mamba_cache_mode = kv_cache_spec.mamba_cache_mode + assert self.mamba_cache_mode == "align" or not self.enable_caching, ( + "Mamba prefix caching requires mamba_cache_mode='align'" + ) self.num_speculative_blocks: int = kv_cache_spec.num_speculative_blocks self.has_prefill_checkpoint_blocks = ( self.mamba_cache_mode == "align"