light version of prefix caching for hybrid models gdn attention

vllm/engine/arg_utils.py

@@ -1356,6 +1356,15 @@ def create_engine_config(
             f"dcp_size={self.decode_context_parallel_size}."
         )
 
+        if (
+            envs.VLLM_USE_LIGHTER_MAMBA_CACHE
+            and self.enable_prefix_caching
+            and model_config.is_hybrid
+        ):
+            assert self.enable_chunked_prefill, (
+                "Prefix caching for hybrid models requires chunked prefill."
+            )
+
         cache_config = CacheConfig(
             block_size=self.block_size,
             gpu_memory_utilization=self.gpu_memory_utilization,

vllm/envs.py

@@ -244,6 +244,7 @@
     VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD: int = 256
     VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary"
     VLLM_USE_V2_MODEL_RUNNER: bool = False
+    VLLM_USE_LIGHTER_MAMBA_CACHE: bool = False
 
 
 def get_default_cache_root():
@@ -1565,6 +1566,9 @@ def get_vllm_port() -> int | None:
     "VLLM_USE_V2_MODEL_RUNNER": lambda: bool(
         int(os.getenv("VLLM_USE_V2_MODEL_RUNNER", "0"))
     ),
+    "VLLM_USE_LIGHTER_MAMBA_CACHE": lambda: bool(
+        os.getenv("VLLM_USE_LIGHTER_MAMBA_CACHE", False)
+    ),
 }
 
 # --8<-- [end:env-vars-definition]
@@ -1692,6 +1696,7 @@ def compile_factors() -> dict[str, object]:
         "VLLM_CPU_MOE_PREPACK",
         "VLLM_CPU_SGL_KERNEL",
         "VLLM_TEST_FORCE_LOAD_FORMAT",
+        "VLLM_USE_LIGHTER_MAMBA_CACHE",
         "LOCAL_RANK",
         "CUDA_VISIBLE_DEVICES",
         "NO_COLOR",

vllm/model_executor/models/config.py

@@ -4,6 +4,7 @@
 from math import lcm
 from typing import TYPE_CHECKING
 
+from vllm import envs
 from vllm.attention.backends.registry import AttentionBackendEnum
 from vllm.logger import init_logger
 from vllm.model_executor.models import ModelRegistry
@@ -389,7 +390,10 @@ def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
         mamba_page_size = MambaSpec(
             shapes=model_cls.get_mamba_state_shape_from_config(vllm_config),
             dtypes=model_cls.get_mamba_state_dtype_from_config(vllm_config),
-            block_size=model_config.max_model_len,
+            block_size=model_config.max_model_len
+            if not envs.VLLM_USE_LIGHTER_MAMBA_CACHE
+            else cache_config.block_size,
+            enable_caching=cache_config.enable_prefix_caching,
         ).page_size_bytes
 
         # Model may be marked as is_hybrid

vllm/model_executor/models/qwen3_next.py

@@ -10,6 +10,7 @@
 from torch import nn
 from transformers.activations import ACT2FN
 
+from vllm import envs
 from vllm.attention.backends.abstract import AttentionMetadata
 from vllm.attention.layer import Attention
 from vllm.compilation.decorators import support_torch_compile
@@ -1192,9 +1193,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         cache_config = vllm_config.cache_config
 
         scheduler_config = vllm_config.scheduler_config
-        assert not cache_config.enable_prefix_caching, (
-            "Qwen3Next currently does not support prefix caching"
-        )
+        if not envs.VLLM_USE_LIGHTER_MAMBA_CACHE:
+            assert not cache_config.enable_prefix_caching, (
+                "Qwen3Next currently does not support prefix caching"
+            )
         self.quant_config = vllm_config.quant_config
 
         super().__init__()

vllm/model_executor/models/qwen3_next_mtp.py

@@ -7,6 +7,7 @@
 import torch
 from torch import nn
 
+from vllm import envs
 from vllm.compilation.decorators import support_torch_compile
 from vllm.config import VllmConfig
 from vllm.distributed.parallel_state import get_pp_group
@@ -234,9 +235,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         config = vllm_config.model_config.hf_config
         self.vllm_config = vllm_config
         cache_config = vllm_config.cache_config
