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 9, 07:30 UTC
fix: allow LFM2 MoE prefix caching (align) (#33376)
Signed-off-by: Tianshu Yu <tianshuyu.formal@gmail.com>
tianshu-Michael-yu committed
Jan 30, 2026 at 00:23 UTC
f45870b53fc4cab2b747b04d5d3c47a68377d548
2 files changed
+11
-3
vllm/model_executor/models/lfm2_moe.py
+5
-3
@@ -651,9 +651,11 @@ class Lfm2MoeForCausalLM(
651
quant_config = vllm_config.quant_config
652
cache_config = vllm_config.cache_config
653
654
- assert not cache_config.enable_prefix_caching, (
655
- "Lfm2Moe currently does not support prefix caching"
656
- )
654
+ if cache_config.mamba_cache_mode == "all":
655
+ raise NotImplementedError(
656
+ "Lfm2Moe currently does not support 'all' prefix caching, "
657
+ "please use '--mamba-cache-mode=align' instead"
658
+ )
659
660
super().__init__()
661
self.config = config
vllm/model_executor/models/lfm2_vl.py
+6
@@ -22,6 +22,8 @@ from vllm.config import VllmConfig
22
from vllm.config.multimodal import BaseDummyOptions
23
from vllm.forward_context import set_forward_context
24
from vllm.model_executor.layers.mamba.mamba_utils import (
25
+ MambaStateCopyFunc,
26
+ MambaStateCopyFuncCalculator,
27
MambaStateDtypeCalculator,
28
MambaStateShapeCalculator,
29
)
@@ -584,6 +586,10 @@ class Lfm2VLForConditionalGeneration(
586
conv_kernel=hf_language_config.conv_L_cache,
587
)
588
589
+ @classmethod
590
+ def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc]:
591
+ return MambaStateCopyFuncCalculator.short_conv_state_copy_func()
592
+
593
def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
594
super().__init__()
595
config: Lfm2VlConfig = vllm_config.model_config.hf_config