Repository navigation
[Core][KV Connector] Keep the connector's head order when its KV layout is unsupported - #59594
lucifer1004 wants to merge 1 commit into
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
…order when its KV layout is unsupported
| @property | ||
| def is_head_major(self) -> bool: | ||
| """True when H is outside N, so each head's tokens are contiguous.""" | ||
| return self.value.index(_DIM_H) < self.value.index(_DIM_N) |
There was a problem hiding this comment.
should it be self.value[_DIM_H] < self.value[_DIM_N] instead?
There was a problem hiding this comment.
No, I believe index() is right here. Each member's value lists the logical dims in physical order, so a dim's physical position is value.index(dim), not value[dim]. For example, BHLNC = (1, 2, 0, 3, 4) is physically [B, H, L, N, C]. The neighboring properties read value the same way: is_block_outermost is value[0] == _DIM_B, is_block_contiguous compares value[-3:], and layer_view_order iterates over value.
value[_DIM_H] is the logical dim stored at physical position 2. For the six layouts defined today, the two expressions happen to agree: five of the permutations are their own inverse, and BHLNC gives the same answer either way. They diverge for a layout like LHNBC = (0, 2, 3, 1, 4), physically [L, H, N, B, C]. Each head's tokens are contiguous there, and value.index(2) = 1 < value.index(3) = 2 says so, but value[2] = 3 > value[3] = 1 would report it as not head-major. Happy to add a comment on the convention next to _DIM_* if that would help.
…order when its KV layout is unsupported
…order when its KV layout is unsupported
|
This pull request has merge conflicts that must be resolved before it can be |
…ut is unsupported When the KV connector's requested layout is not supported by every backend, resolve_kv_cache_layout fell back to the backends' first choice. Mooncake and NIXL request LBHNC so a heterogeneous-TP transfer can slice each page by KV head. A hybrid model whose backends only support block-outer layouts (QSA lists BLNHC, BLHNC) then got BLNHC, where heads are interleaved per token, and a page can no longer be cut by head: Mooncake registers it as one opaque row and the TP1 -> TP2 handshake fails on region lengths. Fall back to the first candidate with the same head/token order as the requested layout, here BLHNC, and only then to the backends' first choice. Every connector that requests a layout does so for that order (head-major for transfer and offloading, token-major for the example hidden-states connector), so the rule matches each request's intent. Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
c18d185 to
bca75c9
Compare
…when its KV layout is unsupported Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com>
This branch has not been deployed
2 participants
Overview
When the KV connector's requested layout is not supported by every backend,
resolve_kv_cache_layoutnow falls back to the first candidate that keeps the requested head/token order, instead of the backends' first choice.Claims
LBHNCso heterogeneous-TP transfers can slice pages by KV head. On a hybrid model whose backends only accept block-outer layouts, the fallback now picksBLHNC(head-major) instead ofBLNHC, so a page can still be cut by head.BLNHCthe handshake fails.VLLM_KV_CACHE_LAYOUTset, are unchanged.Validation
New tests in
tests/v1/core/test_kv_cache_utils.py:test_connector_layout_fallback_keeps_head_order(anLBHNCrequest against[BLNHC, BLHNC]resolves toBLHNC, and against[BLNHC]still falls back toBLNHC), andtest_kv_cache_layout_is_head_major. The first case fails without the change.pytest tests/v1/core/test_kv_cache_utils.py tests/distributed/test_kvlayout.py tests/v1/kv_connector/unit/test_mooncake_connector.py tests/v1/kv_connector/unit/test_mooncake_connector_hybrid_mamba.py: 285 passed.pre-commit hooks pass on the changed files.
End to end, on RTX PRO 6000 with Mooncake over RDMA, Qwen3.8-Flash-Next (QSA attention + GDN, 2 KV heads), prefill TP1 → decode TP2. These runs set
VLLM_KV_CACHE_LAYOUT=BLHNC, the layout this change now selects, and used the connector with [Bugfix][KV Connector] Mooncake: heterogeneous-TP support for hybrid GDN/Mamba models #53078 applied for the GDN state.BLNHC(today's fallback): every pull fails the handshake,local=39739392(TP1 page registered as one opaque row) vsremote=1605632.BLHNC: GSM8K (200 questions, 32 in flight) 0.970, 0 errors. A TP2 → TP2 pair scores 0.985 on the same set; the extra misses are coherent reasoning, not garbled output.Performance at equal TP2 (single instance, no connector,
VLLM_KV_CACHE_LAYOUTset, 8K in / 1K out random prompts, output tok/s):BLNHC vs BLHNC
TP2, MTP3 speculative decoding:
TP2, no speculative decoding:
p50 TPOT is the same within noise at every concurrency.
Details
Mooncake and NIXL request
LBHNCbecause a heterogeneous-TP transfer splits each page by KV head, which needs each head's tokens to be contiguous. QSA declares(BLNHC, BLHNC), soLBHNCis not a candidate, and the fallback tookBLNHC, where heads are interleaved per token. At TP2 a page holds one head, so it is still contiguous and equal-TP transfers work. At TP1 it holds two, cannot be sliced, and registers as one row the size of the whole block.The fallback now picks the first candidate whose head/token order (
KVCacheLayout.is_head_major, H outside N) matches the request, and otherwise the backends' first choice as before. The other layout requests follow the same rule: offloading and HiSparse ask for head-major layouts, and the example hidden-states connector asks forLBNHCto keep each token's data together, so an unsupported request now falls back to a token-major candidate.Related: #53078 (Mooncake heterogeneous TP for hybrid GDN models) and #59592 (Mooncake block-size handshake), from the same TP1/TP2/TP4 testing. #58380 changes the layout order the worker reports for mixed page sizes and keeps connector requests taking precedence; it does not touch this fallback.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
AI assistance (Claude Code) was used to write the change, tests, and this description; I reviewed every changed line and ran the tests above.
🤖 Generated with Claude Code