[Bugfix][MLA] Restore DSpark cache-group capability under optimized Python (#55234)

Signed-off-by: Summer Yang <girasoleyang@gmail.com>

tests/v1/attention/test_mla_noncausal.py

@@ -20,7 +20,9 @@ class _NonCausalMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
 
 
 def _metadata(
-    query_start_loc: list[int], num_tokens: int | None = None
+    query_start_loc: list[int],
+    num_tokens: int | None = None,
+    causal: bool = False,
 ) -> CommonAttentionMetadata:
     num_reqs = len(query_start_loc) - 1
     num_tokens = query_start_loc[-1] if num_tokens is None else num_tokens
@@ -38,7 +40,7 @@ def _metadata(
             num_reqs, 3
         ),
         slot_mapping=torch.arange(num_tokens),
-        causal=False,
+        causal=causal,
         seq_lens_cpu_upper_bound=None,
     )
 
@@ -50,6 +52,7 @@ def _builder(marked: bool = True) -> _NonCausalMLAMetadataBuilder:
     builder.query_len_support = QueryLenSupport.SINGLE_ONLY
     builder.non_causal_multi_token_decode = marked
     builder.dcp_world_size = 1
+    builder.use_pcp = False
     builder.metadata_cls = MLACommonMetadata
     builder.model_config = SimpleNamespace(
         dtype=torch.bfloat16, get_head_size=lambda: 576
@@ -57,9 +60,19 @@ def _builder(marked: bool = True) -> _NonCausalMLAMetadataBuilder:
     return builder
 
 
-def test_noncausal_block_uses_decode_without_cpu_lengths():
+def test_group_capability_keeps_runtime_causality_per_step():
+    builder = _builder()
+
+    target_metadata = builder.build(0, _metadata([0, 1, 2], causal=True))
+    assert (
+        target_metadata.num_decodes,
+        target_metadata.num_decode_tokens,
+        target_metadata.num_prefills,
+        target_metadata.causal,
+    ) == (2, 2, 0, True)
+
     common_metadata = _metadata([0, 8, 16])
-    metadata = _builder().build(0, common_metadata)
+    metadata = builder.build(0, common_metadata)
 
     assert metadata.num_decodes == 2
     assert metadata.num_decode_tokens == 16
@@ -109,7 +122,7 @@ def test_noncausal_decode_metadata_keeps_live_request_buffers():
     )
 
 
-def test_mla_cache_marker_is_a_group_property():
+def test_mla_cache_marker_is_promoted_to_group_capability():
     kwargs = {
         "block_size": 64,
         "num_kv_heads": 1,
@@ -123,5 +136,4 @@ def test_mla_cache_marker_is_a_group_property():
     assert not MLAAttentionSpec.merge(
         [unmarked, unmarked]
     ).non_causal_multi_token_decode
-    with pytest.raises(AssertionError, match="non_causal_multi_token_decode"):
-        MLAAttentionSpec.merge([marked, unmarked])
+    assert MLAAttentionSpec.merge([unmarked, marked]).non_causal_multi_token_decode

tests/v1/core/test_kv_cache_utils.py

@@ -3,8 +3,11 @@
 import copy
 import hashlib
 import importlib
+import subprocess
+import sys
 from collections.abc import Callable
 from dataclasses import replace
+from pathlib import Path
 from types import SimpleNamespace
 from typing import Any, cast
 
@@ -1463,6 +1466,28 @@ def test_is_kv_cache_spec_uniform():
     }
     assert not is_kv_cache_spec_uniform(kv_cache_spec)
 
+    script = """
+import sys
+
+from vllm.v1.core.kv_cache_utils import is_kv_cache_spec_uniform
+from vllm.v1.kv_cache_interface import KVCacheSpec
+
+if sys.flags.optimize < 1:
+    raise RuntimeError("subprocess is not running with optimization enabled")
+specs = {
+    "a": KVCacheSpec(block_size=1),
+    "b": KVCacheSpec(block_size=2),
+}
+if is_kv_cache_spec_uniform(specs):
+    raise RuntimeError("different specs were treated as uniform")
+"""
+    subprocess.run(
+        [sys.executable, "-O", "-c", script],
+        cwd=Path(__file__).resolve().parents[3],
+        check=True,
+        timeout=60,
+    )
+
 
 @pytest.mark.parametrize(
     ("model_id", "max_model_len", "want_estimated_max_len"),

vllm/v1/kv_cache_interface.py

@@ -224,9 +224,10 @@ class KVCacheSpec:
         """
         Merge a list of KVCacheSpec objects into a single KVCacheSpec object.
         """
-        assert all(spec == specs[0] for spec in specs[1:]), (
-            "All layers in the same KV cache group must be the same."
-        )
+        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 copy.deepcopy(specs[0])
 
     def is_uniform_with_collection(
@@ -558,7 +559,8 @@ class MLAAttentionSpec(FullAttentionSpec):
     model_version: str | None = None
     storage_block_size: int | None = None
     """Token width used to view storage when it differs from the kernel block."""
-    # Marks draft groups that flatten a non-causal query block into decode rows.
+    # Group capability enabled when any member flattens a non-causal query block
+    # into decode rows. Runtime metadata still selects causal vs. non-causal mode.
     non_causal_multi_token_decode: bool = False
     # MLA stores a single latent vector per state; there is no separate V.
     head_size_v: int = 0
@@ -586,11 +588,6 @@ class MLAAttentionSpec(FullAttentionSpec):
             "quantization method, tokens per state, model version, and storage "
             "block size."
         )
-        non_causal_mtd_set = {spec.non_causal_multi_token_decode for spec in specs}
-        assert len(non_causal_mtd_set) == 1, (
-            "All attention layers in the same KV cache group must agree on "
-            "non_causal_multi_token_decode."
-        )
         merged_spec = cls(
             block_size=specs[0].block_size,
             num_kv_heads=specs[0].num_kv_heads,
@@ -604,7 +601,9 @@ class MLAAttentionSpec(FullAttentionSpec):
             tokens_per_state=tokens_per_state_set.pop(),
             model_version=model_version_set.pop(),
             storage_block_size=storage_block_size_set.pop(),
-            non_causal_multi_token_decode=non_causal_mtd_set.pop(),
+            non_causal_multi_token_decode=any(
+                spec.non_causal_multi_token_decode for spec in specs
+            ),
         )
         for spec in specs:
             for f in fields(AttentionSpec):