[KV-Offloading][TP] : Expand replicated_layout detection to multi-group MLA

[varun-sundar-rabindranath]

Purpose

KV offloading for MLA models + TP need only offload KV values from one TP rank.
This is controlled by replicated_layout determination in vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py

Building on top of @almogtavor 's work #50667 (which was eventually closed) , this PR expand the replicated_layout check to include,

• SlidingWindowMLASpec
• Multi-group MLASpec / SlidingWindowMLASpec
  models.

Notably,

• deepseek-ai/DeepSeek-V4-Flash
• deepseek-ai/DeepSeek-V3.2-Exp and
• zai-org/GLM-5.3
  fall in this category.

Changes:

• AttentionSpec.max_tp_shards: int | None = None: how many distinct shards this cache splits into across TP. A value of means, all TP ranks are 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).
• KVCacheConfig.kv_tp_replicas: int = 1 : 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.
• Offloading: replicated_layout requires tp > 1 and kv_tp_replicas == tp. T

Justification for adding a KVCacheConfig member

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.

Test Plan

Add unit tests.

lm-eval tests : eval script
eval.sh
Methodology : Run models with vllm serve and disk-offloading. Run lm-eval 3 times.
lm-eval run 1 : cold run. The disk is filled.

lm-eval run 2 : warm run. offloaded KV data is loaded from Disk

lm-eval run 3 : warm run. offloadead KV data is loaded from Disk (this is just a confirmation run)

Test Result

Unit tests pass.

lm-evals:

DSV2-lite

local-completions ({'base_url': 'http://localhost:8000/v1/completions', 'model': 'deepseek-ai/DeepSeek-V2-Lite', 'tokenized_requests': False, 'num_concurrent': 1000}), gen_kwargs: ({'temperature': 0.0}), limit: 500.0, num_fewshot: 25, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  | 0.37|±  |0.0216|
|     |       |strict-match    |    25|exact_match|↑  | 0.37|±  |0.0216|

local-completions ({'base_url': 'http://localhost:8000/v1/completions', 'model': 'deepseek-ai/DeepSeek-V2-Lite', 'tokenized_requests': False, 'num_concurrent': 1000}), gen_kwargs: ({'temperature': 0.0}), limit: 500.0, num_fewshot: 25, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  |0.366|±  |0.0216|
|     |       |strict-match    |    25|exact_match|↑  |0.368|±  |0.0216|

local-completions ({'base_url': 'http://localhost:8000/v1/completions', 'model': 'deepseek-ai/DeepSeek-V2-Lite', 'tokenized_requests': False, 'num_concurrent': 1000}), gen_kwargs: ({'temperature': 0.0}), limit: 500.0, num_fewshot: 25, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  |0.366|±  |0.0216|
|     |       |strict-match    |    25|exact_match|↑  |0.368|±  |0.0216|

DSV4-Flash

local-completions ({'base_url': 'http://localhost:8000/v1/completions', 'model': 'deepseek-ai/DeepSeek-V4-Flash', 'tokenized_requests': False, 'num_concurrent': 1000}), gen_kwargs: ({'temperature': 0.0}), limit: 500.0, num_fewshot: 25, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  |0.956|±  |0.0092|
|     |       |strict-match    |    25|exact_match|↑  |0.956|±  |0.0092|

local-completions ({'base_url': 'http://localhost:8000/v1/completions', 'model': 'deepseek-ai/DeepSeek-V4-Flash', 'tokenized_requests': False, 'num_concurrent': 1000}), gen_kwargs: ({'temperature': 0.0}), limit: 500.0, num_fewshot: 25, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  |0.958|±  | 0.009|
|     |       |strict-match    |    25|exact_match|↑  |0.958|±  | 0.009|

local-completions ({'base_url': 'http://localhost:8000/v1/completions', 'model': 'deepseek-ai/DeepSeek-V4-Flash', 'tokenized_requests': False, 'num_concurrent': 1000}), gen_kwargs: ({'temperature': 0.0}), limit: 500.0, num_fewshot: 25, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  |0.958|±  | 0.009|
|     |       |strict-match    |    25|exact_match|↑  |0.958|±  | 0.009|

DSV3.2

local-completions ({'base_url': 'http://localhost:8000/v1/completions', 'model': 'deepseek-ai/DeepSeek-V3.2-Exp', 'tokenized_requests': False, 'num_concurrent': 1000}), gen_kwargs: ({'temperature': 0.0}), limit: 500.0, num_fewshot: 25, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  |0.952|±  |0.0096|
|     |       |strict-match    |    25|exact_match|↑  |0.952|±  |0.0096|

local-completions ({'base_url': 'http://localhost:8000/v1/completions', 'model': 'deepseek-ai/DeepSeek-V3.2-Exp', 'tokenized_requests': False, 'num_concurrent': 1000}), gen_kwargs: ({'temperature': 0.0}), limit: 500.0, num_fewshot: 25, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  |0.954|±  |0.0094|
|     |       |strict-match    |    25|exact_match|↑  |0.952|±  |0.0096|

