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

[Hybrid] Added supports_mamba_prefix_caching Protocol (#27339)

Browse Source 
Signed-off-by: asafg <39553475+Josephasafg@users.noreply.github.com>
This commit is contained in:
Asaf Joseph Gardin
2025-10-27 15:05:20 +02:00
committed by GitHub
parent f4e8154076
commit 9273754222
10 changed files with 93 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

10
vllm/model_executor/models/mamba2.py
View File

@@ -25,7 +25,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
VocabParallelEmbedding,
)
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
from vllm.model_executor.models.interfaces import HasInnerState, IsAttentionFree
from vllm.model_executor.models.interfaces import (
HasInnerState,
IsAttentionFree,
SupportsMambaPrefixCaching,
)
from vllm.sequence import IntermediateTensors
from .utils import (
@@ -189,7 +193,9 @@ class Mamba2Model(nn.Module):
return loaded_params
class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):
class Mamba2ForCausalLM(
nn.Module, HasInnerState, IsAttentionFree, SupportsMambaPrefixCaching
):
@classmethod
def get_mamba_state_dtype_from_config(
cls,