[#57652 suggestion] AttentionSpec.max_tp_shards + KVCacheConfig.kv_tp_replicas for TP-replication detection

[LucasWilkinson]

Suggested follow-up for vllm-project#57652.

What

• AttentionSpec.max_tp_shards: int | None = None: how many distinct shards this cache splits into across TP. TP ranks beyond that count hold replicas. None means unknown and is treated as not replicated.
  • It's opt-in: no spec class has a default that claims replication, so a layer that forgets it only loses dedup, never correctness. This is the same pattern as MambaSpec.tp_replicated.
  • max_tp_shards=1 is set only where the cache has been checked to be identical on every TP rank:
    • generic MLAAttention (mla_attention.py), which covers the DSv3 / V3.2 / GLM-5.3 latent cache;
    • the DSv3.2 indexer cache (deepseek_v2.py; wk_weights_proj uses disable_tp=True);
    • DeepSeek-V4 attention and indexer caches (deepseek_v4/attention.py);
    • the DeepSeek-V4 compressor state (fused_wkv_wgate uses disable_tp=True);
    • the DeepSeek-V4 SWA cache (sparse_swa.py).
  • Kimi-K3, MiniMax-M3, Qwen4-exp and DSv4.1 attention stay None until their owners opt in.
  • It's passed through every merge().
• KVCacheConfig.kv_tp_replicas: int = 1 replaces all_groups_are_tp_replicated. It's the number of consecutive TP ranks holding identical KV for every layer. get_kv_cache_configs sets it through kv_cache_groups_tp_replicas(groups, tp_size, dcp_size): the gcd of each layer's tp_size // max_tp_shards, taken per layer through iter_layer_specs. A layer counts as 1 if it isn't an attention spec, has no value set, or is DCP-sharded under DCP. The type(spec) in (MLA, SWA-MLA) check is gone.
• Offloading: replicated_layout requires tp > 1 and kv_tp_replicas == tp. The TP-only / no-PP / no-CP limits stay in build_offloading_config, as on main.

Why

Why a per-layer value on KVCacheConfig: MLAAttentionSpec shares uniform_type_base_spec=FullAttentionSpec, so one UniformTypeKVCacheSpecs group can mix MLA and GQA layers. generate_scheduler_kv_cache_config keeps an arbitrary layer spec, so the value has to be computed before that. Otherwise the scheduler could see an MLA spec and wrongly de-duplicate head-sharded KV.

Why a count, not a bool: it can express partial replication later. For example, GQA with 8 KV heads at TP16 is max_tp_shards=8, which gives kv_tp_replicas=2. The config field wouldn't need to change. Offloading could then use world_size // replicas copies, slot rank // replicas, and one writer per replica group. GQA and Mamba (tp_replicated) aren't wired up here; they default to 1, which is today's behaviour.

max_tp_shards is on the spec rather than tp_replicas because it doesn't depend on the TP size, so a layer can set it without knowing the TP size. The TP-dependent replica count is computed once, in get_kv_cache_configs.

Testing

pytest tests/v1/kv_connector/unit/offloading_connector/test_config.py tests/v1/core/test_kv_cache_utils.py
219 passed (3 PP tests deselected: they need more visible devices and fail identically on the base branch)

New test_kv_cache_groups_tp_replicas cases:

Case	TP	DCP	kv_tp_replicas
MLA	8	1	8
MLA	8	2	1
GQA, max_tp_shards=2	8	1	4
GQA, max_tp_shards=8	4	1	1
MLA + GQA, max_tp_shards=4	8	1	2 (gcd)
MLA, max_tp_shards unset	8	1	1
hidden-state	8	1	1

pre-commit (ruff, ruff-format, mypy-3.10) and mypy-3.12 pass on the changed files.

AI assistance (Claude Code) was used; I reviewed the changes.

🤖 Generated with Claude Code