[Bugfix][Attention] Zero reused FlashInfer MLA KV blocks

tests/v1/attention/test_mla_backends.py

@@ -52,6 +52,7 @@
 from vllm.v1.kv_cache_interface import (
     KVQuantMode,
     MLAAttentionSpec,
+    SlidingWindowMLASpec,
 )
 from vllm.v1.worker.block_table import get_block_table_width
 
@@ -1091,6 +1092,68 @@ def test_flashinfer_mla_dspark_dcp_supports_target_and_draft(monkeypatch):
     assert backend.supports_non_causal_dcp()
 
 
+@pytest.mark.skip_global_cleanup
+def test_only_flashinfer_mla_backends_require_kv_cache_zeroing():
+    flashinfer_mla_module = pytest.importorskip(
+        "vllm.v1.attention.backends.mla.flashinfer_mla"
+    )
+    flashinfer_sparse_module = pytest.importorskip(
+        "vllm.v1.attention.backends.mla.flashinfer_mla_sparse"
+    )
+    flashinfer_sparse_sm90_module = pytest.importorskip(
+        "vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm90"
+    )
+    flashmla_sparse_module = pytest.importorskip(
+        "vllm.v1.attention.backends.mla.flashmla_sparse"
+    )
+    deepseek_v4_module = pytest.importorskip(
+        "vllm.models.deepseek_v4.nvidia.flashinfer_sparse"
+    )
+    deepseek_v41_module = pytest.importorskip(
+        "vllm.models.deepseek_v41.nvidia.flashinfer_sparse"
+    )
+    full_spec = MLAAttentionSpec(
+        block_size=64,
+        num_kv_heads=1,
+        head_size=576,
+        dtype=torch.bfloat16,
+    )
+    sliding_spec = SlidingWindowMLASpec(
+        block_size=64,
+        num_kv_heads=1,
+        head_size=576,
+        dtype=torch.bfloat16,
+        sliding_window=128,
+    )
+
+    flashinfer_backends = (
+        (flashinfer_mla_module.FlashInferMLABackend, full_spec),
+        (flashinfer_sparse_module.FlashInferMLASparseTRTLLMBackend, full_spec),
+        (flashinfer_sparse_module.FlashInferMLASparseSM120Backend, full_spec),
+        (flashinfer_sparse_sm90_module.FlashInferMLASparseSM90Backend, full_spec),
+        (deepseek_v4_module.DeepseekV4FlashInferMLASparseBackend, full_spec),
+        (deepseek_v4_module.DeepseekSparseSWAFlashInferBackend, sliding_spec),
+        (deepseek_v41_module.DeepseekV4FlashInferMLASparseBackend, full_spec),
+        (deepseek_v41_module.DeepseekSparseSWAFlashInferBackend, sliding_spec),
+    )
+
+    for backend, spec in flashinfer_backends:
+        assert backend.customize_spec(spec).requires_kv_cache_zeroing
+
+    assert not mla_attention_module.MLACommonBackend.customize_spec(
+        full_spec
+    ).requires_kv_cache_zeroing
+    assert not flashmla_module.FlashMLABackend.customize_spec(
+        full_spec
+    ).requires_kv_cache_zeroing
+    assert not flashmla_sparse_module.FlashMLASparseBackend.customize_spec(
+        full_spec
+    ).requires_kv_cache_zeroing
+    assert not mla_attention_module.MLACommonBackend.customize_spec(
+        sliding_spec
+    ).requires_kv_cache_zeroing
+
+
 @pytest.mark.parametrize(
     ("causal", "tokens_per_decode", "dcp_world_size", "dcp_rank"),
     [

tests/v1/attention/test_mm_prefix.py

@@ -673,6 +673,58 @@ def test_composite_routes_queries_that_need_image_masking(
     assert requires_mm_prefix(common, unclamped_window=True) is unclamped_expected
 
 
+@pytest.mark.skip_global_cleanup
+def test_composite_backend_preserves_kv_cache_zeroing_requirement():
+    from vllm.v1.attention.backend import AttentionBackend
+    from vllm.v1.attention.backends.composite import (
+        MMPrefixAttentionRouting,
+        create_composite_attention_backend,
+    )
+    from vllm.v1.kv_cache_interface import FullAttentionSpec
+
+    class Builder:
+        requires_block_table_width = False
+
+    class DefaultBackend(AttentionBackend):
+        forward_includes_kv_cache_update = False
+
+        @staticmethod
+        def get_name():
+            return "DEFAULT"
+
+        @staticmethod
+        def get_impl_cls():
+            return None
+
+        @staticmethod
+        def get_builder_cls():
+            return Builder
+
+    class ZeroingBackend(DefaultBackend):
+        requires_kv_cache_zeroing = True
+
+        @staticmethod
+        def get_name():
+            return "ZEROING"
+
+    backend = create_composite_attention_backend(
+        DefaultBackend,
+        ZeroingBackend,
+        name="ZeroingCompositeBackend",
+        backend_name="ZEROING_COMPOSITE",
+        module=__name__,
+        routing_policy=MMPrefixAttentionRouting,
+    )
+    spec = FullAttentionSpec(
+        block_size=16,
+        num_kv_heads=1,
+        head_size=64,
+        dtype=torch.bfloat16,
+    )
+
+    assert backend.customize_spec(spec).requires_kv_cache_zeroing
+
+
 @pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA attention kernels")
 @pytest.mark.parametrize("head_size", [256, 512])
 @pytest.mark.parametrize(

tests/v1/core/test_kv_cache_utils.py

@@ -60,6 +60,7 @@
     get_hisparse_gpu_memory_usage,
 )
 from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
     ChunkedLocalAttentionSpec,
     CircularBufferSpec,
     FullAttentionSpec,
@@ -3199,6 +3200,56 @@ def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None):
     )
 
 
+@pytest.mark.parametrize(
+    "spec",
+    [
+        new_kv_cache_spec(),
+        new_sliding_window_spec(),
+        new_mla_spec(),
+        new_swa_mla_spec(),
+    ],
+    ids=["full-attention", "sliding-window", "mla", "sliding-window-mla"],
+)
+@pytest.mark.parametrize("wrap_uniform", [False, True])
+@pytest.mark.parametrize("requires_zeroing", [False, True])
+@pytest.mark.skip_global_cleanup
+def test_attention_kv_cache_zeroing_follows_backend_requirement(
+    spec: AttentionSpec,
+    wrap_uniform: bool,
+    requires_zeroing: bool,
+) -> None:
+    spec = spec.merge([spec, replace(spec, requires_kv_cache_zeroing=requires_zeroing)])
+    group_spec = (
+        UniformTypeKVCacheSpecs(
+            block_size=spec.block_size,
+            kv_cache_specs={
+                "unmarked": replace(spec, requires_kv_cache_zeroing=False),
+                "layer": spec,
+            },
+        )
+        if wrap_uniform
+        else spec
+    )
+    worker_config = KVCacheConfig(
+        num_blocks=10,
+        kv_cache_tensors=[],
+        kv_cache_groups=[
+            KVCacheGroupSpec(
+                list(group_spec.kv_cache_specs)
+                if isinstance(group_spec, UniformTypeKVCacheSpecs)
+                else ["layer"],
+                group_spec,
+            )
+        ],
+    )
+    scheduler_config = generate_scheduler_kv_cache_config([worker_config])
+
+    assert worker_config.has_attention_layers_requiring_zeroing is requires_zeroing
+    assert worker_config.needs_kv_cache_zeroing is requires_zeroing
+    assert scheduler_config.has_attention_layers_requiring_zeroing is requires_zeroing
+    assert scheduler_config.needs_kv_cache_zeroing is requires_zeroing
+
+
 def new_indexer_mla_spec(block_size=16):
     # Sparse-attention indexer k_cache: an MLAAttentionSpec with a much smaller
     # page size than the main MLA attention (uint8, small head), so their pages

vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py

@@ -107,6 +107,7 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):
     Inherits the base and backend reuses its``DeepseekV4SparseMLAMetadataBuilder``
     """
 
+    requires_kv_cache_zeroing: ClassVar[bool] = True
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
     supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
         "auto",
@@ -193,6 +194,8 @@ class DeepseekSparseSWAFlashInferMetadataBuilder(DeepseekSparseSWAMetadataBuilde
 
 
 class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend):
+    requires_kv_cache_zeroing: ClassVar[bool] = True
+
     @staticmethod
     def get_builder_cls() -> type[DeepseekSparseSWAFlashInferMetadataBuilder]:
         return DeepseekSparseSWAFlashInferMetadataBuilder

vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py

@@ -104,6 +104,7 @@ class DeepseekV4FlashInferMLASparseBackend(DeepseekV4SparseMLABackend):
     Inherits the base and backend reuses its``DeepseekV4SparseMLAMetadataBuilder``
     """
 
