[Bugfix][NIXL] Fix host-buffer h2d sync overwriting received KV

[lucamotz]

Overview

With kv_buffer_device="cpu", sync_recved_kv_to_device overwrites D's received KV with host pages NIXL never filled, for models with several KV cache groups or KV-sharing layers. Each group now copies only its own layers' host buffers; the multi-group case regressed in #53780.

Claims

• With host buffers, D keeps the received KV for models with several KV cache groups or KV-sharing layers.
• Each layer of a transfer group is still copied once per request, as after [KV Connector][NIXL] Coalesce host-buffer KV copies across cache groups #54483; KV-sharing layers are no longer copied.

Validation

CPU, this PR (ff55fa35ed) on main efaca9547d:

• pytest tests/v1/kv_connector/unit/test_nixl_connector_hma.py -k test_sync_recved_kv_copies_each_group_own_layers: [KV Connector][NIXL] Coalesce host-buffer KV copies across cache groups #54483's parametrized test now checks which layers and block ids each h2d copy gets, with a KV-sharing layer added. All 3 rows fail on main; with this PR they pass. The rest of test_nixl_connector_hma.py, test_nixl_connector.py and test_nixl_push_connector.py has the same pass/fail set as main.
• A local byte-level check runs real allocate_kv_cache tensors, save_kv_to_host, NIXL replayed between host buffers and sync_recved_kv_to_device on toy models at TP 1/1 and 2/2. On main, D's KV differs from P's in all 16 multi-group cases, which mix sliding-window and full-attention layers, and matches in all 4 full-attention-only cases. With this PR it matches in all 20.

GB10, Gemma-4-E2B-it, which has sliding-window and full-attention groups and 20 KV-sharing layers, P TP1 → D TP1, kv_buffer_device="cpu" on both sides, against the same run with device buffers:

• Main a3e0243b1c: 0/4 outputs identical, max |Δlogprob| per prompt 2.44, 2.63, 2.56, 3.04.
• This PR (f62c2fdc78 on main a3e0243b1c): 4/4 outputs identical, Δlogprob 0 on all 4 prompts, also in a repeat run.

GB10 setup

Both runs, P and D: --attention-backend TRITON_ATTN --enforce-eager --max-model-len 8192 --block-size 64 --no-enable-prefix-caching, greedy, 4 prompts of 1100–1300 tokens; D read 4880 tokens over NIXL.
Host-buffer run, P and D: "kv_buffer_device": "cpu" and --num-gpu-blocks-override 2048; both logged use_host_buffer: True.

Details

On D, with kv_buffer_device="cpu":

1. NIXL READ writes into D's per-layer host buffers. Each host buffer is its own region, and since [2/N][KV Connector][NIXL] Support per-region transfer geometry #53780 each group's descriptors cover only its own regions, so a layer's host buffer is filled only at its own group's blocks. Layers that reuse another layer's KV cache have host buffers but no transfer group, so NIXL never fills them.
2. sync_recved_kv_to_device copies the host buffers into the device tensors, which layers of different KV cache groups share, as does a KV-sharing layer with its target.
   • Main copies every layer's host buffer at all groups' blocks, so unfilled host pages overwrite the other groups' received blocks. A KV-sharing layer, copied after its target, overwrites the target's received KV, even with one group.
   • This PR copies each group's layers at that group's blocks only.

The byte-level check matches in all 10 multi-group cases that build at #53780's parent and differs in all 10 at #53780.

save_kv_to_host still copies every layer for all groups' blocks; NIXL never reads the extra host pages.

#59037, a 60-file draft for hidden-state extraction with packed layouts, contains the same per-group copy; this PR is the standalone fix.

---

Pull Request Checklist

• I used vLLM's /pr-checklist skill. (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 for the analysis, patch and tests. I reviewed every changed line.

🤖 Generated with Claude Code

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