From 5463fe4962785cdc3383477bf3af6533e7647dfd Mon Sep 17 00:00:00 2001 From: Varun Sundar Rabindranath Date: Wed, 30 Sep 2026 09:11:03 -0400 Subject: [PATCH] [KV-Offloading][TP] : Expand replicated_layout detection to multi-group MLA (#57652) --- tests/v1/core/test_kv_cache_utils.py | 97 ++++ .../unit/offloading_connector/test_config.py | 484 +++++++++++++++--- .../kv_connector/v1/offloading/config.py | 16 +- .../layers/attention/mla_attention.py | 1 + vllm/model_executor/models/deepseek_v2.py | 1 + vllm/models/deepseek_v4/attention.py | 2 + vllm/models/deepseek_v4/compressor.py | 1 + vllm/v1/attention/backends/mla/sparse_swa.py | 1 + vllm/v1/core/kv_cache_utils.py | 29 +- vllm/v1/kv_cache_interface.py | 12 + 10 files changed, 553 insertions(+), 91 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index fb82d28c36..0c56e89d59 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -52,6 +52,7 @@ from vllm.v1.core.kv_cache_utils import ( hash_block_tokens, init_none_hash, is_kv_cache_spec_uniform, + kv_cache_groups_tp_replicas, make_block_hash_with_group_id, tensor_data, to_event_extra_keys, @@ -3274,6 +3275,7 @@ def new_mla_spec(cache_dtype_str=None, block_size: int = 16): return MLAAttentionSpec( block_size=block_size, num_kv_heads=1, + max_tp_shards=1, head_size=576, dtype=torch.float32, cache_dtype_str=cache_dtype_str, @@ -3284,6 +3286,7 @@ def new_swa_mla_spec(head_size=576, sliding_window=128, model_version=None): return SlidingWindowMLASpec( block_size=16, num_kv_heads=1, + max_tp_shards=1, head_size=head_size, dtype=torch.float32, sliding_window=sliding_window, @@ -4637,3 +4640,97 @@ def test_trailing_layer_fallback_requires_exact_partition(): _annotate_eagle_groups(config, specs, trimmed, use_trailing_layer_fallback=True) assert not any(g.is_eagle_group for g in trimmed) + + +_GQA_SPEC = FullAttentionSpec( + block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32 +) +_GQA_SWA_SPEC = SlidingWindowSpec( + block_size=16, num_kv_heads=8, head_size=64, dtype=torch.float32, sliding_window=128 +) + + +@pytest.mark.parametrize( + "specs,expected", + [ + pytest.param({"l.0": new_mla_spec(), "l.1": new_mla_spec()}, 4, id="mla"), + pytest.param( + { + "l.0": new_mla_spec(), + "l.1": new_mla_spec(), + "s.0": new_swa_mla_spec(), + "s.1": new_swa_mla_spec(), + }, + 4, + id="mla-and-swa-mla", + ), + pytest.param({"l.0": _GQA_SPEC, "l.1": _GQA_SPEC}, 1, id="gqa"), + pytest.param( + {"l.0": new_mla_spec(), "l.1": _GQA_SPEC}, + 1, + id="mla-and-gqa-uniform-group", + ), + pytest.param( + { + "l.0": new_mla_spec(), + "l.1": new_mla_spec(), + "s.0": _GQA_SWA_SPEC, + "s.1": _GQA_SWA_SPEC, + }, + 1, + id="mla-and-gqa-swa", + ), + ], +) +def test_kv_tp_replicas(monkeypatch, specs, expected): + """Resolved per layer before scheduler flattening hides mixed groups.""" + from vllm.config import ParallelConfig + from vllm.platforms import current_platform + + monkeypatch.setattr(current_platform, "device_count", lambda: 4) + parallel_config = ParallelConfig(tensor_parallel_size=4) + vllm_config = VllmConfig( + model_config=ModelConfig(max_model_len=16), parallel_config=parallel_config + ) + vllm_config.cache_config.kv_cache_layout = "LBNHC" + mem = sum(s.page_size_bytes for s in specs.values()) * 10 + configs = get_kv_cache_configs(vllm_config, [specs], [mem]) + assert configs[0].kv_tp_replicas == expected + scheduler = generate_scheduler_kv_cache_config(configs) + assert scheduler.kv_tp_replicas == expected + + +@pytest.mark.parametrize( + "specs,tp_size,dcp_size,expected", + [ + pytest.param([new_mla_spec()], 8, 1, 8, id="mla"), + pytest.param([new_mla_spec()], 8, 2, 1, id="mla-dcp"), + pytest.param( + [replace(new_mla_spec(), max_tp_shards=None)], 8, 1, 1, id="mla-unset" + ), + pytest.param([replace(_GQA_SPEC, max_tp_shards=2)], 8, 1, 4, id="gqa-partial"), + pytest.param([replace(_GQA_SPEC, max_tp_shards=8)], 4, 1, 1, id="gqa-sharded"), + pytest.param( + [new_mla_spec(), replace(_GQA_SWA_SPEC, max_tp_shards=4)], + 8, + 1, + 2, + id="mla-and-gqa-partial", + ), + pytest.param( + [ + HiddenStateCacheSpec( + block_size=16, num_kv_heads=1, head_size=64, dtype=torch.float32 + ) + ], + 8, + 1, + 1, + id="hidden-state", + ), + ], +) +def test_kv_cache_groups_tp_replicas(specs, tp_size, dcp_size, expected): + """Replicas are the gcd of each layer's tp_size // max_tp_shards.""" + groups = [KVCacheGroupSpec([f"l.{i}"], spec) for i, spec in enumerate(specs)] + assert kv_cache_groups_tp_replicas(groups, tp_size, dcp_size) == expected diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index fe0d258462..237b376004 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -17,7 +17,10 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( SchedulerOffloadConfig, ) from vllm.platforms import current_platform -from vllm.v1.core.kv_cache_utils import generate_scheduler_kv_cache_config +from vllm.v1.core.kv_cache_utils import ( + generate_scheduler_kv_cache_config, + kv_cache_groups_tp_replicas, +) from vllm.v1.kv_cache_interface import ( CircularBufferSpec, FullAttentionSpec, @@ -164,6 +167,7 @@ def _mla_spec( return MLAAttentionSpec( block_size=block_size, num_kv_heads=1, + max_tp_shards=1, head_size=head_size, dtype=dtype, tokens_per_state=tokens_per_state, @@ -321,7 +325,184 @@ def _replicated_layout( config.parallel_config.nnodes = nnodes if world_size is not None: config.parallel_config.world_size = world_size - return build_offloading_config(config, kv_cache_config).replicated_layout + kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas( + kv_cache_config.kv_cache_groups, + config.parallel_config.tensor_parallel_size, + config.parallel_config.decode_context_parallel_size, + ) + + # Scheduler and Worker offload config + worker_offload_config = build_offloading_config(config, kv_cache_config) + scheduler_offload_config = build_offloading_config( + config, generate_scheduler_kv_cache_config([kv_cache_config]) + ) + # Verify that the Worker and Scheduler agree on replication. + # and worker-bytes-per-block + assert ( + worker_offload_config.replicated_layout + == scheduler_offload_config.replicated_layout + ) + assert ( + worker_offload_config.worker_kv_bytes_per_block + == scheduler_offload_config.worker_kv_bytes_per_block + ) + + return worker_offload_config.replicated_layout + + +_SWA_MLA_PAGE = SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + max_tp_shards=1, + head_size=512, + dtype=torch.float32, + sliding_window=128, +).page_size_bytes + + +def _make_swa_mla_kv_cache_config( + num_blocks: int = 4, +) -> KVCacheConfig: + """Single SlidingWindowMLASpec group — DSV4 sliding-window layers.""" + spec = SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + max_tp_shards=1, + head_size=512, + dtype=torch.float32, + sliding_window=128, + ) + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=_SWA_MLA_PAGE * num_blocks, + layers=["layer"], + layer_stride=_SWA_MLA_PAGE * num_blocks, + block_stride=_SWA_MLA_PAGE, + ) + ], + kv_cache_groups=[KVCacheGroupSpec(["layer"], spec)], + ) + + +def _make_dsv4_flash_kv_cache_config(num_blocks: int = 4) -> KVCacheConfig: + """Minimal fake of DeepSeek-V4-Flash's five UniformTypeKVCacheSpecs groups. + + The real model has the following groups (layer counts from runtime logs); + 2 per group is enough to exercise the structural properties that govern + replicated_layout: + group[0]: 22 SlidingWindowMLASpec layers, sliding_window=128 + group[1]: 21 SlidingWindowMLASpec layers, sliding_window=128 + group[2]: 62 MLAAttentionSpec layers (MLA + indexer.k_cache) + group[3]: 42 SlidingWindowMLASpec layers, sliding_window=8 + group[4]: 20 SlidingWindowMLASpec layers, sliding_window=128 + All inner specs have num_kv_heads=1 → every group is TP-replicated. + """ + block_size = 16 + dtype = torch.float32 + head_size = 512 + + def _swa_group(tag: str, sw: int) -> KVCacheGroupSpec: + names = [f"{tag}_0", f"{tag}_1"] + spec = SlidingWindowMLASpec( + block_size=block_size, + num_kv_heads=1, + max_tp_shards=1, + head_size=head_size, + dtype=dtype, + sliding_window=sw, + ) + return KVCacheGroupSpec( + names, + UniformTypeKVCacheSpecs( + block_size=block_size, + kv_cache_specs={name: spec for name in names}, + ), + ) + + def _mla_group(tag: str) -> KVCacheGroupSpec: + names = [f"{tag}_0", f"{tag}_1"] + spec = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + max_tp_shards=1, + head_size=head_size, + dtype=dtype, + ) + return KVCacheGroupSpec( + names, + UniformTypeKVCacheSpecs( + block_size=block_size, + kv_cache_specs={name: spec for name in names}, + ), + ) + + groups = [ + _swa_group("swa_128a", 128), # group[0]: SWA-MLA, sw=128 + _swa_group("swa_128b", 128), # group[1]: SWA-MLA, sw=128 + _mla_group("mla"), # group[2]: MLA + indexer layers + _swa_group("swa_8", 8), # group[3]: SWA-MLA, sw=8 + _swa_group("swa_128c", 128), # group[4]: SWA-MLA, sw=128 + ] + # All specs share the same page size (num_kv_heads=1, head_size=512, fp32). + page = MLAAttentionSpec( + block_size=block_size, num_kv_heads=1, head_size=head_size, dtype=dtype + ).page_size_bytes + all_layers = [name for g in groups for name in g.layer_names] + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=page * len(all_layers) * num_blocks, + layers=all_layers, + layer_stride=page * num_blocks, + block_stride=page, + ) + ], + kv_cache_groups=groups, + ) + + +def _make_dsv3_2_kv_cache_config(num_blocks: int = 4) -> KVCacheConfig: + """Minimal fake of DeepSeek-V3.2-Exp's single UniformTypeKVCacheSpecs group. + + The real model has 1 group (layer counts from runtime logs): + group[0]: 122 MLAAttentionSpec layers (self_attn.attn + indexer.k_cache) + All inner specs have num_kv_heads=1 → the group is TP-replicated. + + GLM-5.3 has the same topology (one UniformType group of MLAAttentionSpec + layers, num_kv_heads=1) with 99 layers instead of 122, so this config + covers both models. + """ + block_size = 16 + names = [f"mla_{i}" for i in range(4)] # 4 layers; real model has 122 + spec = MLAAttentionSpec( + block_size=block_size, + num_kv_heads=1, + max_tp_shards=1, + head_size=512, + dtype=torch.float32, + ) + group = KVCacheGroupSpec( + names, + UniformTypeKVCacheSpecs( + block_size=block_size, + kv_cache_specs={name: spec for name in names}, + ), + ) + return KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=spec.page_size_bytes * len(names) * num_blocks, + layers=names, + layer_stride=spec.page_size_bytes * num_blocks, + block_stride=spec.page_size_bytes, + ) + ], + kv_cache_groups=[group], + ) @pytest.mark.parametrize("packed", [False, True]) @@ -588,35 +769,64 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node( ) +@pytest.mark.parametrize("world_size", [2, 4]) +def test_replicated_layout_enabled_for_swa_mla(world_size: int): + """SlidingWindowMLASpec groups are also replicated and qualify.""" + assert _replicated_layout( + _make_swa_mla_kv_cache_config(), + tensor_parallel_size=world_size, + ) + + +@pytest.mark.parametrize("world_size", [2, 4, 8]) +def test_replicated_layout_dsv4_flash(world_size: int): + """DeepSeek-V4-Flash: all 5 UniformType groups are MLA/SWA-MLA with + num_kv_heads=1, so replicated_layout must be True under TP/mp single-node.""" + assert _replicated_layout( + _make_dsv4_flash_kv_cache_config(), + tensor_parallel_size=world_size, + ) + + +@pytest.mark.parametrize("world_size", [2, 4, 8]) +def test_replicated_layout_dsv3_2(world_size: int): + """DeepSeek-V3.2 / GLM-5.3: single UniformType group of MLAAttentionSpec + layers (122 and 99 respectively), num_kv_heads=1 → replicated_layout must + be True under TP/mp single-node.""" + assert _replicated_layout( + _make_dsv3_2_kv_cache_config(), + tensor_parallel_size=world_size, + ) + + +def test_replicated_layout_enabled_for_uniform_wrapper_mla(): + """UniformTypeKVCacheSpecs wrapping MLA layers of different sizes qualifies: + each layer is replicated and the page accounting sums correctly.""" + mla_layers = ["layer0", "layer1"] + spec = UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={"layer0": _mla_spec(), "layer1": _mla_spec(head_size=256)}, + ) + total_page = spec.page_size_bytes + num_blocks = 4 + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=total_page * num_blocks, + layers=mla_layers, + layer_stride=total_page * num_blocks, + block_stride=total_page, + ) + ], + kv_cache_groups=[KVCacheGroupSpec(mla_layers, spec)], + ) + assert _replicated_layout(kv_cache_config) + + @pytest.mark.parametrize( ("kv_cache_config", "case"), [ - ( - KVCacheConfig( - num_blocks=4, - kv_cache_tensors=[ - KVCacheTensor( - size=_MLA_PAGE * 4, - layers=["layer"], - layer_stride=_MLA_PAGE * 4, - block_stride=_MLA_PAGE, - ) - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer"], - SlidingWindowMLASpec( - block_size=16, - num_kv_heads=1, - head_size=512, - dtype=torch.float32, - sliding_window=128, - ), - ) - ], - ), - "sliding-window-mla", - ), ( KVCacheConfig( num_blocks=4, @@ -642,40 +852,6 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node( ), "hidden-state", ), - ( - # One group, two page sizes: layer1's run starts past layer0's region. - KVCacheConfig( - num_blocks=4, - kv_cache_tensors=[ - KVCacheTensor( - size=(_MLA_PAGE + _HALF_MLA_PAGE) * 4, - layers=["layer0"], - layer_stride=_MLA_PAGE * 4, - block_stride=_MLA_PAGE, - ), - KVCacheTensor( - size=(_MLA_PAGE + _HALF_MLA_PAGE) * 4, - layers=["layer1"], - layer_stride=_HALF_MLA_PAGE * 4, - block_stride=_HALF_MLA_PAGE, - offset=_MLA_PAGE * 4, - ), - ], - kv_cache_groups=[ - KVCacheGroupSpec( - ["layer0", "layer1"], - UniformTypeKVCacheSpecs( - block_size=16, - kv_cache_specs={ - "layer0": _mla_spec(), - "layer1": _mla_spec(head_size=256), - }, - ), - ) - ], - ), - "uniform-wrapper", - ), ( # Overlaid groups with different page sizes: a block is a window of # the largest group's packing. @@ -726,6 +902,84 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node( ), "mla-mamba-hybrid", ), + ], + ids=[ + "hidden-state", + "mla-full-hybrid", + "mla-mamba-hybrid", + ], +) +def test_replicated_layout_excludes_unproven_cache_shapes( + kv_cache_config: KVCacheConfig, + case: str, +): + assert not _replicated_layout(kv_cache_config), case + + +@pytest.mark.parametrize( + ("kv_cache_config", "case"), + [ + ( + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=_MLA_PAGE * 4, + layers=["layer"], + layer_stride=_MLA_PAGE * 4, + block_stride=_MLA_PAGE, + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer"], + SlidingWindowMLASpec( + block_size=16, + num_kv_heads=1, + max_tp_shards=1, + head_size=512, + dtype=torch.float32, + sliding_window=128, + ), + ) + ], + ), + "sliding-window-mla", + ), + ( + # One group, two page sizes: layer1's run starts past layer0's region. + KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=(_MLA_PAGE + _HALF_MLA_PAGE) * 4, + layers=["layer0"], + layer_stride=_MLA_PAGE * 4, + block_stride=_MLA_PAGE, + ), + KVCacheTensor( + size=(_MLA_PAGE + _HALF_MLA_PAGE) * 4, + layers=["layer1"], + layer_stride=_HALF_MLA_PAGE * 4, + block_stride=_HALF_MLA_PAGE, + offset=_MLA_PAGE * 4, + ), + ], + kv_cache_groups=[ + KVCacheGroupSpec( + ["layer0", "layer1"], + UniformTypeKVCacheSpecs( + block_size=16, + kv_cache_specs={ + "layer0": _mla_spec(), + "layer1": _mla_spec(head_size=256), + }, + ), + ) + ], + ), + "uniform-wrapper", + ), ( KVCacheConfig( num_blocks=4, @@ -748,31 +1002,54 @@ def test_replicated_layout_enabled_for_pure_mla_tp_mp_single_node( ], ids=[ "sliding-window-mla", - "hidden-state", "uniform-wrapper", - "mla-full-hybrid", - "mla-mamba-hybrid", "multi-group-mla", ], ) -def test_replicated_layout_excludes_unproven_cache_shapes( +def test_replicated_layout_supported_cache_shapes( kv_cache_config: KVCacheConfig, case: str, ): - assert not _replicated_layout(kv_cache_config), case + assert _replicated_layout(kv_cache_config), case -def test_replicated_layout_rejects_bare_mla_with_mixed_page_accounting(): +@pytest.mark.parametrize( + ("main_spec", "add_spec", "expected"), + [ + ( + _mla_spec(head_size=512), + _mla_spec(head_size=128, dtype=torch.uint8), # indexer.k_cache spec + True, + ), + ( + _mla_spec(head_size=512), + _full_attention_spec(), + False, + ), + ], + ids=["uniform-mla-indexer", "uniform-mla-fa"], +) +def test_replicated_layout_for_dsa_style_groups( + main_spec: KVCacheSpec, add_spec: KVCacheSpec, expected: bool +): + """KV Cache groups with packed-specs inside a UniformTypeKVCacheSpec.""" num_blocks = 4 - main_spec = _mla_spec(head_size=512) - indexer_spec = _mla_spec(head_size=128, dtype=torch.uint8) main_layers = [f"main_{i}" for i in range(61)] - indexer_layers = [f"indexer_{i}" for i in range(61)] - # A DSA-style group: the main pages and the smaller indexer pages are packed one - # after the other, so a block holds more than 61 MLA pages. + add_layers = [f"add_{i}" for i in range(61)] + main_bytes = main_spec.page_size_bytes * len(main_layers) - indexer_bytes = indexer_spec.page_size_bytes * len(indexer_layers) - size = (main_bytes + indexer_bytes) * num_blocks + add_bytes = add_spec.page_size_bytes * len(add_layers) + size = (main_bytes + add_bytes) * num_blocks + + kv_cache_specs = { + **{layer: main_spec for layer in main_layers}, + **{layer: add_spec for layer in add_layers}, + } + layer_names = main_layers + add_layers + group_spec: KVCacheSpec = UniformTypeKVCacheSpecs( + block_size=16, kv_cache_specs=kv_cache_specs + ) + kv_cache_config = KVCacheConfig( num_blocks=num_blocks, kv_cache_tensors=[ @@ -784,16 +1061,16 @@ def test_replicated_layout_rejects_bare_mla_with_mixed_page_accounting(): ), KVCacheTensor( size=size, - layers=indexer_layers, - layer_stride=indexer_spec.page_size_bytes * num_blocks, - block_stride=indexer_spec.page_size_bytes, + layers=add_layers, + layer_stride=add_spec.page_size_bytes * num_blocks, + block_stride=add_spec.page_size_bytes, offset=main_bytes * num_blocks, ), ], - kv_cache_groups=[KVCacheGroupSpec(main_layers + indexer_layers, main_spec)], + kv_cache_groups=[KVCacheGroupSpec(layer_names, group_spec)], ) - assert not _replicated_layout(kv_cache_config) + assert _replicated_layout(kv_cache_config) == expected @pytest.mark.parametrize( @@ -841,6 +1118,7 @@ _SWA_SPEC = SlidingWindowSpec( _SWA_MLA_SPEC = SlidingWindowMLASpec( block_size=16, num_kv_heads=1, + max_tp_shards=1, head_size=576, dtype=torch.float32, sliding_window=128, @@ -904,6 +1182,52 @@ def test_canonical_layout_gate(kv_cache_groups, certified): assert _parallelism_agnostic(kv_cache_groups, canonical=True) is certified +@pytest.mark.parametrize( + ("spec"), + [_SWA_MLA_SPEC, _mla_spec(tokens_per_state=2)], + ids=["swa", "compressed-mla"], +) +def test_uncertifiable_canonical_layout_resets_replication(spec: KVCacheSpec): + def _get_offloading_config(canonical: bool): + config = _make_vllm_config( + tensor_parallel_size=2, + extra_config={"canonical_layout": True} if canonical else None, + ) + config.model_config.use_mla = True + + PAGE = spec.page_size_bytes + kv_config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=PAGE * 4, + layers=["layer"], + layer_stride=PAGE * 4, + block_stride=PAGE, + ) + ], + kv_cache_groups=[ + KVCacheGroupSpec(["layer"], spec), + ], + ) + + kv_config.kv_tp_replicas = kv_cache_groups_tp_replicas( + kv_config.kv_cache_groups, + config.parallel_config.tensor_parallel_size, + config.parallel_config.decode_context_parallel_size, + ) + + return build_offloading_config(config, kv_config) + + offload_config = _get_offloading_config(canonical=False) + assert offload_config.replicated_layout + assert not offload_config.parallel.is_parallelism_agnostic + + canon_offload_config = _get_offloading_config(canonical=True) + assert not canon_offload_config.replicated_layout + assert not canon_offload_config.parallel.is_parallelism_agnostic + + def test_canonical_layout_certifies_v2_model_runner(): """Canonical bytes are certified per layer against live tensor strides at registration, so the static gate must not depend on the model-runner diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index e20c1189af..186590b956 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py @@ -157,15 +157,10 @@ def build_offloading_config( ) replicated_layout = ( vllm_config.model_config.use_mla - # Exact type: fail closed on wrappers and sliding-window variants. - and type(single_group_spec) is MLAAttentionSpec - # Page accounting: one MLA page per layer, no packed/mixed rows. - and worker_kv_bytes_per_block > 0 - and worker_kv_bytes_per_block - == single_group_spec.page_size_bytes - * len(kv_cache_config.kv_cache_groups[0].layer_names) - # Safe MVP boundary: TP-only, no other parallel axes. and parallel_config.tensor_parallel_size > 1 + and kv_cache_config.kv_tp_replicas == parallel_config.tensor_parallel_size + and worker_kv_bytes_per_block > 0 + # Safe MVP boundary: TP-only, no other parallel axes. and parallel_config.pipeline_parallel_size == 1 and parallel_config.prefill_context_parallel_size == 1 and parallel_config.decode_context_parallel_size == 1 @@ -233,9 +228,10 @@ def build_offloading_config( and parallel_config.world_size == tp_size ) - if canonical_layout and is_parallelism_agnostic: + if canonical_layout: replicated_layout = ( - all( + is_parallelism_agnostic + and all( type(spec) is MLAAttentionSpec for _, group in selected_groups for spec in iter_layer_specs(group.kv_cache_spec) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 2ae111c1c3..e016e5daff 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -1350,6 +1350,7 @@ class MLAAttention(nn.Module, AttentionLayerBase): common_kwargs = dict( block_size=vllm_config.cache_config.block_size, num_kv_heads=1, + max_tp_shards=1, head_size=self.head_size, dtype=kv_cache_dtype, cache_dtype_str=self.kv_cache_dtype, diff --git a/vllm/model_executor/models/deepseek_v2.py b/vllm/model_executor/models/deepseek_v2.py index 01010a13fe..70d87529f6 100644 --- a/vllm/model_executor/models/deepseek_v2.py +++ b/vllm/model_executor/models/deepseek_v2.py @@ -658,6 +658,7 @@ class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase): return MLAAttentionSpec( block_size=self.cache_config.block_size, num_kv_heads=1, + max_tp_shards=1, head_size=self.head_dim, dtype=self.dtype, cache_role=SparseCacheRole.INDEXER, diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index 2315bedfeb..4f756538fb 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -865,6 +865,7 @@ class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC): return MLAAttentionSpec( block_size=vllm_config.cache_config.block_size, num_kv_heads=1, + max_tp_shards=1, head_size=self.head_dim, dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype, tokens_per_state=self.compress_ratio, @@ -910,6 +911,7 @@ class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase): return MLAAttentionSpec( block_size=self.cache_config.block_size, num_kv_heads=1, + max_tp_shards=1, head_size=self.head_dim, dtype=self.dtype, tokens_per_state=self.compress_ratio, diff --git a/vllm/models/deepseek_v4/compressor.py b/vllm/models/deepseek_v4/compressor.py index e6e0fd959e..ccb864faa2 100644 --- a/vllm/models/deepseek_v4/compressor.py +++ b/vllm/models/deepseek_v4/compressor.py @@ -184,6 +184,7 @@ class CompressorStateCache(torch.nn.Module, AttentionLayerBase): return SlidingWindowMLASpec( # only has one vector instead of K + V block_size=self.block_size, num_kv_heads=1, + max_tp_shards=1, head_size=self.state_dim, dtype=self.dtype, sliding_window=self.sliding_window, diff --git a/vllm/v1/attention/backends/mla/sparse_swa.py b/vllm/v1/attention/backends/mla/sparse_swa.py index b0ef9e65a3..0f5bff4144 100644 --- a/vllm/v1/attention/backends/mla/sparse_swa.py +++ b/vllm/v1/attention/backends/mla/sparse_swa.py @@ -124,6 +124,7 @@ class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): bounded_replay=self.bounded_replay, block_size=self.block_size, num_kv_heads=1, + max_tp_shards=1, head_size=self.head_dim, dtype=self.dtype, sliding_window=self.window_size, diff --git a/vllm/v1/core/kv_cache_utils.py b/vllm/v1/core/kv_cache_utils.py index 7ef71dd471..a82d3e8113 100644 --- a/vllm/v1/core/kv_cache_utils.py +++ b/vllm/v1/core/kv_cache_utils.py @@ -9,7 +9,7 @@ import os from collections import defaultdict from collections.abc import Callable, Iterable, Iterator, Sequence from dataclasses import dataclass, field, replace -from functools import partial +from functools import partial, reduce from typing import TYPE_CHECKING, Any, NamedTuple, NewType, TypeAlias, cast, overload from vllm import envs @@ -2409,6 +2409,26 @@ def get_kv_cache_groups( return groups +def _layer_tp_replicas(spec: KVCacheSpec, tp_size: int, dcp_size: int) -> int: + if not isinstance(spec, AttentionSpec) or spec.max_tp_shards is None: + return 1 + if spec.dcp_sharded and dcp_size > 1: + return 1 + return max(1, tp_size // spec.max_tp_shards) + + +def kv_cache_groups_tp_replicas( + groups: list[KVCacheGroupSpec], tp_size: int, dcp_size: int = 1 +) -> int: + """Consecutive TP ranks holding identical KV for every layer.""" + specs = [spec for g in groups for spec in iter_layer_specs(g.kv_cache_spec)] + if not specs: + return 1 + return reduce( + math.gcd, (_layer_tp_replicas(s, tp_size, dcp_size) for s in specs), tp_size + ) + + def generate_scheduler_kv_cache_config( kv_cache_configs: list[KVCacheConfig], ) -> KVCacheConfig: @@ -2823,6 +2843,13 @@ def get_kv_cache_configs( vllm_config, groups, min_num_blocks * _pool_bytes_per_block(groups) ) + for kv_cache_config in kv_cache_configs: + kv_cache_config.kv_tp_replicas = kv_cache_groups_tp_replicas( + kv_cache_config.kv_cache_groups, + vllm_config.parallel_config.tensor_parallel_size, + vllm_config.parallel_config.decode_context_parallel_size, + ) + return kv_cache_configs diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 5791704b2b..cabbbfadc3 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -484,6 +484,9 @@ class HiSparseResidentSpec(KVCacheSpec): @dataclass(frozen=True, kw_only=True) class AttentionSpec(KVCacheSpec): dcp_sharded: bool = True + max_tp_shards: int | None = None + """Distinct shards this cache splits into across TP; TP ranks beyond this + hold replicas. None: unknown, treated as not replicated.""" num_kv_heads: int head_size: int dtype: torch.dtype @@ -618,6 +621,7 @@ class FullAttentionSpec(AttentionSpec): dtype=specs[0].dtype, kv_quant_mode=specs[0].kv_quant_mode, dcp_sharded=specs[0].dcp_sharded, + max_tp_shards=specs[0].max_tp_shards, page_size_padded=specs[0].page_size_padded, num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, @@ -705,6 +709,7 @@ class MLAAttentionSpec(FullAttentionSpec): dtype=specs[0].dtype, kv_quant_mode=specs[0].kv_quant_mode, dcp_sharded=specs[0].dcp_sharded, + max_tp_shards=specs[0].max_tp_shards, page_size_padded=specs[0].page_size_padded, num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, @@ -768,6 +773,7 @@ class RSWASpec(FullAttentionSpec): dtype=base.dtype, kv_quant_mode=base.kv_quant_mode, dcp_sharded=base.dcp_sharded, + max_tp_shards=base.max_tp_shards, page_size_padded=base.page_size_padded, num_head_slots=base.num_head_slots, state_content_bytes=base.state_content_bytes, @@ -960,6 +966,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec): bounded_replay_set = set(spec.bounded_replay for spec in specs) block_stride_alignment_set = {spec.block_stride_alignment for spec in specs} assert len({spec.dcp_sharded for spec in specs}) == 1 + assert len({spec.max_tp_shards for spec in specs}) == 1 assert ( len(cache_dtype_str_set) == 1 and len(tokens_per_state_set) == 1 @@ -980,6 +987,7 @@ class SlidingWindowMLASpec(SlidingWindowSpec): head_size=specs[0].head_size, dtype=specs[0].dtype, dcp_sharded=specs[0].dcp_sharded, + max_tp_shards=specs[0].max_tp_shards, page_size_padded=specs[0].page_size_padded, num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, @@ -1197,6 +1205,7 @@ class SinkFullAttentionSpec(FullAttentionSpec): dtype=specs[0].dtype, kv_quant_mode=specs[0].kv_quant_mode, dcp_sharded=specs[0].dcp_sharded, + max_tp_shards=specs[0].max_tp_shards, page_size_padded=specs[0].page_size_padded, num_head_slots=specs[0].num_head_slots, state_content_bytes=specs[0].state_content_bytes, @@ -1483,6 +1492,9 @@ class KVCacheConfig: hisparse_shared_host_pool: bool = False """Whether local TP ranks share one physical HiSparse host pool.""" + kv_tp_replicas: int = 1 + """Consecutive TP ranks holding identical KV for every layer (1: none).""" + @cached_property def transfer_group_ids(self) -> tuple[int, ...]: """IDs of cache groups that participate in external KV transfer."""