[Bugfix][MLA] Restore DSpark cache-group capability under optimized Python (#55234)
Signed-off-by: Summer Yang <girasoleyang@gmail.com>
Summer Yang committed
Sep 4, 2026 at 10:24 UTC
8cd95f7de70df6ff15002168a048891c708428fd
3 files changed
+53
-17
tests/v1/attention/test_mla_noncausal.py
+19
-7
@@ -20,7 +20,9 @@ class _NonCausalMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]):
20
21
22
def _metadata(
23
- query_start_loc: list[int], num_tokens: int | None = None
23
+ query_start_loc: list[int],
24
+ num_tokens: int | None = None,
25
+ causal: bool = False,
26
) -> CommonAttentionMetadata:
27
num_reqs = len(query_start_loc) - 1
28
num_tokens = query_start_loc[-1] if num_tokens is None else num_tokens
@@ -38,7 +40,7 @@ def _metadata(
40
num_reqs, 3
41
),
42
slot_mapping=torch.arange(num_tokens),
41
- causal=False,
43
+ causal=causal,
44
seq_lens_cpu_upper_bound=None,
45
)
46
@@ -50,6 +52,7 @@ def _builder(marked: bool = True) -> _NonCausalMLAMetadataBuilder:
52
builder.query_len_support = QueryLenSupport.SINGLE_ONLY
53
builder.non_causal_multi_token_decode = marked
54
builder.dcp_world_size = 1
55
+ builder.use_pcp = False
56
builder.metadata_cls = MLACommonMetadata
57
builder.model_config = SimpleNamespace(
58
dtype=torch.bfloat16, get_head_size=lambda: 576
@@ -57,9 +60,19 @@ def _builder(marked: bool = True) -> _NonCausalMLAMetadataBuilder:
60
return builder
61
62
60
-def test_noncausal_block_uses_decode_without_cpu_lengths():
63
+def test_group_capability_keeps_runtime_causality_per_step():
64
+ builder = _builder()
65
+
66
+ target_metadata = builder.build(0, _metadata([0, 1, 2], causal=True))
67
+ assert (
68
+ target_metadata.num_decodes,
69
+ target_metadata.num_decode_tokens,
70
+ target_metadata.num_prefills,
71
+ target_metadata.causal,
72
+ ) == (2, 2, 0, True)
73
+
74
common_metadata = _metadata([0, 8, 16])
62
- metadata = _builder().build(0, common_metadata)
75
+ metadata = builder.build(0, common_metadata)
76
77
assert metadata.num_decodes == 2
78
assert metadata.num_decode_tokens == 16
@@ -109,7 +122,7 @@ def test_noncausal_decode_metadata_keeps_live_request_buffers():
122
)
123
124
112
-def test_mla_cache_marker_is_a_group_property():
125
+def test_mla_cache_marker_is_promoted_to_group_capability():
126
kwargs = {
127
"block_size": 64,
128
"num_kv_heads": 1,
@@ -123,5 +136,4 @@ def test_mla_cache_marker_is_a_group_property():
136
assert not MLAAttentionSpec.merge(
137
[unmarked, unmarked]
138
).non_causal_multi_token_decode
126
- with pytest.raises(AssertionError, match="non_causal_multi_token_decode"):
127
- MLAAttentionSpec.merge([marked, unmarked])
139
+ assert MLAAttentionSpec.merge([unmarked, marked]).non_causal_multi_token_decode
tests/v1/core/test_kv_cache_utils.py
+25
@@ -3,8 +3,11 @@
3
import copy
4
import hashlib
5
import importlib
6
+import subprocess
7
+import sys
8
from collections.abc import Callable
9
from dataclasses import replace
10
+from pathlib import Path
11
from types import SimpleNamespace
12
from typing import Any, cast
13
@@ -1463,6 +1466,28 @@ def test_is_kv_cache_spec_uniform():
1466
}
1467
assert not is_kv_cache_spec_uniform(kv_cache_spec)
1468
1469
+ script = """
1470
+import sys
1471
+
1472
+from vllm.v1.core.kv_cache_utils import is_kv_cache_spec_uniform
1473
+from vllm.v1.kv_cache_interface import KVCacheSpec
1474
+
1475
+if sys.flags.optimize < 1:
1476
+ raise RuntimeError("subprocess is not running with optimization enabled")
1477
+specs = {
1478
+ "a": KVCacheSpec(block_size=1),
1479
+ "b": KVCacheSpec(block_size=2),
1480
+}
1481
+if is_kv_cache_spec_uniform(specs):
1482
+ raise RuntimeError("different specs were treated as uniform")
1483
+"""
1484
+ subprocess.run(
1485
+ [sys.executable, "-O", "-c", script],
1486
+ cwd=Path(__file__).resolve().parents[3],
1487
+ check=True,
1488
+ timeout=60,
1489
+ )
1490
+
1491
1492
@pytest.mark.parametrize(
1493
("model_id", "max_model_len", "want_estimated_max_len"),
vllm/v1/kv_cache_interface.py
+9
-10
@@ -224,9 +224,10 @@ class KVCacheSpec:
224
"""
225
Merge a list of KVCacheSpec objects into a single KVCacheSpec object.
226
"""
227
- assert all(spec == specs[0] for spec in specs[1:]), (
228
- "All layers in the same KV cache group must be the same."
229
- )
227
+ if not all(spec == specs[0] for spec in specs[1:]):
228
+ raise AssertionError(
229
+ "All layers in the same KV cache group must be the same."
230
+ )
231
return copy.deepcopy(specs[0])
232
233
def is_uniform_with_collection(
@@ -558,7 +559,8 @@ class MLAAttentionSpec(FullAttentionSpec):
559
model_version: str | None = None
560
storage_block_size: int | None = None
561
"""Token width used to view storage when it differs from the kernel block."""
561
- # Marks draft groups that flatten a non-causal query block into decode rows.
562
+ # Group capability enabled when any member flattens a non-causal query block
563
+ # into decode rows. Runtime metadata still selects causal vs. non-causal mode.
564
non_causal_multi_token_decode: bool = False
565
# MLA stores a single latent vector per state; there is no separate V.
566
head_size_v: int = 0
@@ -586,11 +588,6 @@ class MLAAttentionSpec(FullAttentionSpec):
588
"quantization method, tokens per state, model version, and storage "
589
"block size."
590
)
589
- non_causal_mtd_set = {spec.non_causal_multi_token_decode for spec in specs}
590
- assert len(non_causal_mtd_set) == 1, (
591
- "All attention layers in the same KV cache group must agree on "
592
- "non_causal_multi_token_decode."
593
- )
591
merged_spec = cls(
592
block_size=specs[0].block_size,
593
num_kv_heads=specs[0].num_kv_heads,
@@ -604,7 +601,9 @@ class MLAAttentionSpec(FullAttentionSpec):
601
tokens_per_state=tokens_per_state_set.pop(),
602
model_version=model_version_set.pop(),
603
storage_block_size=storage_block_size_set.pop(),
607
- non_causal_multi_token_decode=non_causal_mtd_set.pop(),
604
+ non_causal_multi_token_decode=any(
605
+ spec.non_causal_multi_token_decode for spec in specs
606
+ ),
607
)
608
for spec in specs:
609
for f in fields(AttentionSpec):