-        assert not cache_config.enable_prefix_caching, (
-            "Qwen3NextMTP currently does not support prefix caching"
-        )
+        if not envs.VLLM_USE_LIGHTER_MAMBA_CACHE:
+            assert not cache_config.enable_prefix_caching, (
+                "Qwen3NextMTP currently does not support prefix caching"
+            )
 
         self.quant_config = vllm_config.quant_config
 

vllm/v1/attention/backends/gdn_attn.py

@@ -6,6 +6,7 @@
 
 import torch
 
+from vllm import envs
 from vllm.attention.backends.abstract import AttentionBackend
 from vllm.attention.backends.utils import PAD_SLOT_ID
 from vllm.config import VllmConfig
@@ -325,6 +326,17 @@ def build(  # type: ignore[override]
             non_spec_query_start_loc = self.non_spec_query_start_loc[: batch_size + 1]
             non_spec_query_start_loc[num_decodes + 1 :].fill_(non_spec_num_query_tokens)
 
+            if envs.VLLM_USE_LIGHTER_MAMBA_CACHE:
+                # NOTE: With Mamba prefix-caching support, a request can consist of
+                # multiple blocks. This makes the state_indices non-contiguous, so
+                # we must explicitly make them contiguous here.
+                if spec_state_indices_tensor is not None:
+                    spec_state_indices_tensor = spec_state_indices_tensor.contiguous()
+                if non_spec_state_indices_tensor is not None:
+                    non_spec_state_indices_tensor = (
+                        non_spec_state_indices_tensor.contiguous()
+                    )
+
         attn_metadata = GDNAttentionMetadata(
             num_prefills=num_prefills,
             num_prefill_tokens=num_prefill_tokens,

vllm/v1/attention/backends/linear_attn.py

@@ -4,6 +4,7 @@
 
 import torch
 
+from vllm import envs
 from vllm.attention.backends.abstract import AttentionBackend
 from vllm.config import VllmConfig
 from vllm.v1.attention.backends.utils import (
@@ -59,6 +60,9 @@ def build(
 
         state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
 
+        if envs.VLLM_USE_LIGHTER_MAMBA_CACHE:
+            state_indices_tensor = state_indices_tensor.contiguous()
+
         num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
             split_decodes_and_prefills(
                 common_attn_metadata, decode_threshold=self.reorder_batch_threshold

vllm/v1/attention/backends/mamba1_attn.py

@@ -5,6 +5,7 @@
 
 import torch
 
+from vllm import envs
 from vllm.attention.backends.abstract import AttentionBackend
 from vllm.attention.backends.utils import PAD_SLOT_ID
 from vllm.config import VllmConfig
@@ -73,7 +74,10 @@ def build(
 
         # TODO(@Josephasafg) Mamba1 and Mamba2 have a lot of code in common here.
         # We should consolidate this code
-        if self.vllm_config.cache_config.enable_prefix_caching:
+        if (
+            not envs.VLLM_USE_LIGHTER_MAMBA_CACHE
+            and self.vllm_config.cache_config.enable_prefix_caching
+        ):
             # Return a tensor of shape (#requests, #max blocks)
             state_indices_tensor = common_attn_metadata.block_table_tensor
             mamba_block_size = self.kv_cache_spec.block_size
@@ -108,7 +112,10 @@ def build(
                 common_attn_metadata.query_start_loc.device
             )
 
-            if self.vllm_config.cache_config.enable_prefix_caching:
+            if (
+                not envs.VLLM_USE_LIGHTER_MAMBA_CACHE
+                and self.vllm_config.cache_config.enable_prefix_caching
+            ):
                 assert num_computed_tokens is not None
                 num_computed_tokens_p = num_computed_tokens[
                     num_reqs - num_prefills : num_reqs
@@ -129,7 +136,10 @@ def build(
             state_indices_tensor = self.state_indices_tensor[:num_decode_tokens]
             state_indices_tensor[num_decodes:] = PAD_SLOT_ID
 
-            if self.vllm_config.cache_config.enable_prefix_caching:
+            if (
+                not envs.VLLM_USE_LIGHTER_MAMBA_CACHE
+                and self.vllm_config.cache_config.enable_prefix_caching
+            ):
                 self.block_idx_last_scheduled_token[:num_decodes].copy_(
                     block_idx_last_scheduled_token, non_blocking=True
                 )

vllm/v1/attention/backends/mamba2_attn.py

@@ -5,6 +5,7 @@
 
 import torch
 
+from vllm import envs
 from vllm.attention.backends.abstract import AttentionBackend
 from vllm.config import VllmConfig
 from vllm.utils.math_utils import cdiv
@@ -172,7 +173,10 @@ def build(
         block_idx_first_scheduled_token = None
         block_idx_first_scheduled_token_p = None
 
-        if self.vllm_config.cache_config.enable_prefix_caching:
+        if (
+            not envs.VLLM_USE_LIGHTER_MAMBA_CACHE
+            and self.vllm_config.cache_config.enable_prefix_caching
+        ):
             # Return a tensor of shape (#requests, #max blocks)
             state_indices_tensor = common_attn_metadata.block_table_tensor
             # Additional cache-related varaiables:
@@ -190,6 +194,11 @@ def build(
         else:
             # Always return just a single block per each request:
             state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
+            if envs.VLLM_USE_LIGHTER_MAMBA_CACHE:
+                # NOTE: With Mamba prefix-caching support, a request can consist of
+                # multiple blocks. This makes the state_indices non-contiguous, so
+                # we must explicitly make them contiguous here.
+                state_indices_tensor = state_indices_tensor.contiguous()
             # Additional cache-related varaiables:
             block_idx_last_scheduled_token = None
             block_idx_last_computed_token = None
@@ -219,7 +228,10 @@ def build(
                 - num_decode_tokens
             )
 
-            if self.vllm_config.cache_config.enable_prefix_caching:
+            if (
+                not envs.VLLM_USE_LIGHTER_MAMBA_CACHE
+                and self.vllm_config.cache_config.enable_prefix_caching
+            ):
                 assert num_computed_tokens is not None
                 num_computed_tokens_p = num_computed_tokens[
                     num_reqs - num_prefills : num_reqs
@@ -308,7 +320,10 @@ def build(
             )
             state_indices_tensor = self.state_indices_tensor[:num_decode_tokens]
 
-            if self.vllm_config.cache_config.enable_prefix_caching:
+            if (
+                not envs.VLLM_USE_LIGHTER_MAMBA_CACHE
+                and self.vllm_config.cache_config.enable_prefix_caching
+            ):
                 self.block_idx_last_scheduled_token[:num_decodes].copy_(
                     block_idx_last_scheduled_token, non_blocking=True
                 )

vllm/v1/core/block_pool.py

@@ -206,6 +206,50 @@ def get_cached_block(
             cached_blocks.append(block)
         return cached_blocks
 
+    def cache_full_block(
+        self,
+        request: Request,
+        block: KVCacheBlock,
+        cached_block_index: int,
+        block_size: int,
+        kv_cache_group_id: int,
+    ) -> None:
+        """Cache a full block for prefix caching."""
+
+        assert cached_block_index >= 0
+        assert len(request.block_hashes) > cached_block_index
+        new_block_hash: BlockHash = request.block_hashes[cached_block_index]
+        new_hashes: list[ExternalBlockHash] | None = (
+            [] if self.enable_kv_cache_events else None
+        )
+        assert block.block_hash is None
+
+        # Update and added the full block to the cache.
+        block_hash_with_group_id: BlockHashWithGroupId = make_block_hash_with_group_id(
+            new_block_hash, kv_cache_group_id
+        )
+        block.block_hash = block_hash_with_group_id
+        self.cached_block_hash_to_block.insert(block_hash_with_group_id, block)
+        if new_hashes is not None:
+            new_hashes.append(maybe_convert_block_hash(new_block_hash))
+
+        if self.enable_kv_cache_events:
+            parent_block_hash: ExternalBlockHash | None = None
+
+            self.kv_event_queue.append(
+                BlockStored(
+                    block_hashes=new_hashes,
+                    parent_block_hash=parent_block_hash,
+                    token_ids=request.all_token_ids[
+                        cached_block_index * block_size : (cached_block_index + 1)
+                        * block_size
+                    ],
+                    block_size=block_size,
+                    lora_id=request.lora_request.id if request.lora_request else None,
+                    medium=MEDIUM_GPU,
+                )
+            )
+
     def cache_full_blocks(
         self,
         request: Request,