[stu-cao]

Stacked on #58471 and #57662 (both open). This branch carries their commits. Only the last three commits belong to this PR (+328/−19 across base_worker.py, test_nixl_connector.py and the connector docs):

• 57978684fc: opt-in whole-row regions;
• edb4a3c77b: reject draft pages packed into the target row;
• 9adc138cfc: let the test fake take a mixed-spec group.

Reviewing #57662 and #58471 first will shrink this diff to those three commits.

Purpose

Fixes #60561.

Since #53781, a device-resident BLHNC MLA cache registers one NIXL region per layer, because page_contiguous is true for every layer. For GLM-5.3 that is 99 descriptors per block instead of one, so BLHNC transfers cost more than LBHNC's (104 vs 62 strides on a decode rank).

This PR adds kv_connector_extra_config.whole_row_regions (default false):

• Replicated (MLA) layers of a block-major row register as one whole-row region, as they did before [3/N] HiSparse: host-resident sparse-MLA decode hot-buffering #53781. The row length is the block stride, as on main.
• Head-sharded layers (a GQA draft under an MLA target) never take the whole-row branch. They keep per-layer SPLIT regions, so each TP rank reads its own head slice ([Bugfix][KV Connector][NIXL] Shard a GQA draft's KV by head under an MLA target #58471).
• Non-MLA models are unaffected; the flag only changes MLA regions.
• Rejected configurations. The flag errors out at registration for:
  • kv_buffer_device=cpu;
  • HiSparse host-resident groups, so [3/N] HiSparse: host-resident sparse-MLA decode hot-buffering #53781's case keeps its current behaviour;
  • a layout that packs head-sharded draft pages into the target's group and row. This is BLHNC with --disable-hybrid-kv-cache-manager. The row's size depends on TP there, so prefill and decode rows would differ, and without the check this only failed later as a handshake assert.

Why #57662 is needed. Under BLHNC every KV cache group starts at byte 0 of the shared row, so the first draft page aliases the first target page's base address. Without #57662's (base_addr, block_len) key, the draft region folds into a target region. The rebase keeps main's rule that a packed PP producer never dedups aliased regions.

Not a duplicate:

• [Bugfix][NIXL] Dedup transfer regions by (base_addr, block_len) #57662 fixes only the dedup key.
• [Bugfix][KVConnector] Split MLA kernel rows larger than the NIXL block into dense pages #59586 splits MLA kernel rows larger than the NIXL block into dense pages. That changes per-layer page geometry; it does not register one region per block row.
• [Bugfix][NIXL] Preserve block geometry for packed MLA caches #50717 preserves packed-MLA block geometry.
• [Review only][NIXL] Simplify draft head slicing from #58471 #60368 simplifies [Bugfix][KV Connector][NIXL] Shard a GQA draft's KV by head under an MLA target #58471's head slicing.

None of them restores whole-row registration for page-contiguous caches in pull mode.

Separate issue: the DFlash2 acceptance drop seen on BLHNC deployments (#60560) does not come from NIXL registration. It comes from the hybrid KV cache manager's sliding-window draft group, and is fixed by #60767. This PR neither causes nor fixes it.

Test Plan

pytest tests/v1/kv_connector/unit/test_nixl_connector.py tests/v1/kv_connector/unit/test_nixl_desc_geometry.py tests/v1/kv_connector/unit/test_tp_mapping.py

New tests:

• test_blhnc_gqa_draft_reads_own_head_slice_from_tp1_prefill, parametrized over whole_row in {False, True} and tp_rank in {0, 3}. A TP1 prefill and a TP4 decode each register a BLHNC MLA target plus a GQA draft. With whole-row on, the target row must be one region. Every draft layer must be its own region and read its own head slice.
• test_whole_row_regions_rejects_host_buffer.
• test_whole_row_regions_rejects_draft_packed_into_target_row.

End to end: GLM-5.3 (NVFP4) with a DFlash2 draft on 2× B300, DEP8 prefill to 2× TP4 decode over NIXL (RDMA), BLHNC block size 64, whole_row_regions on both roles, agentic multi-turn traffic (AgentX) at 32 concurrent sessions.

Test Result

Unit tests. One B300 GPU, test_nixl_connector.py, test_nixl_desc_geometry.py and test_tp_mapping.py. Each tree was built with VLLM_USE_PRECOMPILED=1 on the kernels of nightly 51ef239768.

• main (81d9c45e4b): 424 passed, 2 failed.
• [Bugfix][KV Connector][NIXL] Shard a GQA draft's KV by head under an MLA target #58471 rebased: 438 passed, 2 failed.
• This PR's head: 445 passed, 2 failed.
• Same two failures everywhere. test_abort_timeout_on_prefiller[None] and [ray] start a real engine that wants 50% of GPU memory, and only 17.6 of 268 GiB were free on the shared GPU. No failures are introduced by this PR.
• New tests (-k "blhnc_gqa_draft or whole_row_regions"): 6 of 6 pass.
• Lint: ruff and mypy pass on the changed files.

End to end, with this PR's code. These are from the B300 runs above. The NIXL connector's own metrics are summed over the profiling window of each run (~5,700 transfers each):

BLHNC, block 64	regions	descriptors / transfer	descriptors / MB	transfer time / GB
whole_row_regions on (4 runs)	target row: 1; draft: per layer	1,309 – 1,860	0.48 – 0.54	48 – 69 ms
per-layer regions	1 per layer	46,692	49.3	77 ms

The per-layer row ran with --disable-hybrid-kv-cache-manager, which changes what is transferred (decode prefix hits rise and less KV is pulled), so its transfer time is indicative only. The descriptor reduction is the direct effect of this PR.

Earlier variant, for context. On GB200 (128 sessions), a predecessor of this change produced:

• Registration: 1 stride / 26,104 descriptors on a decode rank, vs 62 / 2,488,080 with LBHNC block 1024.
• Transfer time: mean NIXL transfer time 85.23 → 22.30 ms.
• TTFT: p50 848 → 730 ms.

That predecessor also folded the head-sharded draft into the whole-row region, which this PR deliberately avoids, so these are not this PR's numbers. The target's registration is the same; this PR adds the draft's per-layer regions on top.

Not covered: push mode (main's packed-MLA push path is unchanged) and non-MLA models (the flag does nothing there).

AI assistance

AI assistance (Claude) was used to write the code, tests and this description, and to run and analyze the experiments. @stu-cao reviewed the change and is responsible for it.

🤖 Generated with Claude Code

[mergify]

Documentation preview: https://vllm--60564.org.readthedocs.build/en/60564/

[mergify]

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

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

[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.

[stu-cao]

Test results on the final head 4f86d7da69, vs base c0a10889e0 (#58471 head).

Setup: one B300 GPU; vLLM built from each tree with VLLM_USE_PRECOMPILED=1 (merge-base nightly 711fc55c10).

pytest tests/v1/kv_connector/unit/test_nixl_connector.py -k "blhnc_gqa_draft or whole_row_regions"
  head: 6 passed
  rejection test with e75df93337's base_worker.py (no guard): DID NOT RAISE ValueError (fails, as intended)

pytest tests/v1/kv_connector/unit/test_nixl_connector.py
  base: 153 passed, 3 failed
  head: 159 passed, 3 failed   (same 3 failures as base; no regressions)

All three failures are environmental:

• test_kv_transfer_handshake: a live engine held the inherited VLLM_NIXL_SIDE_CHANNEL_PORT. It passes on both trees with that variable unset.
• test_abort_timeout_on_prefiller[ray] / [None]: they start a real engine that wants 50% of GPU memory, and only 17.6 of 268 GiB were free on the shared GPU.

[mergify]

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

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

[ivanium]

Thanks for the PR and sry for the delayed review! IIUC, #61074 fixed the whole row issue and can supersede this PR. #57662 may still be needed though but that one seems still ongoing

[stu-cao]

@ivanium Yes, I validated the fix is already in main

┌────────────────────┬────────────────────────────┬───────────┬─────────────┬────────────┬─────────┐
│        Run         │    Regions per worker      │ Transfers │ Descriptors │    Per     │   MB    │
│                    │     (prefill / decode)     │           │             │  transfer  │  moved  │
├────────────────────┼────────────────────────────┼───────────┼─────────────┼────────────┼─────────┤
│ main (#61074),     │ 1 / 1                      │       192 │         432 │        2.2 │  37,671 │
│ BLHNC              │                            │           │             │            │         │
├────────────────────┼────────────────────────────┼───────────┼─────────────┼────────────┼─────────┤
│ #60564 rebuild,    │ 1 / 1                      │       192 │         420 │        2.2 │  37,671 │
│ flag off           │                            │           │             │            │         │
├────────────────────┼────────────────────────────┼───────────┼─────────────┼────────────┼─────────┤
│ #60564 rebuild,    │ 1 / 1                      │       192 │         436 │        2.3 │  37,671 │
│ flag on            │                            │           │             │            │         │
├────────────────────┼────────────────────────────┼───────────┼─────────────┼────────────┼─────────┤
│ main, LBHNC        │ 27 / 27                    │       192 │      12,420 │       64.7 │  37,671 │
│ (control)          │                            │           │             │            │         │
└────────────────────┴────────────────────────────┴───────────┴─────────────┴────────────┴─────────┘

Closing as superseded by #61074