fix: allow LFM2 MoE prefix caching (align) (vllm-project#33376)

Signed-off-by: Tianshu Yu <[email protected]>

Author: tianshu-Michael-yu

vllm/model_executor/models/lfm2_moe.py

@@ -651,9 +651,11 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
         quant_config = vllm_config.quant_config
         cache_config = vllm_config.cache_config
 
-        assert not cache_config.enable_prefix_caching, (
-            "Lfm2Moe currently does not support prefix caching"
-        )
+        if cache_config.mamba_cache_mode == "all":
+            raise NotImplementedError(
+                "Lfm2Moe currently does not support 'all' prefix caching, "
+                "please use '--mamba-cache-mode=align' instead"
+            )
 
         super().__init__()
         self.config = config

vllm/model_executor/models/lfm2_vl.py

@@ -22,6 +22,8 @@
 from vllm.config.multimodal import BaseDummyOptions
 from vllm.forward_context import set_forward_context
 from vllm.model_executor.layers.mamba.mamba_utils import (
+    MambaStateCopyFunc,
+    MambaStateCopyFuncCalculator,
     MambaStateDtypeCalculator,
     MambaStateShapeCalculator,
 )
@@ -584,6 +586,10 @@ def get_mamba_state_shape_from_config(
             conv_kernel=hf_language_config.conv_L_cache,
         )
 
+    @classmethod
+    def get_mamba_state_copy_func(cls) -> tuple[MambaStateCopyFunc]:
+        return MambaStateCopyFuncCalculator.short_conv_state_copy_func()
+
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = "model"):
         super().__init__()
         config: Lfm2VlConfig = vllm_config.model_config.hf_config