[Bugfix] Report FULL_ATTENTION for uniform-base UniformTypeKVCacheSpecs groups instead of UNKNOWN (#51218)
Signed-off-by: Yifan Jiang <19356972+yifjiang@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai>
Yifan Jiang committed
Aug 12, 2026 at 18:25 UTC
5936bac7d746dd51ad0f6cf391ce2f990cc4986d
2 files changed
+121
tests/v1/test_kv_cache_spec_registry.py
+112
@@ -28,6 +28,7 @@ from vllm.v1.kv_cache_interface import (
28
FullAttentionSpec,
29
HiddenStateCacheSpec,
30
KVCacheSpec,
31
+ KVCacheSpecKind,
32
MambaSpec,
33
MLAAttentionSpec,
34
SinkFullAttentionSpec,
@@ -35,6 +36,7 @@ from vllm.v1.kv_cache_interface import (
36
SlidingWindowSpec,
37
TQFullAttentionSpec,
38
UniformTypeKVCacheSpecs,
39
+ get_kv_cache_spec_kind,
40
)
41
from vllm.v1.kv_cache_spec_registry import (
42
_REGISTRY_KVCACHESPEC_LIST,
@@ -320,3 +322,113 @@ class TestKVCacheSpecRegistry:
322
)
323
324
assert are_uniform_specs(custom_spec, make_spec(FullAttentionSpec))
325
+
326
+
327
+class TestGetKVCacheSpecKind:
328
+ """get_kv_cache_spec_kind() for merged UniformTypeKVCacheSpecs groups."""
329
+
330
+ @pytest.mark.parametrize(
331
+ "spec_cls,expected_kind",
332
+ [
333
+ (FullAttentionSpec, KVCacheSpecKind.FULL_ATTENTION),
334
+ (MLAAttentionSpec, KVCacheSpecKind.MLA_ATTENTION),
335
+ (SlidingWindowSpec, KVCacheSpecKind.SLIDING_WINDOW),
336
+ ],
337
+ )
338
+ def test_homogeneous_group_keeps_member_kind(self, spec_cls, expected_kind):
339
+ """A group whose members all share one kind reports that kind."""
340
+ spec = make_spec(spec_cls)
341
+ group = UniformTypeKVCacheSpecs(
342
+ block_size=64, kv_cache_specs={"layer_0": spec, "layer_1": spec}
343
+ )
344
+
345
+ assert get_kv_cache_spec_kind(group) == expected_kind
346
+
347
+ @pytest.mark.parametrize("other_cls", [MLAAttentionSpec, SinkFullAttentionSpec])
348
+ def test_mixed_kinds_sharing_full_attention_base(self, other_cls):
349
+ """
350
+ Members with differing kinds but a single registered
351
+ uniform_type_base_spec of FullAttentionSpec report FULL_ATTENTION.
352
+
353
+ This is the shape produced by a model that pairs a main K/V cache with a
354
+ differently shaped auxiliary cache (e.g. a K-only index cache). The two
355
+ merge into one group precisely because they share a base spec, so
356
+ falling back to UNKNOWN would discard what the merge established.
357
+ """
358
+ group = UniformTypeKVCacheSpecs(
359
+ block_size=64,
360
+ kv_cache_specs={
361
+ "layer_0": make_spec(FullAttentionSpec),
362
+ "layer_1": make_spec(other_cls),
363
+ },
364
+ )
365
+
366
+ # Precondition: the members really do differ in kind but share a base.
367
+ assert (
368
+ len({get_kv_cache_spec_kind(s) for s in group.kv_cache_specs.values()}) == 2
369
+ )
370
+ assert {
371
+ KVCacheSpecRegistry.get_uniform_type_base_spec(s)
372
+ for s in group.kv_cache_specs.values()
373
+ } == {FullAttentionSpec}
374
+
375
+ assert get_kv_cache_spec_kind(group) == KVCacheSpecKind.FULL_ATTENTION
376
+
377
+ def test_mixed_base_specs_report_unknown(self):
378
+ """Members resolving to different base specs remain UNKNOWN."""
379
+ group = UniformTypeKVCacheSpecs(
380
+ block_size=64,
381
+ kv_cache_specs={
382
+ "layer_0": make_spec(FullAttentionSpec),
383
+ "layer_1": make_spec(SlidingWindowSpec),
384
+ },
385
+ )
386
+
387
+ assert (
388
+ len(
389
+ {
390
+ KVCacheSpecRegistry.get_uniform_type_base_spec(s)
391
+ for s in group.kv_cache_specs.values()
392
+ }
393
+ )
394
+ == 2
395
+ )
396
+ assert get_kv_cache_spec_kind(group) == KVCacheSpecKind.UNKNOWN
397
+
398
+ def test_uniform_non_full_attention_base_reports_unknown(self):
399
+ """
400
+ The FULL_ATTENTION result is gated on the shared base spec being
401
+ FullAttentionSpec; a uniform base of any other type stays UNKNOWN.
402
+ """
403
+
404
+ # SlidingWindowMLASpec subclasses SlidingWindowSpec but carries a
405
+ # distinct kind, so declaring SlidingWindowSpec as the base gives a
406
+ # group with differing kinds yet a single, non-full-attention base.
407
+ @register_kv_cache_spec(
408
+ manager_class=SlidingWindowManager,
409
+ uniform_type_base_spec=SlidingWindowSpec,
410
+ )
411
+ @dataclass(frozen=True, kw_only=True)
412
+ class _SlidingBasedMLASpec(SlidingWindowMLASpec):
413
+ custom_param: int = 16
414
+
415
+ custom_spec = _SlidingBasedMLASpec(**spec_args_map[SlidingWindowMLASpec])
416
+ group = UniformTypeKVCacheSpecs(
417
+ block_size=64,
418
+ kv_cache_specs={
419
+ "layer_0": custom_spec,
420
+ "layer_1": make_spec(SlidingWindowSpec),
421
+ },
422
+ )
423
+
424
+ # Two differing kinds, but a single shared base that is not
425
+ # FullAttentionSpec -- so the group must not be upgraded.
426
+ assert (
427
+ len({get_kv_cache_spec_kind(s) for s in group.kv_cache_specs.values()}) == 2
428
+ )
429
+ assert {
430
+ KVCacheSpecRegistry.get_uniform_type_base_spec(s)
431
+ for s in group.kv_cache_specs.values()
432
+ } == {SlidingWindowSpec}
433
+
434
+ assert get_kv_cache_spec_kind(group) == KVCacheSpecKind.UNKNOWN
vllm/v1/kv_cache_interface.py
+9
@@ -934,6 +934,15 @@ def get_kv_cache_spec_kind(kv_cache_spec: KVCacheSpec) -> KVCacheSpecKind:
934
}
935
if len(inner_kinds) == 1:
936
return next(iter(inner_kinds))
937
+ # A group is only formed when all members share one registered
938
+ # uniform_type_base_spec, so UNKNOWN would discard what the merge
939
+ # already established.
940
+ base_specs = {
941
+ KVCacheSpecRegistry.get_uniform_type_base_spec(spec)
942
+ for spec in kv_cache_spec.kv_cache_specs.values()
943
+ }
944
+ if len(base_specs) == 1 and next(iter(base_specs)) is FullAttentionSpec:
945
+ return KVCacheSpecKind.FULL_ATTENTION
946
return KVCacheSpecKind.UNKNOWN
947
# Keep subclass checks before base classes so specialized specs keep their
948
# more precise kind.