A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[Hybrid] Added supports_mamba_prefix_caching Protocol (#27339)
Signed-off-by: asafg <39553475+Josephasafg@users.noreply.github.com>
Asaf Joseph Gardin committed
Oct 27, 2025 at 15:05 UTC
9273754222d4e4ada346cdff4f0f4cd967155e28
10 files changed
+93
-20
vllm/config/model.py
+4
@@ -1656,6 +1656,10 @@ class ModelConfig:
1656
def has_inner_state(self):
1657
return self._model_info.has_inner_state
1658
1659
+ @property
1660
+ def supports_mamba_prefix_caching(self) -> bool:
1661
+ return self._model_info.supports_mamba_prefix_caching
1662
+
1663
@property
1664
def use_mla(self) -> bool:
1665
return self.is_deepseek_mla and not envs.VLLM_MLA_DISABLE
vllm/model_executor/models/bamba.py
+15
-2
@@ -37,7 +37,14 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
37
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
38
from vllm.sequence import IntermediateTensors
39
40
-from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP, SupportsQuant
40
+from .interfaces import (
41
+ HasInnerState,
42
+ IsHybrid,
43
+ SupportsLoRA,
44
+ SupportsMambaPrefixCaching,
45
+ SupportsPP,
46
+ SupportsQuant,
47
+)
48
from .utils import (
49
AutoWeightsLoader,
50
is_pp_missing_parameter,
@@ -394,7 +401,13 @@ class BambaModel(nn.Module):
401
402
403
class BambaForCausalLM(
397
- nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHybrid, SupportsQuant
404
+ nn.Module,
405
+ HasInnerState,
406
+ SupportsLoRA,
407
+ SupportsPP,
408
+ IsHybrid,
409
+ SupportsQuant,
410
+ SupportsMambaPrefixCaching,
411
):
412
packed_modules_mapping = {
413
"qkv_proj": [
vllm/model_executor/models/config.py
+1
-10
@@ -295,17 +295,8 @@ class MambaModelConfig(VerifyAndUpdateConfig):
295
# override by prefix caching logic later)
296
cache_config.mamba_block_size = model_config.max_model_len
297
298
- # TODO(@tdoublep) find a better way to do this than whitelist
299
- MAMBA2_MODELS = [
300
- "BambaForCausalLM",
301
- "FalconH1ForCausalLM",
302
- "GraniteMoeHybridForCausalLM",
303
- "Mamba2ForCausalLM",
304
- "NemotronHForCausalLM",
305
- "Zamba2ForCausalLM",
306
- ]
298
if cache_config.enable_prefix_caching:
308
- if model_config.architecture in MAMBA2_MODELS:
299
+ if model_config.supports_mamba_prefix_caching:
300
logger.info(
301
"Warning: Prefix caching is currently enabled. "
302
"Its support for Mamba2 layers is experimental. "
vllm/model_executor/models/falcon_h1.py
+15
-2
@@ -37,7 +37,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
37
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
38
from vllm.sequence import IntermediateTensors
39
40
-from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP
40
+from .interfaces import (
41
+ HasInnerState,
42
+ IsHybrid,
43
+ SupportsLoRA,
44
+ SupportsMambaPrefixCaching,
45
+ SupportsPP,
46
+)
47
from .utils import (
48
PPMissingLayer,
49
is_pp_missing_parameter,
@@ -495,7 +501,14 @@ class FalconH1Model(nn.Module):
501
return hidden_states
502
503
498
-class FalconH1ForCausalLM(nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHybrid):
504
+class FalconH1ForCausalLM(
505
+ nn.Module,
506
+ HasInnerState,
507
+ SupportsLoRA,
508
+ SupportsPP,
509
+ IsHybrid,
510
+ SupportsMambaPrefixCaching,
511
+):
512
packed_modules_mapping = {
513
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
514
"gate_up_proj": ["gate_proj", "up_proj"],
vllm/model_executor/models/granitemoehybrid.py
+15
-2
@@ -34,7 +34,14 @@ from vllm.sequence import IntermediateTensors
34
35
from .granitemoe import GraniteMoeMoE
36
from .granitemoeshared import GraniteMoeSharedMLP
37
-from .interfaces import HasInnerState, IsHybrid, SupportsLoRA, SupportsPP, SupportsQuant
37
+from .interfaces import (
38
+ HasInnerState,
39
+ IsHybrid,
40
+ SupportsLoRA,
41
+ SupportsMambaPrefixCaching,
42
+ SupportsPP,
43
+ SupportsQuant,
44
+)
45
from .utils import (
46
AutoWeightsLoader,
47
is_pp_missing_parameter,
@@ -584,7 +591,13 @@ class GraniteMoeHybridModel(nn.Module):
591
592
593
class GraniteMoeHybridForCausalLM(
587
- nn.Module, HasInnerState, SupportsLoRA, SupportsPP, IsHybrid, SupportsQuant
594
+ nn.Module,
595
+ HasInnerState,
596
+ SupportsLoRA,
597
+ SupportsPP,
598
+ IsHybrid,
599
+ SupportsQuant,
600
+ SupportsMambaPrefixCaching,
601
):
602
packed_modules_mapping = {
603
"qkv_proj": [
vllm/model_executor/models/interfaces.py
+28
@@ -697,6 +697,34 @@ def has_noops(
697
return getattr(model, "has_noops", False)
698
699
700
+@runtime_checkable
701
+class SupportsMambaPrefixCaching(Protocol):
702
+ """The interface for models whose mamba layers support prefix caching.
703
+
704
+ This is currently experimental.
705
+ """
706
+
707
+ supports_mamba_prefix_caching: ClassVar[Literal[True]] = True
708
+
709
+
710
+@overload
711
+def supports_mamba_prefix_caching(
712
+ model: object,
713
+) -> TypeIs[SupportsMambaPrefixCaching]: ...
714
+
715
+
716
+@overload
717
+def supports_mamba_prefix_caching(
718
+ model: type[object],
719
+) -> TypeIs[type[SupportsMambaPrefixCaching]]: ...
720
+
721
+
722
+def supports_mamba_prefix_caching(
723
+ model: type[object] | object,
724
+) -> TypeIs[type[SupportsMambaPrefixCaching]] | TypeIs[SupportsMambaPrefixCaching]:
725
+ return getattr(model, "supports_mamba_prefix_caching", False)
726
+
727
+
728
@runtime_checkable
729
class SupportsCrossEncoding(Protocol):
730
"""The interface required for all models that support cross encoding."""
vllm/model_executor/models/mamba2.py
+8
-2
@@ -25,7 +25,11 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
25
VocabParallelEmbedding,
26
)
27
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
28
-from vllm.model_executor.models.interfaces import HasInnerState, IsAttentionFree
28
+from vllm.model_executor.models.interfaces import (
29
+ HasInnerState,
30
+ IsAttentionFree,
31
+ SupportsMambaPrefixCaching,
32
+)
33
from vllm.sequence import IntermediateTensors
34
35
from .utils import (
@@ -189,7 +193,9 @@ class Mamba2Model(nn.Module):
193
return loaded_params
194
195
192
-class Mamba2ForCausalLM(nn.Module, HasInnerState, IsAttentionFree):
196
+class Mamba2ForCausalLM(
197
+ nn.Module, HasInnerState, IsAttentionFree, SupportsMambaPrefixCaching
198
+):
199
@classmethod
200
def get_mamba_state_dtype_from_config(
201
cls,
vllm/model_executor/models/nemotron_h.py
+2
@@ -62,6 +62,7 @@ from vllm.model_executor.models.interfaces import (
62
IsHybrid,
63
MixtureOfExperts,
64
SupportsLoRA,
65
+ SupportsMambaPrefixCaching,
66
SupportsPP,
67
SupportsQuant,
68
)
@@ -695,6 +696,7 @@ class NemotronHForCausalLM(
696
IsHybrid,
697
SupportsQuant,
698
MixtureOfExperts,
699
+ SupportsMambaPrefixCaching,
700
):
701
hf_to_vllm_mapper = WeightsMapper(
702
orig_to_new_prefix={"backbone": "model"},
vllm/model_executor/models/registry.py
+3
@@ -39,6 +39,7 @@ from .interfaces import (
39
is_attention_free,
40
is_hybrid,
41
supports_cross_encoding,
42
+ supports_mamba_prefix_caching,
43
supports_multimodal,
44
supports_multimodal_encoder_tp_data,
45
supports_multimodal_raw_input_only,
@@ -496,6 +497,7 @@ class _ModelInfo:
497
is_attention_free: bool
498
is_hybrid: bool
499
has_noops: bool
500
+ supports_mamba_prefix_caching: bool
501
supports_transcription: bool
502
supports_transcription_only: bool
503
@@ -518,6 +520,7 @@ class _ModelInfo:
520
has_inner_state=has_inner_state(model),
521
is_attention_free=is_attention_free(model),
522
is_hybrid=is_hybrid(model),
523
+ supports_mamba_prefix_caching=supports_mamba_prefix_caching(model),
524
supports_transcription=supports_transcription(model),
525
supports_transcription_only=(
526
supports_transcription(model) and model.supports_transcription_only
vllm/model_executor/models/zamba2.py
+2
-2
@@ -45,7 +45,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
45
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
46
from vllm.sequence import IntermediateTensors
47
48
-from .interfaces import HasInnerState, IsHybrid
48
+from .interfaces import HasInnerState, IsHybrid, SupportsMambaPrefixCaching
49
from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
50
51
@@ -824,7 +824,7 @@ class Zamba2Model(nn.Module):
824
return loaded_params
825
826
827
-class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid):
827
+class Zamba2ForCausalLM(nn.Module, HasInnerState, IsHybrid, SupportsMambaPrefixCaching):
828
"""Zamba2 model with causal language modeling head.
829
830
This class wraps the core Zamba2 model and adds: