[KV-Offloading][TP] : Expand replicated_layout detection to multi-group MLA (#57652)
Varun Sundar Rabindranath committed
Sep 30, 2026 at 09:11 UTC
5463fe4962785cdc3383477bf3af6533e7647dfd
10 files changed
+547
-85
tests/v1/core/test_kv_cache_utils.py
+97
@@ -52,6 +52,7 @@ from vllm.v1.core.kv_cache_utils import (
52
hash_block_tokens,
53
init_none_hash,
54
is_kv_cache_spec_uniform,
55
+ kv_cache_groups_tp_replicas,
56
make_block_hash_with_group_id,
57
tensor_data,
58
to_event_extra_keys,
@@ -3274,6 +3275,7 @@ def new_mla_spec(cache_dtype_str=None, block_size: int = 16):
3275
return MLAAttentionSpec(
3276
block_size=block_size,
3277
num_kv_heads=1,
3278
+ max_tp_shards=1,
3279
head_size=576,
3280
dtype=torch.float32,
3281
cache_dtype_str=cache_dtype_str,
@@ -3284,6 +3286,7 @@ def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None):
3286
return SlidingWindowMLASpec(
3287
block_size=16,
3288
num_kv_heads=1,
3289
+ max_tp_shards=1,
3290
head_size=head_size,
3291
dtype=torch.float32,
3292
sliding_window=sliding_window,
@@ -4637,3 +4640,97 @@ def test_trailing_layer_fallback_requires_exact_partition():
4640
_annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True)
4641
4642
assert not any(g.is_eagle_group for g in trimmed)
4643
+
4644
+
4645
+_GQA_SPEC = FullAttentionSpec(
4646
+ block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32
4647
+)
4648
+_GQA_SWA_SPEC = SlidingWindowSpec(
4649
+ block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32, sliding_window=128
4650
+)
4651
+
4652
+
4653
+@pytest.mark.parametrize(
4654
+ "specs,expected",
4655
+ [
4656
+ pytest.param({"l.0": new_mla_spec(), "l.1": new_mla_spec()}, 4, id="mla"),
4657
+ pytest.param(
4658
+ {
4659
+ "l.0": new_mla_spec(),
4660
+ "l.1": new_mla_spec(),
4661
+ "s.0": new_swa_mla_spec(),
4662
+ "s.1": new_swa_mla_spec(),
4663
+ },
4664
+ 4,
4665
+ id="mla-and-swa-mla",
4666
+ ),
4667
+ pytest.param({"l.0": _GQA_SPEC, "l.1": _GQA_SPEC}, 1, id="gqa"),
4668
+ pytest.param(
4669
+ {"l.0": new_mla_spec(), "l.1": _GQA_SPEC},
4670
+ 1,
4671
+ id="mla-and-gqa-uniform-group",
4672
+ ),
4673
+ pytest.param(
4674
+ {
4675
+ "l.0": new_mla_spec(),
4676
+ "l.1": new_mla_spec(),
4677
+ "s.0": _GQA_SWA_SPEC,
4678
+ "s.1": _GQA_SWA_SPEC,
4679
+ },
4680
+ 1,
4681
+ id="mla-and-gqa-swa",
4682
+ ),
4683
+ ],
4684
+)
4685
+def test_kv_tp_replicas(monkeypatch, specs, expected):
4686
+ """Resolved per layer before scheduler flattening hides mixed groups."""
4687
+ from vllm.config import ParallelConfig
4688
+ from vllm.platforms import current_platform
4689
+
4690
+ monkeypatch.setattr(current_platform, "device_count", lambda: 4)
4691
+ parallel_config = ParallelConfig(tensor_parallel_size=4)
4692
+ vllm_config = VllmConfig(
4693
+ model_config=ModelConfig(max_model_len=16), parallel_config=parallel_config
4694
+ )
4695
+ vllm_config.cache_config.kv_cache_layout = "LBNHC"
4696
+ mem = sum(s.page_size_bytes for s in specs.values()) * 10
4697
+ configs = get_kv_cache_configs(vllm_config, [specs], [mem])
4698
+ assert configs[0].kv_tp_replicas == expected
4699
+ scheduler = generate_scheduler_kv_cache_config(configs)
4700
+ assert scheduler.kv_tp_replicas == expected
4701
+
4702
+
4703
+@pytest.mark.parametrize(
4704
+ "specs,tp_size,dcp_size,expected",
4705
+ [
4706
+ pytest.param([new_mla_spec()], 8, 1, 8, id="mla"),
4707
+ pytest.param([new_mla_spec()], 8, 2, 1, id="mla-dcp"),
4708
+ pytest.param(
4709
+ [replace(new_mla_spec(), max_tp_shards=None)], 8, 1, 1, id="mla-unset"
4710
+ ),
4711
+ pytest.param([replace(_GQA_SPEC, max_tp_shards=2)], 8, 1, 4, id="gqa-partial"),
4712
+ pytest.param([replace(_GQA_SPEC, max_tp_shards=8)], 4, 1, 1, id="gqa-sharded"),
4713
+ pytest.param(
4714
+ [new_mla_spec(), replace(_GQA_SWA_SPEC, max_tp_shards=4)],
4715
+ 8,
4716
+ 1,
4717
+ 2,
4718
+ id="mla-and-gqa-partial",
4719
+ ),
4720
+ pytest.param(
4721
+ [
4722
+ HiddenStateCacheSpec(
4723
+ block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32
4724
+ )
4725
+ ],
4726
+ 8,
4727
+ 1,
4728
+ 1,
4729
+ id="hidden-state",
4730
+ ),
4731
+ ],
4732
+)
4733
+def test_kv_cache_groups_tp_replicas(specs, tp_size, dcp_size, expected):
4734
+ """Replicas are the gcd of each layer's tp_size // max_tp_shards."""
4735
+ groups = [KVCacheGroupSpec([f"l.{i}"], spec) for i, spec in enumerate(specs)]
4736
+ assert kv_cache_groups_tp_replicas(groups, tp_size, dcp_size) == expected
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+399
-75
@@ -17,7 +17,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
17
SchedulerOffloadConfig,
18
)
19
from vllm.platforms import current_platform
20
-from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config
20
+from vllm.v1.core.kv_cache_utils import (
21
+ generate_scheduler_kv_cache_config,
22
+ kv_cache_groups_tp_replicas,
23
+)
24
from vllm.v1.kv_cache_interface import (
25
CircularBufferSpec,
26
FullAttentionSpec,
@@ -164,6 +167,7 @@ def _mla_spec(
167
return MLAAttentionSpec(
168
block_size=block_size,
169
num_kv_heads=1,
170
+ max_tp_shards=1,
171
head_size=head_size,
172
dtype=dtype,
173
tokens_per_state=tokens_per_state,
@@ -321,7 +325,184 @@ def _replicated_layout(
325
config.parallel_config.nnodes = nnodes
326
if world_size is not None:
327
config.parallel_config.world_size = world_size
324
- return build_offloading_config(config, kv_cache_config).replicated_layout
328
+ kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas(
329
+ kv_cache_config.kv_cache_groups,
330
+ config.parallel_config.tensor_parallel_size,
331
+ config.parallel_config.decode_context_parallel_size,
332
+ )
333
+
334
+ # Scheduler and Worker offload config
335
+ worker_offload_config = build_offloading_config(config, kv_cache_config)
336
+ scheduler_offload_config = build_offloading_config(
337
+ config, generate_scheduler_kv_cache_config([kv_cache_config])
338
+ )
339
+ # Verify that the Worker and Scheduler agree on replication.
340
+ # and worker-bytes-per-block
341
+ assert (
342
+ worker_offload_config.replicated_layout
343
+ == scheduler_offload_config.replicated_layout
344
+ )
345
+ assert (
346
+ worker_offload_config.worker_kv_bytes_per_block
347
+ == scheduler_offload_config.worker_kv_bytes_per_block
348
+ )
349
+
350
+ return worker_offload_config.replicated_layout
351
+
352
+
353
+_SWA_MLA_PAGE = SlidingWindowMLASpec(
354
+ block_size=16,
355
+ num_kv_heads=1,
356
+ max_tp_shards=1,
357
+ head_size=512,
358
+ dtype=torch.float32,
359
+ sliding_window=128,
360
+).page_size_bytes
361
+
362
+
363
+def _make_swa_mla_kv_cache_config(
364
+ num_blocks: int = 4,
365
+) -> KVCacheConfig:
366
+ """Single SlidingWindowMLASpec group — DSV4 sliding-window layers."""
367
+ spec = SlidingWindowMLASpec(
368
+ block_size=16,
369
+ num_kv_heads=1,
370
+ max_tp_shards=1,
371
+ head_size=512,
372
+ dtype=torch.float32,
373
+ sliding_window=128,
374
+ )
375
+ return KVCacheConfig(
376
+ num_blocks=num_blocks,
377
+ kv_cache_tensors=[
378
+ KVCacheTensor(
379
+ size=_SWA_MLA_PAGE * num_blocks,
380
+ layers=["layer"],
381
+ layer_stride=_SWA_MLA_PAGE * num_blocks,
382
+ block_stride=_SWA_MLA_PAGE,
383
+ )
384
+ ],
385
+ kv_cache_groups=[KVCacheGroupSpec(["layer"], spec)],
386
+ )
387
+
388
+
389
+def _make_dsv4_flash_kv_cache_config(num_blocks: int = 4) -> KVCacheConfig:
390
+ """Minimal fake of DeepSeek-V4-Flash's five UniformTypeKVCacheSpecs groups.
391
+
392
+ The real model has the following groups (layer counts from runtime logs);
393
+ 2 per group is enough to exercise the structural properties that govern
394
+ replicated_layout:
395
+ group[0]: 22 SlidingWindowMLASpec layers, sliding_window=128
396
+ group[1]: 21 SlidingWindowMLASpec layers, sliding_window=128
397
+ group[2]: 62 MLAAttentionSpec layers (MLA + indexer.k_cache)
398
+ group[3]: 42 SlidingWindowMLASpec layers, sliding_window=8
399
+ group[4]: 20 SlidingWindowMLASpec layers, sliding_window=128
400
+ All inner specs have num_kv_heads=1 → every group is TP-replicated.
401
+ """
402
+ block_size = 16
403
+ dtype = torch.float32
404
+ head_size = 512
405
+
406
+ def _swa_group(tag: str, sw: int) -> KVCacheGroupSpec:
407
+ names = [f"{tag}_0", f"{tag}_1"]
408
+ spec = SlidingWindowMLASpec(
409
+ block_size=block_size,
410
+ num_kv_heads=1,
411
+ max_tp_shards=1,
412
+ head_size=head_size,
413
+ dtype=dtype,
414
+ sliding_window=sw,
415
+ )
416
+ return KVCacheGroupSpec(
417
+ names,
418
+ UniformTypeKVCacheSpecs(
419
+ block_size=block_size,
420
+ kv_cache_specs={name: spec for name in names},
421
+ ),
422
+ )
423
+
424
+ def _mla_group(tag: str) -> KVCacheGroupSpec:
425
+ names = [f"{tag}_0", f"{tag}_1"]
426
+ spec = MLAAttentionSpec(
427
+ block_size=block_size,
428
+ num_kv_heads=1,
429
+ max_tp_shards=1,
430
+ head_size=head_size,
431
+ dtype=dtype,
432
+ )
433
+ return KVCacheGroupSpec(
434
+ names,
435
+ UniformTypeKVCacheSpecs(
436
+ block_size=block_size,
437
+ kv_cache_specs={name: spec for name in names},
438
+ ),
439
+ )
440
+
441
+ groups = [
442
+ _swa_group("swa_128a", 128), # group[0]: SWA-MLA, sw=128
443
+ _swa_group("swa_128b", 128), # group[1]: SWA-MLA, sw=128
444
+ _mla_group("mla"), # group[2]: MLA + indexer layers
445
+ _swa_group("swa_8", 8), # group[3]: SWA-MLA, sw=8
446
+ _swa_group("swa_128c", 128), # group[4]: SWA-MLA, sw=128
447
+ ]
448
+ # All specs share the same page size (num_kv_heads=1, head_size=512, fp32).
449
+ page = MLAAttentionSpec(
450
+ block_size=block_size, num_kv_heads=1, head_size=head_size, dtype=dtype
451
+ ).page_size_bytes
452
+ all_layers = [name for g in groups for name in g.layer_names]
453
+ return KVCacheConfig(
454
+ num_blocks=num_blocks,
455
+ kv_cache_tensors=[
456
+ KVCacheTensor(
457
+ size=page * len(all_layers) * num_blocks,
458
+ layers=all_layers,
459
+ layer_stride=page * num_blocks,
460
+ block_stride=page,
461
+ )
462
+ ],
463
+ kv_cache_groups=groups,
464
+ )
465
+
466
+
467
+def _make_dsv3_2_kv_cache_config(num_blocks: int = 4) -> KVCacheConfig:
468
+ """Minimal fake of DeepSeek-V3.2-Exp's single UniformTypeKVCacheSpecs group.
469
+
470
+ The real model has 1 group (layer counts from runtime logs):
471
+ group[0]: 122 MLAAttentionSpec layers (self_attn.attn + indexer.k_cache)
472
+ All inner specs have num_kv_heads=1 → the group is TP-replicated.
473
+
474
+ GLM-5.3 has the same topology (one UniformType group of MLAAttentionSpec
475
+ layers, num_kv_heads=1) with 99 layers instead of 122, so this config
476
+ covers both models.
477
+ """
478
+ block_size = 16
479
+ names = [f"mla_{i}" for i in range(4)] # 4 layers; real model has 122
480
+ spec = MLAAttentionSpec(
481
+ block_size=block_size,
482
+ num_kv_heads=1,
483
+ max_tp_shards=1,
484
+ head_size=512,
485
+ dtype=torch.float32,
486
+ )
487
+ group = KVCacheGroupSpec(
488
+ names,
489
+ UniformTypeKVCacheSpecs(
490
+ block_size=block_size,
491
+ kv_cache_specs={name: spec for name in names},
492
+ ),
493
+ )
494
+ return KVCacheConfig(
495
+ num_blocks=num_blocks,
496
+ kv_cache_tensors=[
497
+ KVCacheTensor(
498
+ size=spec.page_size_bytes * len(names) * num_blocks,
499
+ layers=names,
500
+ layer_stride=spec.page_size_bytes * num_blocks,
501
+ block_stride=spec.page_size_bytes,
502
+ )
503
+ ],
504
+ kv_cache_groups=[group],
505
+ )
506
507
508
@pytest.mark.parametrize("packed", [False, True])
@@ -588,6 +769,61 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
769
)
770
771
772
+@pytest.mark.parametrize("world_size", [2, 4])
773
+def test_replicated_layout_enabled_for_swa_mla(world_size: int):
774
+ """SlidingWindowMLASpec groups are also replicated and qualify."""
775
+ assert _replicated_layout(
776
+ _make_swa_mla_kv_cache_config(),
777
+ tensor_parallel_size=world_size,
778
+ )
779
+
780
+
781
+@pytest.mark.parametrize("world_size", [2, 4, 8])
782
+def test_replicated_layout_dsv4_flash(world_size: int):
783
+ """DeepSeek-V4-Flash: all 5 UniformType groups are MLA/SWA-MLA with
784
+ num_kv_heads=1, so replicated_layout must be True under TP/mp single-node."""
785
+ assert _replicated_layout(
786
+ _make_dsv4_flash_kv_cache_config(),
787
+ tensor_parallel_size=world_size,
788
+ )
789
+
790
+
791
+@pytest.mark.parametrize("world_size", [2, 4, 8])
792
+def test_replicated_layout_dsv3_2(world_size: int):
793
+ """DeepSeek-V3.2 / GLM-5.3: single UniformType group of MLAAttentionSpec
794
+ layers (122 and 99 respectively), num_kv_heads=1 → replicated_layout must
795
+ be True under TP/mp single-node."""
796
+ assert _replicated_layout(
797
+ _make_dsv3_2_kv_cache_config(),
798
+ tensor_parallel_size=world_size,
799
+ )
800
+
801
+
802
+def test_replicated_layout_enabled_for_uniform_wrapper_mla():
803
+ """UniformTypeKVCacheSpecs wrapping MLA layers of different sizes qualifies:
804
+ each layer is replicated and the page accounting sums correctly."""
805
+ mla_layers = ["layer0", "layer1"]
806
+ spec = UniformTypeKVCacheSpecs(
807
+ block_size=16,
808
+ kv_cache_specs={"layer0": _mla_spec(), "layer1": _mla_spec(head_size=256)},
809
+ )
810
+ total_page = spec.page_size_bytes
811
+ num_blocks = 4
812
+ kv_cache_config = KVCacheConfig(
813
+ num_blocks=num_blocks,
814
+ kv_cache_tensors=[
815
+ KVCacheTensor(
816
+ size=total_page * num_blocks,
817
+ layers=mla_layers,
818
+ layer_stride=total_page * num_blocks,
819
+ block_stride=total_page,
820
+ )
821
+ ],
822
+ kv_cache_groups=[KVCacheGroupSpec(mla_layers, spec)],
823
+ )
824
+ assert _replicated_layout(kv_cache_config)
825
+
826
+
827
@pytest.mark.parametrize(
828
("kv_cache_config", "case"),
829
[
@@ -605,18 +841,84 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
841
kv_cache_groups=[
842
KVCacheGroupSpec(
843
["layer"],
608
- SlidingWindowMLASpec(
844
+ HiddenStateCacheSpec(
845
block_size=16,
846
num_kv_heads=1,
847
head_size=512,
848
dtype=torch.float32,
613
- sliding_window=128,
849
),
850
)
851
],
852
),
618
- "sliding-window-mla",
853
+ "hidden-state",
854
+ ),
855
+ (
856
+ # Overlaid groups with different page sizes: a block is a window of
857
+ # the largest group's packing.
858
+ KVCacheConfig(
859
+ num_blocks=4,
860
+ kv_cache_tensors=[
861
+ KVCacheTensor(
862
+ size=_FULL_PAGE * 4,
863
+ layers=["mla"],
864
+ layer_stride=_MLA_PAGE,
865
+ block_stride=_FULL_PAGE,
866
+ ),
867
+ KVCacheTensor(
868
+ size=_FULL_PAGE * 4,
869
+ layers=["full"],
870
+ layer_stride=_FULL_PAGE,
871
+ block_stride=_FULL_PAGE,
872
+ ),
873
+ ],
874
+ kv_cache_groups=[
875
+ KVCacheGroupSpec(["mla"], _mla_spec()),
876
+ KVCacheGroupSpec(["full"], _full_attention_spec()),
877
+ ],
878
+ ),
879
+ "mla-full-hybrid",
880
),
881
+ (
882
+ KVCacheConfig(
883
+ num_blocks=4,
884
+ kv_cache_tensors=[
885
+ KVCacheTensor(
886
+ size=_MLA_PAGE * 4,
887
+ layers=["mla"],
888
+ layer_stride=_MLA_PAGE,
889
+ block_stride=_MLA_PAGE,
890
+ ),
891
+ KVCacheTensor(
892
+ size=_MLA_PAGE * 4,
893
+ layers=["mamba"],
894
+ layer_stride=_MAMBA_PAGE,
895
+ block_stride=_MLA_PAGE,
896
+ ),
897
+ ],
898
+ kv_cache_groups=[
899
+ KVCacheGroupSpec(["mla"], _mla_spec()),
900
+ KVCacheGroupSpec(["mamba"], _MAMBA_SPEC),
901
+ ],
902
+ ),
903
+ "mla-mamba-hybrid",
904
+ ),
905
+ ],
906
+ ids=[
907
+ "hidden-state",
908
+ "mla-full-hybrid",
909
+ "mla-mamba-hybrid",
910
+ ],
911
+)
912
+def test_replicated_layout_excludes_unproven_cache_shapes(
913
+ kv_cache_config: KVCacheConfig,
914
+ case: str,
915
+):
916
+ assert not _replicated_layout(kv_cache_config), case
917
+
918
+
919
+@pytest.mark.parametrize(
920
+ ("kv_cache_config", "case"),
921
+ [
922
(
923
KVCacheConfig(
924
num_blocks=4,
@@ -631,16 +933,18 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
933
kv_cache_groups=[
934
KVCacheGroupSpec(
935
["layer"],
634
- HiddenStateCacheSpec(
936
+ SlidingWindowMLASpec(
937
block_size=16,
938
num_kv_heads=1,
939
+ max_tp_shards=1,
940
head_size=512,
941
dtype=torch.float32,
942
+ sliding_window=128,
943
),
944
)
945
],
946
),
643
- "hidden-state",
947
+ "sliding-window-mla",
948
),
949
(
950
# One group, two page sizes: layer1's run starts past layer0's region.
@@ -676,56 +980,6 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
980
),
981
"uniform-wrapper",
982
),
679
- (
680
- # Overlaid groups with different page sizes: a block is a window of
681
- # the largest group's packing.
682
- KVCacheConfig(
683
- num_blocks=4,
684
- kv_cache_tensors=[
685
- KVCacheTensor(
686
- size=_FULL_PAGE * 4,
687
- layers=["mla"],
688
- layer_stride=_MLA_PAGE,
689
- block_stride=_FULL_PAGE,
690
- ),
691
- KVCacheTensor(
692
- size=_FULL_PAGE * 4,
693
- layers=["full"],
694
- layer_stride=_FULL_PAGE,
695
- block_stride=_FULL_PAGE,
696
- ),
697
- ],
698
- kv_cache_groups=[
699
- KVCacheGroupSpec(["mla"], _mla_spec()),
700
- KVCacheGroupSpec(["full"], _full_attention_spec()),
701
- ],
702
- ),
703
- "mla-full-hybrid",
704
- ),
705
- (
706
- KVCacheConfig(
707
- num_blocks=4,
708
- kv_cache_tensors=[
709
- KVCacheTensor(
710
- size=_MLA_PAGE * 4,
711
- layers=["mla"],
712
- layer_stride=_MLA_PAGE,
713
- block_stride=_MLA_PAGE,
714
- ),
715
- KVCacheTensor(
716
- size=_MLA_PAGE * 4,
717
- layers=["mamba"],
718
- layer_stride=_MAMBA_PAGE,
719
- block_stride=_MLA_PAGE,
720
- ),
721
- ],
722
- kv_cache_groups=[
723
- KVCacheGroupSpec(["mla"], _mla_spec()),
724
- KVCacheGroupSpec(["mamba"], _MAMBA_SPEC),
725
- ],
726
- ),
727
- "mla-mamba-hybrid",
728
- ),
983
(
984
KVCacheConfig(
985
num_blocks=4,
@@ -748,31 +1002,54 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node(
1002
],
1003
ids=[
1004
"sliding-window-mla",
751
- "hidden-state",
1005
"uniform-wrapper",
753
- "mla-full-hybrid",
754
- "mla-mamba-hybrid",
1006
"multi-group-mla",
1007
],
1008
)
758
-def test_replicated_layout_excludes_unproven_cache_shapes(
1009
+def test_replicated_layout_supported_cache_shapes(
1010
kv_cache_config: KVCacheConfig,
1011
case: str,
1012
):
762
- assert not _replicated_layout(kv_cache_config), case
1013
+ assert _replicated_layout(kv_cache_config), case
1014
1015
765
-def test_replicated_layout_rejects_bare_mla_with_mixed_page_accounting():
1016
+@pytest.mark.parametrize(
1017
+ ("main_spec", "add_spec", "expected"),
1018
+ [
1019
+ (
1020
+ _mla_spec(head_size=512),
1021
+ _mla_spec(head_size=128, dtype=torch.uint8), # indexer.k_cache spec
1022
+ True,
1023
+ ),
1024
+ (
1025
+ _mla_spec(head_size=512),
1026
+ _full_attention_spec(),
1027
+ False,
1028
+ ),
1029
+ ],
1030
+ ids=["uniform-mla-indexer", "uniform-mla-fa"],
1031
+)
1032
+def test_replicated_layout_for_dsa_style_groups(
1033
+ main_spec: KVCacheSpec, add_spec: KVCacheSpec, expected: bool
1034
+):
1035
+ """KV Cache groups with packed-specs inside a UniformTypeKVCacheSpec."""
1036
num_blocks = 4
767
- main_spec = _mla_spec(head_size=512)
768
- indexer_spec = _mla_spec(head_size=128, dtype=torch.uint8)
1037
main_layers = [f"main_{i}" for i in range(61)]
770
- indexer_layers = [f"indexer_{i}" for i in range(61)]
771
- # A DSA-style group: the main pages and the smaller indexer pages are packed one
772
- # after the other, so a block holds more than 61 MLA pages.
1038
+ add_layers = [f"add_{i}" for i in range(61)]
1039
+
1040
main_bytes = main_spec.page_size_bytes * len(main_layers)
774
- indexer_bytes = indexer_spec.page_size_bytes * len(indexer_layers)
775
- size = (main_bytes + indexer_bytes) * num_blocks
1041
+ add_bytes = add_spec.page_size_bytes * len(add_layers)
1042
+ size = (main_bytes + add_bytes) * num_blocks
1043
+
1044
+ kv_cache_specs = {
1045
+ **{layer: main_spec for layer in main_layers},
1046
+ **{layer: add_spec for layer in add_layers},
1047
+ }
1048
+ layer_names = main_layers + add_layers
1049
+ group_spec: KVCacheSpec = UniformTypeKVCacheSpecs(
1050
+ block_size=16, kv_cache_specs=kv_cache_specs
1051
+ )
1052
+
1053
kv_cache_config = KVCacheConfig(
1054
num_blocks=num_blocks,
1055
kv_cache_tensors=[
@@ -784,16 +1061,16 @@ def test_replicated_layout_rejects_bare_mla_with_mixed_page_accounting():
1061
),
1062
KVCacheTensor(
1063
size=size,
787
- layers=indexer_layers,
788
- layer_stride=indexer_spec.page_size_bytes * num_blocks,
789
- block_stride=indexer_spec.page_size_bytes,
1064
+ layers=add_layers,
1065
+ layer_stride=add_spec.page_size_bytes * num_blocks,
1066
+ block_stride=add_spec.page_size_bytes,
1067
offset=main_bytes * num_blocks,
1068
),
1069
],
793
- kv_cache_groups=[KVCacheGroupSpec(main_layers + indexer_layers, main_spec)],
1070
+ kv_cache_groups=[KVCacheGroupSpec(layer_names, group_spec)],
1071
)
1072
796
- assert not _replicated_layout(kv_cache_config)
1073
+ assert _replicated_layout(kv_cache_config) == expected
1074
1075
1076
@pytest.mark.parametrize(
@@ -841,6 +1118,7 @@ _SWA_SPEC = SlidingWindowSpec(
1118
_SWA_MLA_SPEC = SlidingWindowMLASpec(
1119
block_size=16,
1120
num_kv_heads=1,
1121
+ max_tp_shards=1,
1122
head_size=576,
1123
dtype=torch.float32,
1124
sliding_window=128,
@@ -904,6 +1182,52 @@ def test_canonical_layout_gate(kv_cache_groups, certified):
1182
assert _parallelism_agnostic(kv_cache_groups, canonical=True) is certified
1183
1184
1185
+@pytest.mark.parametrize(
1186
+ ("spec"),
1187
+ [_SWA_MLA_SPEC, _mla_spec(tokens_per_state=2)],
1188
+ ids=["swa", "compressed-mla"],
1189
+)
1190
+def test_uncertifiable_canonical_layout_resets_replication(spec: KVCacheSpec):
1191
+ def _get_offloading_config(canonical: bool):
1192
+ config = _make_vllm_config(
1193
+ tensor_parallel_size=2,
1194
+ extra_config={"canonical_layout": True} if canonical else None,
1195
+ )
1196
+ config.model_config.use_mla = True
1197
+
1198
+ PAGE = spec.page_size_bytes
1199
+ kv_config = KVCacheConfig(
1200
+ num_blocks=4,
1201
+ kv_cache_tensors=[
1202
+ KVCacheTensor(
1203
+ size=PAGE * 4,
1204
+ layers=["layer"],
1205
+ layer_stride=PAGE * 4,
1206
+ block_stride=PAGE,
1207
+ )
1208
+ ],
1209
+ kv_cache_groups=[
1210
+ KVCacheGroupSpec(["layer"], spec),
1211
+ ],
1212
+ )
1213
+
1214
+ kv_config.kv_tp_replicas = kv_cache_groups_tp_replicas(
1215
+ kv_config.kv_cache_groups,
1216
+ config.parallel_config.tensor_parallel_size,
1217
+ config.parallel_config.decode_context_parallel_size,
1218
+ )
1219
+
1220
+ return build_offloading_config(config, kv_config)
1221
+
1222
+ offload_config = _get_offloading_config(canonical=False)
1223
+ assert offload_config.replicated_layout
1224
+ assert not offload_config.parallel.is_parallelism_agnostic
1225
+
1226
+ canon_offload_config = _get_offloading_config(canonical=True)
1227
+ assert not canon_offload_config.replicated_layout
1228
+ assert not canon_offload_config.parallel.is_parallelism_agnostic
1229
+
1230
+
1231
def test_canonical_layout_certifies_v2_model_runner():
1232
"""Canonical bytes are certified per layer against live tensor strides at
1233
registration, so the static gate must not depend on the model-runner
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+5
-9
@@ -157,15 +157,10 @@ def build_offloading_config(
157
)
158
replicated_layout = (
159
vllm_config.model_config.use_mla
160
- # Exact type: fail closed on wrappers and sliding-window variants.
161
- and type(single_group_spec) is MLAAttentionSpec
162
- # Page accounting: one MLA page per layer, no packed/mixed rows.
160
+ and parallel_config.tensor_parallel_size > 1
161
+ and kv_cache_config.kv_tp_replicas == parallel_config.tensor_parallel_size
162
and worker_kv_bytes_per_block > 0
164
- and worker_kv_bytes_per_block
165
- == single_group_spec.page_size_bytes
166
- * len(kv_cache_config.kv_cache_groups[0].layer_names)
163
# Safe MVP boundary: TP-only, no other parallel axes.
168
- and parallel_config.tensor_parallel_size > 1
164
and parallel_config.pipeline_parallel_size == 1
165
and parallel_config.prefill_context_parallel_size == 1
166
and parallel_config.decode_context_parallel_size == 1
@@ -233,9 +228,10 @@ def build_offloading_config(
228
and parallel_config.world_size == tp_size
229
)
230
236
- if canonical_layout and is_parallelism_agnostic:
231
+ if canonical_layout:
232
replicated_layout = (
238
- all(
233
+ is_parallelism_agnostic
234
+ and all(
235
type(spec) is MLAAttentionSpec
236
for _, group in selected_groups
237
for spec in iter_layer_specs(group.kv_cache_spec)
vllm/model_executor/layers/attention/mla_attention.py
+1
@@ -1350,6 +1350,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):
1350
common_kwargs = dict(
1351
block_size=vllm_config.cache_config.block_size,
1352
num_kv_heads=1,
1353
+ max_tp_shards=1,
1354
head_size=self.head_size,
1355
dtype=kv_cache_dtype,
1356
cache_dtype_str=self.kv_cache_dtype,
vllm/model_executor/models/deepseek_v2.py
+1
@@ -658,6 +658,7 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):
658
return MLAAttentionSpec(
659
block_size=self.cache_config.block_size,
660
num_kv_heads=1,
661
+ max_tp_shards=1,
662
head_size=self.head_dim,
663
dtype=self.dtype,
664
cache_role=SparseCacheRole.INDEXER,
vllm/models/deepseek_v4/attention.py
+2
@@ -865,6 +865,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
865
return MLAAttentionSpec(
866
block_size=vllm_config.cache_config.block_size,
867
num_kv_heads=1,
868
+ max_tp_shards=1,
869
head_size=self.head_dim,
870
dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
871
tokens_per_state=self.compress_ratio,
@@ -910,6 +911,7 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
911
return MLAAttentionSpec(
912
block_size=self.cache_config.block_size,
913
num_kv_heads=1,
914
+ max_tp_shards=1,
915
head_size=self.head_dim,
916
dtype=self.dtype,
917
tokens_per_state=self.compress_ratio,
vllm/models/deepseek_v4/compressor.py
+1
@@ -184,6 +184,7 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
184
return SlidingWindowMLASpec( # only has one vector instead of K + V
185
block_size=self.block_size,
186
num_kv_heads=1,
187
+ max_tp_shards=1,
188
head_size=self.state_dim,
189
dtype=self.dtype,
190
sliding_window=self.sliding_window,
vllm/v1/attention/backends/mla/sparse_swa.py
+1
@@ -124,6 +124,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
124
bounded_replay=self.bounded_replay,
125
block_size=self.block_size,
126
num_kv_heads=1,
127
+ max_tp_shards=1,
128
head_size=self.head_dim,
129
dtype=self.dtype,
130
sliding_window=self.window_size,
vllm/v1/core/kv_cache_utils.py
+28
-1
@@ -9,7 +9,7 @@ import os
9
from collections import defaultdict
10
from collections.abc import Callable, Iterable, Iterator, Sequence
11
from dataclasses import dataclass, field, replace
12
-from functools import partial
12
+from functools import partial, reduce
13
from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, cast, overload
14
15
from vllm import envs
@@ -2409,6 +2409,26 @@ def get_kv_cache_groups(
2409
return groups
2410
2411
2412
+def _layer_tp_replicas(spec: KVCacheSpec, tp_size: int, dcp_size: int) -> int:
2413
+ if not isinstance(spec, AttentionSpec) or spec.max_tp_shards is None:
2414
+ return 1
2415
+ if spec.dcp_sharded and dcp_size > 1:
2416
+ return 1
2417
+ return max(1, tp_size // spec.max_tp_shards)
2418
+
2419
+
2420
+def kv_cache_groups_tp_replicas(
2421
+ groups: list[KVCacheGroupSpec], tp_size: int, dcp_size: int = 1
2422
+) -> int:
2423
+ """Consecutive TP ranks holding identical KV for every layer."""
2424
+ specs = [spec for g in groups for spec in iter_layer_specs(g.kv_cache_spec)]
2425
+ if not specs:
2426
+ return 1
2427
+ return reduce(
2428
+ math.gcd, (_layer_tp_replicas(s, tp_size, dcp_size) for s in specs), tp_size
2429
+ )
2430
+
2431
+
2432
def generate_scheduler_kv_cache_config(
2433
kv_cache_configs: list[KVCacheConfig],
2434
) -> KVCacheConfig:
@@ -2823,6 +2843,13 @@ def get_kv_cache_configs(
2843
vllm_config, groups, min_num_blocks * _pool_bytes_per_block(groups)
2844
)
2845
2846
+ for kv_cache_config in kv_cache_configs:
2847
+ kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas(
2848
+ kv_cache_config.kv_cache_groups,
2849
+ vllm_config.parallel_config.tensor_parallel_size,
2850
+ vllm_config.parallel_config.decode_context_parallel_size,
2851
+ )
2852
+
2853
return kv_cache_configs
2854
2855
vllm/v1/kv_cache_interface.py
+12
@@ -484,6 +484,9 @@ class HiSparseResidentSpec(KVCacheSpec):
484
@dataclass(frozen=True, kw_only=True)
485
class AttentionSpec(KVCacheSpec):
486
dcp_sharded: bool = True
487
+ max_tp_shards: int | None = None
488
+ """Distinct shards this cache splits into across TP; TP ranks beyond this
489
+ hold replicas. None: unknown, treated as not replicated."""
490
num_kv_heads: int
491
head_size: int
492
dtype: torch.dtype
@@ -618,6 +621,7 @@ class FullAttentionSpec(AttentionSpec):
621
dtype=specs[0].dtype,
622
kv_quant_mode=specs[0].kv_quant_mode,
623
dcp_sharded=specs[0].dcp_sharded,
624
+ max_tp_shards=specs[0].max_tp_shards,
625
page_size_padded=specs[0].page_size_padded,
626
num_head_slots=specs[0].num_head_slots,
627
state_content_bytes=specs[0].state_content_bytes,
@@ -705,6 +709,7 @@ class MLAAttentionSpec(FullAttentionSpec):
709
dtype=specs[0].dtype,
710
kv_quant_mode=specs[0].kv_quant_mode,
711
dcp_sharded=specs[0].dcp_sharded,
712
+ max_tp_shards=specs[0].max_tp_shards,
713
page_size_padded=specs[0].page_size_padded,
714
num_head_slots=specs[0].num_head_slots,
715
state_content_bytes=specs[0].state_content_bytes,
@@ -768,6 +773,7 @@ class RSWASpec(FullAttentionSpec):
773
dtype=base.dtype,
774
kv_quant_mode=base.kv_quant_mode,
775
dcp_sharded=base.dcp_sharded,
776
+ max_tp_shards=base.max_tp_shards,
777
page_size_padded=base.page_size_padded,
778
num_head_slots=base.num_head_slots,
779
state_content_bytes=base.state_content_bytes,
@@ -960,6 +966,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
966
bounded_replay_set = set(spec.bounded_replay for spec in specs)
967
block_stride_alignment_set = {spec.block_stride_alignment for spec in specs}
968
assert len({spec.dcp_sharded for spec in specs}) == 1
969
+ assert len({spec.max_tp_shards for spec in specs}) == 1
970
assert (
971
len(cache_dtype_str_set) == 1
972
and len(tokens_per_state_set) == 1
@@ -980,6 +987,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec):
987
head_size=specs[0].head_size,
988
dtype=specs[0].dtype,
989
dcp_sharded=specs[0].dcp_sharded,
990
+ max_tp_shards=specs[0].max_tp_shards,
991
page_size_padded=specs[0].page_size_padded,
992
num_head_slots=specs[0].num_head_slots,
993
state_content_bytes=specs[0].state_content_bytes,
@@ -1197,6 +1205,7 @@ class SinkFullAttentionSpec(FullAttentionSpec):
1205
dtype=specs[0].dtype,
1206
kv_quant_mode=specs[0].kv_quant_mode,
1207
dcp_sharded=specs[0].dcp_sharded,
1208
+ max_tp_shards=specs[0].max_tp_shards,
1209
page_size_padded=specs[0].page_size_padded,
1210
num_head_slots=specs[0].num_head_slots,
1211
state_content_bytes=specs[0].state_content_bytes,
@@ -1483,6 +1492,9 @@ class KVCacheConfig:
1492
hisparse_shared_host_pool: bool = False
1493
"""Whether local TP ranks share one physical HiSparse host pool."""
1494
1495
+ kv_tp_replicas: int = 1
1496
+ """Consecutive TP ranks holding identical KV for every layer (1: none)."""
1497
+
1498
@cached_property
1499
def transfer_group_ids(self) -> tuple[int, ...]:
1500
"""IDs of cache groups that participate in external KV transfer."""