diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 3862aa922244..6b6498450670 100644 --- a/vllm/engine/arg_utils.py +++ b/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, diff --git a/vllm/envs.py b/vllm/envs.py index d0f279809626..5d226baf2d4a 100755 --- a/vllm/envs.py +++ b/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", diff --git a/vllm/model_executor/models/config.py b/vllm/model_executor/models/config.py index 4b08472538db..9011b5fcadbc 100644 --- a/vllm/model_executor/models/config.py +++ b/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 diff --git a/vllm/model_executor/models/qwen3_next.py b/vllm/model_executor/models/qwen3_next.py index ccf6cc6e5894..7f5cd2b128e8 100644 --- a/vllm/model_executor/models/qwen3_next.py +++ b/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__() diff --git a/vllm/model_executor/models/qwen3_next_mtp.py b/vllm/model_executor/models/qwen3_next_mtp.py index 83694caa5248..17ee3a9792d8 100644 --- a/vllm/model_executor/models/qwen3_next_mtp.py +++ b/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 diff --git a/vllm/v1/attention/backends/gdn_attn.py b/vllm/v1/attention/backends/gdn_attn.py index ace2cbb0564c..4aec54c3a296 100644 --- a/vllm/v1/attention/backends/gdn_attn.py +++ b/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, diff --git a/vllm/v1/attention/backends/linear_attn.py b/vllm/v1/attention/backends/linear_attn.py index 004baa2d09cd..0c99f728a0fb 100644 --- a/vllm/v1/attention/backends/linear_attn.py +++ b/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 diff --git a/vllm/v1/attention/backends/mamba1_attn.py b/vllm/v1/attention/backends/mamba1_attn.py index fcda6134016b..7c184f311d42 100644 --- a/vllm/v1/attention/backends/mamba1_attn.py +++ b/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 ) diff --git a/vllm/v1/attention/backends/mamba2_attn.py b/vllm/v1/attention/backends/mamba2_attn.py index bf1d8f09ab0a..efe4d16004a3 100644 --- a/vllm/v1/attention/backends/mamba2_attn.py +++ b/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 ) diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index c779e3d34b3e..0c6e801796d0 100644 --- a/vllm/v1/core/block_pool.py +++ b/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, diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 278970ae7ee8..fa799a91b6b4 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -42,7 +42,7 @@ from vllm.v1.core.sched.request_queue import SchedulingPolicy, create_request_queue from vllm.v1.core.sched.utils import check_stop, remove_all from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs -from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.metrics.stats import ( PrefixCacheStats, SchedulerStats, @@ -213,6 +213,14 @@ def __init__( self.use_pp = self.parallel_config.pipeline_parallel_size > 1 self.use_v2_model_runner = envs.VLLM_USE_V2_MODEL_RUNNER + def _has_mamba_spec(self) -> bool: + has_mamba: bool = any( + isinstance(spec.kv_cache_spec, MambaSpec) + for spec in self.kv_cache_config.kv_cache_groups + ) + assert not has_mamba or self.vllm_config.model_config.is_hybrid + return has_mamba + def schedule(self) -> SchedulerOutput: # NOTE(woosuk) on the scheduling algorithm: # There's no "decoding phase" nor "prefill phase" in the scheduler. @@ -263,14 +271,27 @@ def schedule(self) -> SchedulerOutput: req_index += 1 continue - num_new_tokens = ( - request.num_tokens_with_spec - + request.num_output_placeholders - - request.num_computed_tokens - ) + # Ensure new tokens for a request in the prefill phase do not contain + # sps tokens, especially in the last prefill chunk. For a hybrid-model, + # extra sps tokens would corrupt the generated Mamba state. + # TODO: This logic does not yet handle resumed requests. + if request.num_computed_tokens < request.num_prompt_tokens: + num_new_tokens = ( + min( + request.num_tokens_with_spec + request.num_output_placeholders, + request.num_prompt_tokens, + ) + - request.num_computed_tokens + ) + else: + num_new_tokens = ( + request.num_tokens_with_spec + + request.num_output_placeholders + - request.num_computed_tokens + ) + if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min(num_new_tokens, token_budget) # Make sure the input position does not exceed the max model len. # This is necessary when using spec decoding. @@ -278,6 +299,35 @@ def schedule(self) -> SchedulerOutput: num_new_tokens, self.max_model_len - 1 - request.num_computed_tokens ) + if ( + envs.VLLM_USE_LIGHTER_MAMBA_CACHE + and self.cache_config.enable_prefix_caching + and self._has_mamba_spec() + ): + # To enable block-aligned caching of the Mamba state, `num_new_tokens` + # must be a multiple of `block_size`. + # As an exception, if `num_new_tokens` is less than `block_size`, the + # state is simply not cached, requiring no special handling. + # Additionally, when Eagle mode is enabled, FullAttn prunes the last + # matching block. To prevent this from causing a Mamba cache miss, the + # last chunk must be larger than `block_size`. + block_size = self.block_size + max_last_chunk = block_size * (2 if self.use_eagle else 1) + if num_new_tokens < max_last_chunk: + num_new_tokens = min(num_new_tokens, token_budget) + else: + ori_num_new_tokens = num_new_tokens + num_new_tokens = min(num_new_tokens, token_budget) + num_new_tokens = num_new_tokens // block_size * block_size + if ( + self.use_eagle + and ori_num_new_tokens - num_new_tokens < block_size + ): + assert num_new_tokens >= block_size + num_new_tokens -= block_size + else: + num_new_tokens = min(num_new_tokens, token_budget) + # Schedule encoder inputs. encoder_inputs_to_schedule = None external_load_encoder_input: list[int] = [] @@ -306,6 +356,8 @@ def schedule(self) -> SchedulerOutput: # its max_total_tokens or max_model_len. # 2. The encoder budget is exhausted. # 3. The encoder cache is exhausted. + # 4. Insufficient budget for a block-aligned chunk in hybrid + # models with lighter mamba prefix caching. # NOTE(woosuk): Here, by doing `continue` instead of `break`, # we do not strictly follow the FCFS scheduling policy and # allow the lower-priority requests to be scheduled. @@ -538,7 +590,30 @@ def schedule(self) -> SchedulerOutput: # we can stop the scheduling here. break - num_new_tokens = min(num_new_tokens, token_budget) + if ( + envs.VLLM_USE_LIGHTER_MAMBA_CACHE + and self.cache_config.enable_prefix_caching + and self._has_mamba_spec() + ): + block_size = self.block_size + max_last_chunk = block_size * (2 if self.use_eagle else 1) + if num_new_tokens < max_last_chunk: + num_new_tokens = min(num_new_tokens, token_budget) + else: + ori_num_new_tokens = num_new_tokens + num_new_tokens = min(num_new_tokens, token_budget) + num_new_tokens = num_new_tokens // block_size * block_size + if ( + self.use_eagle + and ori_num_new_tokens - num_new_tokens < block_size + ): + assert num_new_tokens >= block_size + num_new_tokens -= block_size + if num_new_tokens == 0: + token_budget = 0 + break + else: + num_new_tokens = min(num_new_tokens, token_budget) assert num_new_tokens > 0 # Schedule encoder inputs. diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 4aeb17a156bb..ad5dbd9b4b53 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -5,6 +5,7 @@ from collections import defaultdict from collections.abc import Sequence +from vllm import envs from vllm.utils.math_utils import cdiv from vllm.v1.core.block_pool import BlockPool from vllm.v1.core.kv_cache_utils import BlockHashList, KVCacheBlock @@ -652,6 +653,12 @@ def get_num_common_prefix_blocks(self, running_request_id: str) -> int: class MambaManager(SingleTypeKVCacheManager): + def __init__(self, kv_cache_spec: MambaSpec, **kwargs) -> None: + super().__init__(kv_cache_spec, **kwargs) + if envs.VLLM_USE_LIGHTER_MAMBA_CACHE: + self._req_to_computed_tokens: dict[str, int] = {} + self._req_to_new_tokens: dict[str, int] = {} + @classmethod def find_longest_cache_hit( cls, @@ -700,6 +707,38 @@ def find_longest_cache_hit( return computed_blocks + def remove_skipped_blocks(self, request_id, num_computed_tokens): + if not envs.VLLM_USE_LIGHTER_MAMBA_CACHE: + return super().remove_skipped_blocks(request_id, num_computed_tokens) + else: + assert isinstance(self.kv_cache_spec, MambaSpec) + self._req_to_computed_tokens[request_id] = num_computed_tokens + # Each request will always have 1 block at this moment, so no need to + # remove blocks. + if not self.kv_cache_spec.enable_caching: + return + blocks: list[KVCacheBlock] = self.req_to_blocks[request_id] + num_blocks = len(blocks) + if num_blocks > 0: + prefix_block: KVCacheBlock = blocks[0] + assert self._req_to_computed_tokens[request_id] != 0 + if not prefix_block.is_null: + self.block_pool.free_blocks([prefix_block]) + blocks[0] = self.block_pool.null_block + + if num_blocks > 2 + self.kv_cache_spec.num_speculative_blocks: + assert num_blocks == 3 + self.kv_cache_spec.num_speculative_blocks + last_new_tokens = self._req_to_new_tokens[request_id] + assert last_new_tokens % self.block_size == 0 + assert last_new_tokens >= self.block_size + self.block_pool.free_blocks(blocks[-1:]) + blocks.pop() + else: + assert ( + num_blocks == 0 + or num_blocks == 2 + self.kv_cache_spec.num_speculative_blocks + ) + def get_num_common_prefix_blocks(self, running_request_id: str) -> int: """ cascade attention is not supported by mamba @@ -715,14 +754,67 @@ def get_num_blocks_to_allocate( # Allocate extra `num_speculative_blocks` blocks for # speculative decoding (MTP/EAGLE) with linear attention. assert isinstance(self.kv_cache_spec, MambaSpec) - if self.kv_cache_spec.num_speculative_blocks > 0: - num_tokens += ( - self.kv_cache_spec.block_size - * self.kv_cache_spec.num_speculative_blocks + if not envs.VLLM_USE_LIGHTER_MAMBA_CACHE: + if self.kv_cache_spec.num_speculative_blocks > 0: + num_tokens += ( + self.kv_cache_spec.block_size + * self.kv_cache_spec.num_speculative_blocks + ) + return super().get_num_blocks_to_allocate( + request_id, num_tokens, new_computed_blocks ) - return super().get_num_blocks_to_allocate( - request_id, num_tokens, new_computed_blocks - ) + else: + num_computed_tokens = self._req_to_computed_tokens.get(request_id, 0) + if num_computed_tokens == 0: + num_new_blocks = 1 + ( + self.kv_cache_spec.num_speculative_blocks + if self.kv_cache_spec.num_speculative_blocks > 0 + else 0 + ) + else: + num_new_blocks = 0 + + if self.kv_cache_spec.enable_caching: + num_new_tokens = ( + num_tokens + - num_computed_tokens + - len(new_computed_blocks) * self.block_size + ) + if num_computed_tokens != 0: + num_new_tokens -= self.kv_cache_spec.num_speculative_blocks + self._req_to_new_tokens[request_id] = num_new_tokens + # NOTE: last chunk may larger than block_size when using eagle. + if ( + num_new_tokens >= self.block_size + and num_new_tokens % self.block_size == 0 + ): + num_new_blocks += 1 + + # If a computed block of a request is an eviction candidate (in the + # free queue and ref_cnt == 0), it will be changed from a free block + # to a computed block when the request is allocated, so we also count + # it as needed to be allocated. + num_evictable_computed_blocks = sum( + blk.ref_cnt == 0 and not blk.is_null for blk in new_computed_blocks + ) + return num_new_blocks + num_evictable_computed_blocks + + def save_new_computed_blocks( + self, request_id: str, new_computed_blocks: Sequence[KVCacheBlock] + ) -> None: + assert isinstance(self.kv_cache_spec, MambaSpec) + if envs.VLLM_USE_LIGHTER_MAMBA_CACHE: + if not self.kv_cache_spec.enable_caching: + return + if request_id not in self.num_cached_block: + if new_computed_blocks: + assert ( + len(new_computed_blocks) == 1 or new_computed_blocks[-2].is_null + ) + new_computed_blocks = list(new_computed_blocks[-1:]) + else: + new_computed_blocks = [self.block_pool.null_block] + super().save_new_computed_blocks(request_id, new_computed_blocks) def allocate_new_blocks( self, request_id: str, num_tokens: int @@ -730,12 +822,76 @@ def allocate_new_blocks( # Allocate extra `num_speculative_blocks` blocks for # speculative decoding (MTP/EAGLE) with linear attention. assert isinstance(self.kv_cache_spec, MambaSpec) - if self.kv_cache_spec.num_speculative_blocks > 0: - num_tokens += ( - self.kv_cache_spec.block_size - * self.kv_cache_spec.num_speculative_blocks - ) - return super().allocate_new_blocks(request_id, num_tokens) + if not envs.VLLM_USE_LIGHTER_MAMBA_CACHE: + if self.kv_cache_spec.num_speculative_blocks > 0: + num_tokens += ( + self.kv_cache_spec.block_size + * self.kv_cache_spec.num_speculative_blocks + ) + return super().allocate_new_blocks(request_id, num_tokens) + else: + num_computed_tokens = self._req_to_computed_tokens.get(request_id, 0) + if num_computed_tokens == 0: + num_new_blocks = 1 + ( + self.kv_cache_spec.num_speculative_blocks + if self.kv_cache_spec.num_speculative_blocks > 0 + else 0 + ) + else: + assert num_tokens >= num_computed_tokens + num_new_blocks = 0 + + if self.kv_cache_spec.enable_caching: + num_new_tokens = self._req_to_new_tokens[request_id] + if ( + num_new_tokens >= self.block_size + and num_new_tokens % self.block_size == 0 + ): + num_new_blocks += 1 + + req_blocks: list[KVCacheBlock] = self.req_to_blocks[request_id] + if num_new_blocks <= 0: + return [] + else: + new_blocks: list[KVCacheBlock] = self.block_pool.get_new_blocks( + num_new_blocks + ) + req_blocks.extend(new_blocks) + return new_blocks + + def cache_blocks(self, request: Request, num_tokens: int) -> None: + assert isinstance(self.kv_cache_spec, MambaSpec) + if envs.VLLM_USE_LIGHTER_MAMBA_CACHE: + num_computed_tokens = request.num_computed_tokens + num_new_tokens = self._req_to_new_tokens[request.request_id] + # NOTE:For sps, an extra block may be allocated but not cached + if ( + num_new_tokens >= self.block_size + and num_new_tokens % self.block_size == 0 + and num_tokens % self.block_size == 0 + ): + assert num_new_tokens % self.block_size == 0 + assert num_computed_tokens % self.block_size == 0 + assert ( + len(self.req_to_blocks[request.request_id]) + == 3 + self.kv_cache_spec.num_speculative_blocks + ) + self.block_pool.cache_full_block( + request=request, + block=self.req_to_blocks[request.request_id][-1], + cached_block_index=(num_tokens // self.block_size - 1), + block_size=self.block_size, + kv_cache_group_id=self.kv_cache_group_id, + ) + self.num_cached_block[request.request_id] += 1 + else: + super().cache_blocks(request, num_tokens) + + def free(self, request_id: str) -> None: + if envs.VLLM_USE_LIGHTER_MAMBA_CACHE: + self._req_to_computed_tokens.pop(request_id, None) + self._req_to_new_tokens.pop(request_id, None) + super().free(request_id) class CrossAttentionManager(SingleTypeKVCacheManager): diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 751862aa9c76..8feaaab7f0eb 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -8,6 +8,7 @@ import torch from typing_extensions import Self +from vllm import envs from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.utils.math_utils import cdiv @@ -247,6 +248,7 @@ class MambaSpec(KVCacheSpec): page_size_padded: int | None = None mamba_type: str = "mamba2" num_speculative_blocks: int = 0 + enable_caching: bool = False @property def page_size_bytes(self) -> int: @@ -260,8 +262,18 @@ def page_size_bytes(self) -> int: return page_size def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int: - max_model_len = vllm_config.model_config.max_model_len - return cdiv(max_model_len, self.block_size) * self.page_size_bytes + # We allocate 1 block for each request now, so max_memory_usage_bytes is + # the same as page_size_bytes. + # Need to update this when supporting prefix caching. + if not envs.VLLM_USE_LIGHTER_MAMBA_CACHE: + max_model_len = vllm_config.model_config.max_model_len + return cdiv(max_model_len, self.block_size) * self.page_size_bytes + else: + # NOTE: We allocate 1 block per request by default. With prefix + # caching enabled, up to 2 additional blocks are required: one + # for reading the matched prefix and one for caching the current + # state. + return self.page_size_bytes * (3 if self.enable_caching else 1) @dataclass(frozen=True) diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 978224faae65..099f64f62ff4 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -109,6 +109,7 @@ reorder_batch_to_split_decodes_and_prefills, split_attn_metadata, ) +from vllm.v1.core.sched.output import CachedRequestData, NewRequestData from vllm.v1.cudagraph_dispatcher import CudagraphDispatcher from vllm.v1.kv_cache_interface import ( AttentionSpec, @@ -1524,6 +1525,7 @@ def _build_attention_metadata( num_tokens: int, num_reqs: int, max_query_len: int, + scheduler_output: SchedulerOutput | None = None, num_tokens_padded: int | None = None, num_reqs_padded: int | None = None, ubatch_slices: UBatchSlices | None = None, @@ -1692,6 +1694,16 @@ def _build_attn_group_metadata( cm.block_table_tensor, cm.slot_mapping = ( _get_block_table_and_slot_mapping(kv_cache_gid) ) + if ( + envs.VLLM_USE_LIGHTER_MAMBA_CACHE + and self.cache_config.enable_prefix_caching + and isinstance(kv_cache_group.kv_cache_spec, MambaSpec) + ): + cm.block_table_tensor = cm.block_table_tensor[:, 1:] + assert scheduler_output is not None + self._preprocess_mamba_prefix( + scheduler_output, kv_cache_gid, kv_cache_group + ) if self.speculative_config and spec_decode_common_attn_metadata is None: if isinstance(self.drafter, EagleProposer): @@ -2904,6 +2916,98 @@ def _register_layerwise_nvtx_hooks(self) -> None: pyt_hooks.register_hooks(self.model, self.model.__class__.__name__) self.layerwise_nvtx_hooks_registered = True + def _mamba_copy_block( + self, + kv_cache_group_spec: KVCacheGroupSpec, + src_block_id: int, + dest_block_id: int, + ): + forward_context = self.compilation_config.static_forward_context + for layer_name in kv_cache_group_spec.layer_names: + kv_caches: list[list[torch.Tensor]] = forward_context[layer_name].kv_cache + for kv_cache in kv_caches: + if isinstance(kv_cache, torch.Tensor): + kv_cache[dest_block_id].copy_(kv_cache[src_block_id]) + elif isinstance(kv_cache, list): + for kv_cache_part in kv_cache: + kv_cache_part[dest_block_id].copy_(kv_cache_part[src_block_id]) + + def _preprocess_mamba_prefix( + self, + scheduler_output: "SchedulerOutput", + kv_cache_group_id: int, + kv_cache_group_spec: KVCacheGroupSpec, + ): + assert isinstance(kv_cache_group_spec.kv_cache_spec, MambaSpec) + assert self.cache_config.enable_prefix_caching + new_reqs: list[NewRequestData] = scheduler_output.scheduled_new_reqs + for new_req in new_reqs: + if new_req.num_computed_tokens == 0: + continue + block_ids: list[int] = new_req.block_ids[kv_cache_group_id] + assert block_ids[0] != 0, f"{block_ids=}" + prefix_block_id, dest_block_id = block_ids[0], block_ids[1] + self._mamba_copy_block(kv_cache_group_spec, prefix_block_id, dest_block_id) + cached_reqs: CachedRequestData = scheduler_output.scheduled_cached_reqs + for i, resumed in enumerate(cached_reqs.resumed_from_preemption): + if not resumed: + continue + group_block_ids: tuple[list[int], ...] | None = cached_reqs.new_block_ids[i] + assert group_block_ids is not None + new_block_ids: list[int] = group_block_ids[kv_cache_group_id] + assert ( + len(new_block_ids) + >= 2 + kv_cache_group_spec.kv_cache_spec.num_speculative_blocks + ) + if cached_reqs.num_computed_tokens[i] == 0: + continue + assert new_block_ids[0] != 0, f"{new_block_ids=}" + prefix_block_id, dest_block_id = new_block_ids[0], new_block_ids[1] + self._mamba_copy_block(kv_cache_group_spec, prefix_block_id, dest_block_id) + + def _postprocess_mamba_cache(self, scheduler_output: "SchedulerOutput"): + assert self.cache_config.enable_prefix_caching + for kv_cache_group_id, kv_cache_group_spec in enumerate( + self.kv_cache_config.kv_cache_groups + ): + if not isinstance(kv_cache_group_spec.kv_cache_spec, MambaSpec): + continue + new_reqs: list[NewRequestData] = scheduler_output.scheduled_new_reqs + num_speculative_blocks = ( + kv_cache_group_spec.kv_cache_spec.num_speculative_blocks + ) + for new_req in new_reqs: + block_ids: list[int] = new_req.block_ids[kv_cache_group_id] + if len(block_ids) <= 2 + num_speculative_blocks: + continue + assert len(block_ids) == 3 + num_speculative_blocks + src_block_id, dest_block_id = block_ids[1], block_ids[-1] + self._mamba_copy_block(kv_cache_group_spec, src_block_id, dest_block_id) + cached_reqs: CachedRequestData = scheduler_output.scheduled_cached_reqs + for i, req_id in enumerate(cached_reqs.req_ids): + group_block_ids: tuple[list[int], ...] | None = ( + cached_reqs.new_block_ids[i] + ) + if group_block_ids is None: + assert not cached_reqs.resumed_from_preemption[i] + continue + new_block_ids: list[int] = group_block_ids[kv_cache_group_id] + if not new_block_ids: + assert not cached_reqs.resumed_from_preemption[i] + continue + if not cached_reqs.resumed_from_preemption[i]: + assert len(new_block_ids) == 1 + block_ids_new: list[int] = self.requests[req_id].block_ids[ + kv_cache_group_id + ] + src_block_id, dest_block_id = block_ids_new[1], new_block_ids[0] + else: + if len(new_block_ids) == 2 + num_speculative_blocks: + continue + assert len(new_block_ids) == 3 + num_speculative_blocks + src_block_id, dest_block_id = new_block_ids[1], new_block_ids[-1] + self._mamba_copy_block(kv_cache_group_spec, src_block_id, dest_block_id) + @torch.inference_mode() def execute_model( self, @@ -3040,6 +3144,7 @@ def execute_model( num_tokens=num_tokens_unpadded, num_tokens_padded=num_tokens_padded if pad_attn else None, num_reqs=num_reqs, + scheduler_output=scheduler_output, num_reqs_padded=num_reqs_padded if pad_attn else None, max_query_len=max_num_scheduled_tokens, ubatch_slices=ubatch_slices_attn, @@ -3205,6 +3310,11 @@ def sample_tokens( apply_grammar_bitmask( scheduler_output, grammar_output, self.input_batch, logits ) + if ( + envs.VLLM_USE_LIGHTER_MAMBA_CACHE + and self.cache_config.enable_prefix_caching + ): + self._postprocess_mamba_cache(scheduler_output) with record_function_or_nullcontext("gpu_model_runner: sample"): sampler_output = self._sample(logits, spec_decode_metadata)