Repository navigation
[KV Connector][NIXL] Register packed MLA KV rows as a single region - #61074
Conversation
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
/ci run --allow-stale |
|
✅ Triggered Buildkite CI #94324 for commit
|
CI selector (shadow): 13 test steps (17 jobs) instead of 19 (25 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (13)
Would skip (today's rules run them) (6)
Would add (today's rules do not run them) (0)none AMD mirrors: would skip (15)
AMD mirrors: would add (0)none 3 changed files · base |
|
/ci run |
|
❌ This PR is 131 commits behind upstream |
|
/ci run |
|
✅ Triggered Buildkite CI #94492 for commit |
|
This pull request has merge conflicts that must be resolved before it can be |
ivanium
left a comment
ivanium
left a comment
There was a problem hiding this comment.
LGTM. One minor comment but I think it's okay to merge with or without it.
| # With HiSparse, P may be on BL* while D is LB*, as only indexer is on HBM. | ||
| # TODO (NickLucche) we may consider not allowing this case at all (P BL*) | ||
| # and just clear this function. | ||
| self._split_remote_packed_rows(nixl_agent_meta) |
There was a problem hiding this comment.
I kinda agree and would prefer to reject this case directly.
With a block-outermost layout (e.g. VLLM_KV_CACHE_LAYOUT=BLHNC) every layer's page of a block lives in one contiguous row. Register that row as one NIXL region in pull mode too, so a block is a single descriptor instead of one per layer. Peers with per-layer regions (HiSparse decode, layer-outer layouts) address pages inside the row via the advertised packed_member_layouts. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
… rows Explain _split_remote_packed_rows with a small diagram, and run the whole-row registration test for the pull worker too, so packed BLHNC rows cannot silently fall back to one region per layer again. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
DeepSeek-V4.1 registers a compressor-state CircularBufferSpec ring, and _is_csa_linear then forced packed_storage off. Its seven transfer groups alias each block row from byte 0, so it fell back to overlapping per-layer regions: about 11k NIXL descriptors per 100k-token request. Gate whole-row transfer on what it needs instead: every page in the row must be TP-replicated, which holds for MLA without SSM state. Qwen3.8-Flash-Next (Mamba) and non-MLA models keep per-layer regions. Cherry-picked from vllm-project#61121 (c673ff0). Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Co-authored-by: Codex <noreply@openai.com>
6ee70a1 to
5f18e79
Compare
|
/ci run |
|
✅ Triggered Buildkite CI #95032 for commit |
…llm-project#61074) Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai> Signed-off-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Yifan Qiao <yifanqiao@inferact.ai> Co-authored-by: Codex <noreply@openai.com>
Purpose
With a block-outermost KV cache layout (
VLLM_KV_CACHE_LAYOUT=BLHNC), all layers' pages of a block sit in one contiguous row.NIXL still registers one region per layer, so each block moves as one descriptor per layer (effectively rendering BL* useless).
This PR registers the row once, fixing registration back to its intended use.
Example: GLM-5.3-NVFP4, TP4, BLHNC. Input to
register_kv_cacheson rank 0 is 99 views over one 43.8 GB storage:Before (
elif storage_is_block_major:branch), each view is its own region:After (the existing whole-row branch, previously push-only), every view emits the same region and dedup folds them into one:
HiSparse
A HiSparse decoder can't register rows: its MLA cache is in host DRAM and its indexer on GPU, so it keeps per-layer regions LB*.
P can still advertise each layer's
(offset, page_size)within the row (BL*).When the peer is per-layer,
_split_remote_packed_rows(added with this PR) maps P's row onto the peer's per-layer regions.BLHNC P + HiSparse D therefore works as before, with the same descriptor count as
main.cc @LucasWilkinson the alternative is to enforce P also uses LB* here, not sure if there's disadvantages (considering .
Out of scope: GQA spec-decode drafts under an MLA target, deferred to #60564.
Test Plan
pytest tests/v1/kv_connector/unit/test_nixl_desc_geometry.py test_nixl_connector.py test_nixl_connector_hma.py test_nixl_push_connector.pytest_packed_mla_pp1_transfers_whole_rowsnow covers pull and push and asserts 1 region / 1 descriptor per block, guarding against falling back to per-layer regions again.tests/evals/gsm8k/gsm8k_eval.py --num-questions 500(5-shot) throughtoy_proxy_server.py.Test Result
main, the new pull-mode cases oftest_packed_mla_pp1_transfers_whole_rowsfail (per-layer regions); with this PR they pass.Related PRs
#60851 and #60564 target the same descriptor blow-up. #60564 puts whole rows behind an opt-in
whole_row_regionsflag that must be set on both P and D, and rejects HiSparse. Here it is the default for BLHNC, needs no flag, and HiSparse decoders keep working.#60564 also handles GQA draft head-sharding, which this PR leaves to it.