From c62582111a7dbf9212e20bea8218975970c9af30 Mon Sep 17 00:00:00 2001 From: Andy Lo Date: Fri, 25 Sep 2026 14:01:43 +0000 Subject: [PATCH 1/2] [Bugfix][Attention] Zero reused FlashInfer MLA KV blocks Co-authored-by: OpenAI Codex Signed-off-by: Andy Lo --- tests/v1/attention/test_mla_backends.py | 62 +++++++++++++++++++ tests/v1/core/test_kv_cache_utils.py | 33 ++++++++++ .../deepseek_v4/nvidia/flashinfer_sparse.py | 3 + .../deepseek_v41/nvidia/flashinfer_sparse.py | 3 + vllm/v1/attention/backend.py | 16 ++++- .../attention/backends/mla/flashinfer_mla.py | 1 + .../backends/mla/flashinfer_mla_sparse.py | 2 + .../mla/flashinfer_mla_sparse_sm90.py | 1 + vllm/v1/kv_cache_interface.py | 32 ++++++++-- 9 files changed, 147 insertions(+), 6 deletions(-) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 9ae0fc216a47..9df11491296e 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/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,67 @@ def test_flashinfer_mla_dspark_dcp_supports_target_and_draft(monkeypatch): assert backend.supports_non_causal_dcp() +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"), [ diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index c2c86053f0f5..c27fc7bfcbdb 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3199,6 +3199,39 @@ def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None): ) +@pytest.mark.parametrize( + "spec", + [new_mla_spec(), new_swa_mla_spec()], + ids=["full-mla", "sliding-window-mla"], +) +@pytest.mark.parametrize("wrap_uniform", [False, True]) +@pytest.mark.parametrize("requires_zeroing", [False, True]) +def test_mla_kv_cache_zeroing_follows_backend_requirement( + spec: MLAAttentionSpec | SlidingWindowMLASpec, + wrap_uniform: bool, + requires_zeroing: bool, +) -> None: + spec = replace(spec, requires_kv_cache_zeroing=requires_zeroing) + group_spec = ( + UniformTypeKVCacheSpecs( + block_size=spec.block_size, kv_cache_specs={"layer": spec} + ) + if wrap_uniform + else spec + ) + worker_config = KVCacheConfig( + num_blocks=10, + kv_cache_tensors=[], + kv_cache_groups=[KVCacheGroupSpec(["layer"], group_spec)], + ) + scheduler_config = generate_scheduler_kv_cache_config([worker_config]) + + assert worker_config.has_mla_layers_requiring_zeroing is requires_zeroing + assert worker_config.needs_kv_cache_zeroing is requires_zeroing + assert scheduler_config.has_mla_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 diff --git a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py index 3e83c280f3c0..6d80a28a13c8 100644 --- a/vllm/models/deepseek_v4/nvidia/flashinfer_sparse.py +++ b/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 diff --git a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py b/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py index 0e02b486a517..6f8d964683b0 100644 --- a/vllm/models/deepseek_v41/nvidia/flashinfer_sparse.py +++ b/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 diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 5d661634d743..95c314b8066c 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -30,7 +30,12 @@ KVQuantMode, ) -from vllm.v1.kv_cache_interface import KVCacheLayout, get_kv_quant_mode +from vllm.v1.kv_cache_interface import ( + KVCacheLayout, + MLAAttentionSpec, + SlidingWindowMLASpec, + get_kv_quant_mode, +) class AttentionType(str, Enum): @@ -64,6 +69,7 @@ class AttentionBackend(ABC): "float16", "bfloat16", ] + requires_kv_cache_zeroing: ClassVar[bool] = False # Does attention's forward() include kv cache update? forward_includes_kv_cache_update: bool = True @@ -145,7 +151,13 @@ 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 + assert isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)), ( + f"{cls.__name__} requires KV cache zeroing but received " + f"unsupported spec type {type(spec).__name__}" + ) + return replace(spec, requires_kv_cache_zeroing=True) @classmethod def get_preferred_block_size(cls, default_block_size: int) -> int: diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py index ca78961eb591..6c48cbbe1e7c 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla.py +++ b/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", diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index 5af50ef1bd21..393d197c03ba 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/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" diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py index a5f7757d3887..604b2b1a32c4 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm90.py +++ b/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", diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 5791704b2bfd..9329d3461f58 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -656,6 +656,8 @@ def _apply_alignment_padding(spec: MLAAttentionSpec | SlidingWindowMLASpec): class MLAAttentionSpec(FullAttentionSpec): # TODO(Lucas/Chen): less hacky way to do this cache_dtype_str: str | None = None + requires_kv_cache_zeroing: bool = False + """Whether the selected backend requires clean pages before cache reuse.""" # DeepseekV4 only fields. Non-DeepseekV4 MLA models leave these at defaults. alignment: int | None = None # Default to None for no padding. model_version: str | None = None @@ -709,6 +711,9 @@ def merge(cls, specs: list[Self]) -> Self: num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, cache_dtype_str=cache_dtype_str_set.pop(), + requires_kv_cache_zeroing=any( + spec.requires_kv_cache_zeroing for spec in specs + ), tokens_per_state=tokens_per_state_set.pop(), model_version=model_version_set.pop(), cache_role=cache_role_set.pop(), @@ -923,6 +928,8 @@ class SlidingWindowMLASpec(SlidingWindowSpec): """Sliding window attention with MLA cache format.""" cache_dtype_str: str | None = None + requires_kv_cache_zeroing: bool = False + """Whether the selected backend requires clean pages before cache reuse.""" # DeepseekV4-only: see MLAAttentionSpec.model_version. alignment: int | None = None # Default to None for no padding. model_version: str | None = None @@ -983,6 +990,9 @@ def merge(cls, specs: list[Self]) -> Self: page_size_padded=specs[0].page_size_padded, num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, + requires_kv_cache_zeroing=any( + spec.requires_kv_cache_zeroing for spec in specs + ), sliding_window=sliding_window_set.pop(), extra_retained_tokens=extra_retained_set.pop(), cache_dtype_str=cache_dtype_str_set.pop(), @@ -1559,6 +1569,15 @@ def has_mamba_layers(self) -> bool: for spec in iter_layer_specs(group.kv_cache_spec) ) + @property + def has_mla_layers_requiring_zeroing(self) -> bool: + return any( + isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)) + and spec.requires_kv_cache_zeroing + for group in self.kv_cache_groups + for spec in iter_layer_specs(group.kv_cache_spec) + ) + @property def has_mixed_precision_kv_cache(self) -> bool: """Whether device attention caches use more than one precision.""" @@ -1578,8 +1597,13 @@ def needs_kv_cache_zeroing(self) -> bool: """Whether newly allocated KV cache blocks must be zeroed before use. Required for Mamba layers, whose state is read before it is fully written - (#35219), and for mixed-precision caches, where a block reused across - groups can be reinterpreted under a different precision and decode stale - bytes to NaN/Inf. Uniform-precision caches skip zeroing. + (#35219), attention backends whose kernels can consume non-finite values + from an unwritten page suffix, and mixed-precision caches, where a block + reused across groups can be reinterpreted under a different precision and + decode stale bytes to NaN/Inf. """ - return self.has_mamba_layers or self.has_mixed_precision_kv_cache + return ( + self.has_mamba_layers + or self.has_mla_layers_requiring_zeroing + or self.has_mixed_precision_kv_cache + ) From 85262f0b64bd9b34a000385d8b36591b5057541b Mon Sep 17 00:00:00 2001 From: Andy Lo Date: Sun, 27 Sep 2026 15:56:02 +0000 Subject: [PATCH 2/2] [Refactor][Attention] Generalize KV cache zeroing capability Co-authored-by: OpenAI Codex Signed-off-by: Andy Lo --- tests/v1/attention/test_mla_backends.py | 1 + tests/v1/attention/test_mm_prefix.py | 52 +++++++++++++++++++++++++ tests/v1/core/test_kv_cache_utils.py | 36 ++++++++++++----- vllm/v1/attention/backend.py | 12 +----- vllm/v1/attention/backends/composite.py | 9 ++++- vllm/v1/core/kv_cache_utils.py | 13 ++++++- vllm/v1/kv_cache_interface.py | 39 +++++++++++++++---- 7 files changed, 131 insertions(+), 31 deletions(-) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 9df11491296e..6fd12a15ff22 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -1092,6 +1092,7 @@ 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" diff --git a/tests/v1/attention/test_mm_prefix.py b/tests/v1/attention/test_mm_prefix.py index f878ede4865f..2d44cbc3fdbc 100644 --- a/tests/v1/attention/test_mm_prefix.py +++ b/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( diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index c27fc7bfcbdb..9d529da35331 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/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, @@ -3201,20 +3202,30 @@ def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None): @pytest.mark.parametrize( "spec", - [new_mla_spec(), new_swa_mla_spec()], - ids=["full-mla", "sliding-window-mla"], + [ + 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]) -def test_mla_kv_cache_zeroing_follows_backend_requirement( - spec: MLAAttentionSpec | SlidingWindowMLASpec, +@pytest.mark.skip_global_cleanup +def test_attention_kv_cache_zeroing_follows_backend_requirement( + spec: AttentionSpec, wrap_uniform: bool, requires_zeroing: bool, ) -> None: - spec = replace(spec, requires_kv_cache_zeroing=requires_zeroing) + spec = spec.merge([spec, replace(spec, requires_kv_cache_zeroing=requires_zeroing)]) group_spec = ( UniformTypeKVCacheSpecs( - block_size=spec.block_size, kv_cache_specs={"layer": spec} + block_size=spec.block_size, + kv_cache_specs={ + "unmarked": replace(spec, requires_kv_cache_zeroing=False), + "layer": spec, + }, ) if wrap_uniform else spec @@ -3222,13 +3233,20 @@ def test_mla_kv_cache_zeroing_follows_backend_requirement( worker_config = KVCacheConfig( num_blocks=10, kv_cache_tensors=[], - kv_cache_groups=[KVCacheGroupSpec(["layer"], group_spec)], + 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_mla_layers_requiring_zeroing is requires_zeroing + 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_mla_layers_requiring_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 diff --git a/vllm/v1/attention/backend.py b/vllm/v1/attention/backend.py index 95c314b8066c..6cf8c3efabf2 100644 --- a/vllm/v1/attention/backend.py +++ b/vllm/v1/attention/backend.py @@ -30,12 +30,7 @@ KVQuantMode, ) -from vllm.v1.kv_cache_interface import ( - KVCacheLayout, - MLAAttentionSpec, - SlidingWindowMLASpec, - get_kv_quant_mode, -) +from vllm.v1.kv_cache_interface import KVCacheLayout, get_kv_quant_mode class AttentionType(str, Enum): @@ -70,6 +65,7 @@ class AttentionBackend(ABC): "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 @@ -153,10 +149,6 @@ def customize_spec(cls, spec: "AttentionSpec") -> "AttentionSpec": """ if not cls.requires_kv_cache_zeroing: return spec - assert isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)), ( - f"{cls.__name__} requires KV cache zeroing but received " - f"unsupported spec type {type(spec).__name__}" - ) return replace(spec, requires_kv_cache_zeroing=True) @classmethod diff --git a/vllm/v1/attention/backends/composite.py b/vllm/v1/attention/backends/composite.py index 0ba40e99f20a..4bf9df6175d7 100644 --- a/vllm/v1/attention/backends/composite.py +++ b/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): diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 3f4f4e7bb164..ee44640cf046 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/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 diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 9329d3461f58..518301df3c1f 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -487,6 +487,8 @@ class AttentionSpec(KVCacheSpec): num_kv_heads: int head_size: int dtype: torch.dtype + requires_kv_cache_zeroing: bool = field(default=False, compare=False) + """Whether the selected backend requires clean pages before cache reuse.""" head_size_v: int = None # type: ignore[assignment] kv_quant_mode: KVQuantMode = KVQuantMode.NONE page_size_padded: int | None = None @@ -506,6 +508,19 @@ def __post_init__(self): if self.head_size_v is None: object.__setattr__(self, "head_size_v", self.head_size) + @classmethod + def merge(cls, specs: list[Self]) -> Self: + if not all(spec == specs[0] for spec in specs[1:]): + raise AssertionError( + "All layers in the same KV cache group must be the same." + ) + return replace( + specs[0], + requires_kv_cache_zeroing=any( + spec.requires_kv_cache_zeroing for spec in specs + ), + ) + @property def num_heads(self) -> int: if self.num_head_slots is not None: @@ -622,6 +637,9 @@ def merge(cls, specs: list[Self]) -> Self: num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, tokens_per_state=specs[0].tokens_per_state, + requires_kv_cache_zeroing=any( + spec.requires_kv_cache_zeroing for spec in specs + ), sliding_window=cls.merge_window_sizes(sliding_window), attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), # If any layer in the group is non-causal, treat the group as @@ -630,6 +648,8 @@ def merge(cls, specs: list[Self]) -> Self: ) for spec in specs: for f in fields(AttentionSpec): + if f.name == "requires_kv_cache_zeroing": + continue assert getattr(spec, f.name) == getattr(merged_spec, f.name), ( "All attention layers in the same KV cache group must have " "the same attention spec." @@ -656,8 +676,6 @@ def _apply_alignment_padding(spec: MLAAttentionSpec | SlidingWindowMLASpec): class MLAAttentionSpec(FullAttentionSpec): # TODO(Lucas/Chen): less hacky way to do this cache_dtype_str: str | None = None - requires_kv_cache_zeroing: bool = False - """Whether the selected backend requires clean pages before cache reuse.""" # DeepseekV4 only fields. Non-DeepseekV4 MLA models leave these at defaults. alignment: int | None = None # Default to None for no padding. model_version: str | None = None @@ -726,6 +744,8 @@ def merge(cls, specs: list[Self]) -> Self: ) for spec in specs: for f in fields(AttentionSpec): + if f.name == "requires_kv_cache_zeroing": + continue assert getattr(spec, f.name) == getattr(merged_spec, f.name), ( "All attention layers in the same KV cache group must have " "the same attention spec." @@ -777,6 +797,7 @@ def merge(cls, specs: list[RSWASpec]) -> RSWASpec: num_head_slots=base.num_head_slots, state_content_bytes=base.state_content_bytes, tokens_per_state=base.tokens_per_state, + requires_kv_cache_zeroing=base.requires_kv_cache_zeroing, sliding_window=base.sliding_window, attention_chunk_size=base.attention_chunk_size, non_causal=base.non_causal, @@ -928,8 +949,6 @@ class SlidingWindowMLASpec(SlidingWindowSpec): """Sliding window attention with MLA cache format.""" cache_dtype_str: str | None = None - requires_kv_cache_zeroing: bool = False - """Whether the selected backend requires clean pages before cache reuse.""" # DeepseekV4-only: see MLAAttentionSpec.model_version. alignment: int | None = None # Default to None for no padding. model_version: str | None = None @@ -1210,12 +1229,17 @@ def merge(cls, specs: list[Self]) -> Self: page_size_padded=specs[0].page_size_padded, num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, + requires_kv_cache_zeroing=any( + spec.requires_kv_cache_zeroing for spec in specs + ), sliding_window=cls.merge_window_sizes(sliding_window), attention_chunk_size=cls.merge_window_sizes(attention_chunk_size), non_causal=any(spec.non_causal for spec in specs), ) for spec in specs: for f in fields(AttentionSpec): + if f.name == "requires_kv_cache_zeroing": + continue assert getattr(spec, f.name) == getattr(merged_spec, f.name), ( "All attention layers in the same KV cache group must have " "the same attention spec." @@ -1570,10 +1594,9 @@ def has_mamba_layers(self) -> bool: ) @property - def has_mla_layers_requiring_zeroing(self) -> bool: + def has_attention_layers_requiring_zeroing(self) -> bool: return any( - isinstance(spec, (MLAAttentionSpec, SlidingWindowMLASpec)) - and spec.requires_kv_cache_zeroing + isinstance(spec, AttentionSpec) and spec.requires_kv_cache_zeroing for group in self.kv_cache_groups for spec in iter_layer_specs(group.kv_cache_spec) ) @@ -1604,6 +1627,6 @@ def needs_kv_cache_zeroing(self) -> bool: """ return ( self.has_mamba_layers - or self.has_mla_layers_requiring_zeroing + or self.has_attention_layers_requiring_zeroing or self.has_mixed_precision_kv_cache )