local-completions ({'base_url': 'http://localhost:8000/v1/completions', 'model': 'deepseek-ai/DeepSeek-V3.2-Exp', 'tokenized_requests': False, 'num_concurrent': 1000}), gen_kwargs: ({'temperature': 0.0}), limit: 500.0, num_fewshot: 25, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  |0.946|±  |0.0101|
|     |       |strict-match    |    25|exact_match|↑  |0.946|±  |0.0101|

GLM5.3

local-completions ({'base_url': 'http://localhost:8000/v1/completions', 'model': 'zai-org/GLM-5.3', 'tokenized_requests': False, 'num_concurrent': 1000}), gen_kwargs: ({'temperature': 0.0}), limit: 500.0, num_fewshot: 25, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  |0.922|±  | 0.012|
|     |       |strict-match    |    25|exact_match|↑  |0.922|±  | 0.012|

local-completions ({'base_url': 'http://localhost:8000/v1/completions', 'model': 'zai-org/GLM-5.3', 'tokenized_requests': False, 'num_concurrent': 1000}), gen_kwargs: ({'temperature': 0.0}), limit: 500.0, num_fewshot: 25, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  |0.946|±  |0.0101|
|     |       |strict-match    |    25|exact_match|↑  |0.946|±  |0.0101|

local-completions ({'base_url': 'http://localhost:8000/v1/completions', 'model': 'zai-org/GLM-5.3', 'tokenized_requests': False, 'num_concurrent': 1000}), gen_kwargs: ({'temperature': 0.0}), limit: 500.0, num_fewshot: 25, batch_size: 1
|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|    25|exact_match|↑  |0.928|±  |0.0116|
|     |       |strict-match    |    25|exact_match|↑  |0.928|±  |0.0116|

Savings compared to main (running the lmeval workloads above)

<style type="text/css"></style>

MAX DISK SIZE = 177G
model	Main (FS Occupied disk size)	PR (FS Occupied disk size)
Dsv2 lite (TP=2)	70G	70G
Dsv4-flash (TP=8)	157G	20G
Glm5.3 (TP=8)	177G (ENOSPC)	103G
DSV3.2 (TP=8)	177G (ENOSPC)	149G
CPU KV cache block size - Note the savings scale with TP size
model	Main (bytes)	PR (bytes)
Dsv2 lite (TP=2)	499712	499712
Dsv4-flash (TP=8)	40103936	5013504
Glm5.3 (TP=8)	32071680	4009984
DSV3.2 (TP=8)	27619328	3452928

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[varun-sundar-rabindranath]

cc @almogtavor @Etelis @bnellnm @tlrmchlsmth

[varun-sundar-rabindranath]

I have another set of evals going - will post results when done. But I think the PR is in a good state to review anyways.
Edit : Updated ✅

[varun-sundar-rabindranath]

also cc @LucasWilkinson

[bnellnm]

LGTM.

[Etelis]

Thanks @varun-sundar-rabindranath!! This addresses the missing dedup support. Can we also make the suggested edits to fix the remaining regressions? We have the canonical fallback fix on our local branch fix/mla-canonical-fallback-sizing.

[Etelis]

A uniform group can contain MLA alongside a regular-attention draft layer. The scheduler has already replaced the group with its first MLA spec, so this enables dedup there while the workers reject it. They then calculate different row sizes and capacities. Could we preserve the per-layer eligibility when building the scheduler config?

Restoring the old size check would also lose the intended support for differently sized MLA/indexer pages.

[varun-sundar-rabindranath]

Hi @Etelis - This is a nice catch. I have been exploring a solution for this, and haven't come up with any simpler fixes.