+    requires_kv_cache_zeroing: ClassVar[bool] = True
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
     supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
         "auto",
@@ -226,6 +227,8 @@ def build(
 
 
 class DeepseekSparseSWAFlashInferBackend(DeepseekSparseSWABackend):
+    requires_kv_cache_zeroing: ClassVar[bool] = True
+
     @staticmethod
     def get_builder_cls() -> type[DeepseekSparseSWAFlashInferMetadataBuilder]:
         return DeepseekSparseSWAFlashInferMetadataBuilder

vllm/v1/attention/backend.py

@@ -64,6 +64,8 @@ class AttentionBackend(ABC):
         "float16",
         "bfloat16",
     ]
+    requires_kv_cache_zeroing: ClassVar[bool] = False
+    """Whether recycled KV cache blocks must be zeroed before reuse."""
 
     # Does attention's forward() include kv cache update?
     forward_includes_kv_cache_update: bool = True
@@ -145,7 +147,9 @@ def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec":
 
         (see: https://github.com/vllm-project/vllm/issues/42449)
         """
-        return spec
+        if not cls.requires_kv_cache_zeroing:
+            return spec
+        return replace(spec, requires_kv_cache_zeroing=True)
 
     @classmethod
     def get_preferred_block_size(cls, default_block_size: int) -> int:

vllm/v1/attention/backends/composite.py

@@ -364,6 +364,10 @@ def set_workspace_buffer(self, workspace_buffer):
 
     class CompositeAttentionBackend(AttentionBackend):
         forward_includes_kv_cache_update = False
+        requires_kv_cache_zeroing = (
+            general_backend.requires_kv_cache_zeroing
+            or causal_backend.requires_kv_cache_zeroing
+        )
         general_backend_cls = general_backend
         causal_backend_cls = causal_backend
 
@@ -499,9 +503,10 @@ def supports_combination(
         @classmethod
         def customize_spec(cls, spec):
             general_spec = general_backend.customize_spec(spec)
-            if general_spec != causal_backend.customize_spec(spec):
+            causal_spec = causal_backend.customize_spec(spec)
+            if general_spec != causal_spec:
                 raise ValueError("Composite backends require the same KV cache spec")
-            return general_spec
+            return super().customize_spec(general_spec)
 
         @classmethod
         def supported_kv_cache_layouts(cls):

vllm/v1/attention/backends/mla/flashinfer_mla.py

@@ -191,6 +191,7 @@ def _build_decode(
 
 
 class FlashInferMLABackend(MLACommonBackend):
+    requires_kv_cache_zeroing: ClassVar[bool] = True
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]
     supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
         "auto",

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py

@@ -45,6 +45,8 @@
 class _FlashInferMLASparseBackendBase(AttentionBackend):
     """Common metadata for concrete FlashInfer sparse MLA backends."""
 
+    requires_kv_cache_zeroing: ClassVar[bool] = True
+
     @staticmethod
     def get_name() -> str:
         return "FLASHINFER_MLA_SPARSE"

vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py

@@ -137,6 +137,7 @@ def _clamp_work_kv_end(
 
 
 class FlashInferMLASparseSM90Backend(AttentionBackend):
+    requires_kv_cache_zeroing: ClassVar[bool] = True
     supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
     supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
         "auto",

vllm/v1/core/kv_cache_utils.py

@@ -2378,9 +2378,18 @@ def generate_scheduler_kv_cache_config(
     for group in cfg.kv_cache_groups:
         if isinstance(group.kv_cache_spec, UniformTypeKVCacheSpecs):
             # All layers in the UniformTypeKVCacheSpecs have the same type,
-            # so use an arbitrary one to initialize the scheduler.
+            # so use one to initialize the scheduler. Prefer a spec requiring
+            # zeroing so this backend capability is not lost when collapsing
+            # per-layer specs.
+            layer_specs = list(group.kv_cache_spec.kv_cache_specs.values())
             group.kv_cache_spec = next(
-                iter(group.kv_cache_spec.kv_cache_specs.values())
+                (
+                    spec
+                    for spec in layer_specs
+                    if isinstance(spec, AttentionSpec)
+                    and spec.requires_kv_cache_zeroing
+                ),
+                layer_specs[0],
             )
     return cfg