biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

[V1] [Hybrid] Mamba1 Automatic Prefix Caching (#26377)

Browse Source 
Signed-off-by: asafg <39553475+Josephasafg@users.noreply.github.com>
This commit is contained in:
Asaf Joseph Gardin
2025-11-02 14:16:23 +02:00
committed by GitHub
parent 73444b7b56
commit 00b31a36a2
16 changed files with 442 additions and 153 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
vllm/model_executor/models/mamba.py
View File

@@ -29,6 +29,7 @@ from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.interfaces import (
HasInnerState,
IsAttentionFree,
SupportsMambaPrefixCaching,
SupportsPP,
)
from vllm.sequence import IntermediateTensors
@@ -193,15 +194,13 @@ class MambaModel(nn.Module):
return loaded_params
class MambaForCausalLM(nn.Module, HasInnerState, IsAttentionFree, SupportsPP):
class MambaForCausalLM(
nn.Module, HasInnerState, IsAttentionFree, SupportsPP, SupportsMambaPrefixCaching
):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
config = vllm_config.model_config.hf_config
cache_config = vllm_config.cache_config
lora_config = vllm_config.lora_config
self.scheduler_config = vllm_config.scheduler_config
assert not cache_config.enable_prefix_caching, (
"Mamba does not support prefix caching"
)
super().__init__()
self.config = config