1. The straight forward fix is to not flatten the Scheduler KV config and update the Scheduler code to resolve UniformKVCacheSpecs at inspection time. But, the Scheduler KV cache is ubiquitous and the audit surface area non-trivial.
2. In the Schdeduler process, we can keep both the raw KV configs and the flattened ones. and use the unflattened ones just in the Offloading Connector. This again is messy.
3. I also explored resolving this as part of KV connector handshaking (Offloading Connector doesn't do any handshaking at the moment) . But we create all the CPU KV cache memory even before the handshaking point. This approach will let us detect config mismatches, instead of running with mismatched configs. But this feels hairy too.

Continuing to see if there are better ways of doing this.

[Etelis]

Thanks Varun. Could we preserve the per-layer specs for offloading before flattening, while leaving the scheduler’s current representation intact?

[varun-sundar-rabindranath]

Hi @Etelis . PTAL this commit 9efc2a2

I moved the replication detection directly into the KVCacheConfig so there is no divergence. it seems like the best among the options I have considered. wdyt ?

[varun-sundar-rabindranath]

I think there is still a gap with the commit 9efc2a2 . The offloading config need only deal with the groups that can be offloaded for TP dedup - i.e.

vllm/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py

Line 71 in 09fe178

selected_groups = tuple(

by moving the replication detection away from the offloading connector, we are being stricter (checking all groups) than necessary. This I think is a downside. Am I thinking about this correctly ?

Edit : Gave it some thought. A cleaner way is to consider only "prefix-cacheable" groups in determining KVCacheConfig.all_groups_are_tp_replicated. Probably could add it as a follow up.

[Etelis]

Yes, agreed. This can be a follow-up. I’d match the offloading group selection and keep checking every inner spec before scheduler flattening. That also covers HiSparse. Thanks for including the canonical fallback fix!

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @varun-sundar-rabindranath.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[Etelis]

Just two small test suggestions below.

[Etelis]

Can we also run this through generate_scheduler_kv_cache_config and check that worker and scheduler agree on replication and CPU row size? One MLA/indexer case and one MLA/full-attention case should cover it. We can reuse the existing helpers to keep this small.

[varun-sundar-rabindranath]

Added test_replicated_layout_for_dsa_style_groupsand removed test_replicated_layout_rejects_bare_mla_with_mixed_page_accounting that is on main. see below.

[varun-sundar-rabindranath]

@Etelis - all replicated_layout checks go through additional scheduler / worker alignment.

[varun-sundar-rabindranath]

@Etelis - I removed the bare case and used UniformTypeKVCacheSpecs to group the specs.
I verified from running DSV4 (this is group 2)

(Worker_TP2 pid=3528120) group[2]  wrapper=UniformTypeKVCacheSpecs  (62 layer specs)
(Worker_TP2 pid=3528120)   layer=model.layers.2.attn.indexer.k_cache  spec=MLAAttentionSpec  num_kv_heads=1
(Worker_TP2 pid=3528120)   layer=model.layers.2.attn  spec=MLAAttentionSpec  num_kv_heads=1
(Worker_TP2 pid=3528120)   layer=model.layers.3.attn  spec=MLAAttentionSpec  num_kv_heads=1
(Worker_TP2 pid=3528120)   layer=model.layers.4.attn.indexer.k_cache  spec=MLAAttentionSpec  num_kv_heads=1
(Worker_TP2 pid=3528120)   layer=model.layers.4.attn  spec=MLAAttentionSpec  num_kv_heads=1
(Worker_TP2 pid=3528120)   layer=model.layers.5.attn  spec=MLAAttentionSpec  num_kv_heads=1

.

However I am not sure if bare mla, the way the test was originally written was a real possibility. Any insight here will be very helpful. Thanks.

[varun-sundar-rabindranath]

Test to verify that canonical_layout resets the replicated_layout setting. Please review.

[Etelis]

LGTM

Cc @orozery

[varun-sundar-rabindranath]

I think this might be cleaner / less brittle if add a max_tp_shards attribute to AttentionSpec.max_tp_shards; thoughts @varun-sundar-rabindranath ?

like: neuralmagic#176

Thanks @LucasWilkinson . I like the idea of declaring TP sharding at the Spec level (much closer to the source and avoids checks on Spec types).
I have merged into this PR.

cc @Etelis @bnellnm Please take a look ! Thanks 🙌

[varun-sundar-rabindranath]

/ci run

[github-actions]

❌ This PR is 17 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[varun-sundar-rabindranath]

/ci run

[github-actions]

✅ Triggered Buildkite CI #91660 for commit be2521833025.

[varun-sundar-rabindranath]

Thanks @LucasWilkinson - I ran the e2e eval test with the new changes and I can repro what is in the PR description.

[varun-sundar-rabindranath]

/ci run

[github-actions]

❌ This PR is 3 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[varun-sundar-rabindranath]

/ci run

[github-actions]

✅ Triggered Buildkite CI #91915 for commit 4d9c27948f34.

[LucasWilkinson]

LGTM thanks!

[varun-sundar-rabindranath]

Update : Having bit of a hard time trying to figure out why this test [buildkite/ci/pr/amd-mi300-lm-eval-kv-offload](https://buildkite.com/vllm/ci/builds/91915#01a0ee85-ace4-4bfc-84a7-ad2f3cc879af) buildkite/ci/pr/amd-mi300-lm-eval-kv-offload https://buildkite.com/vllm/ci/builds/91915 is failing .

The first couple of failures are real (but not sure if the test is just flaky). The tests were added yesterday with #50519 . in subsequent retries the test itself passes, but run-amd-test.sh reports getting a coredump with SIGABRT.

Digging into this further.

jfyi cc @AndreasKaratzas

[Edit]: A new retry passed the test. also, I could get the accuracy to drop on the nemotron model on main on H100.

[AndreasKaratzas]

@varun-sundar-rabindranath this should be fine, we are going to debug this. I see that this is a 134 exit code too.