[NickLucche]

Purpose

Fix a NIXL pull that can silently load the wrong prompt pages when P and D allocate KV differently. The motivating setup has plain NIXL on P and NIXL + HiSparse on D, with --block-size 128 on both servers.

The same command-line block size does not guarantee the same allocation. A logical block is the unit the scheduler allocates; a physical page is the unit used by the attention kernel and transfer. With a kernel that requires 64-token pages, plain P can keep 128-token logical blocks and split each into two physical pages. HiSparse's setup on D resolves its MLA scheduler blocks to the native 64-token size directly. That produces this geometry in the affected path:

	Command-line block size	Scheduler's logical block	Physical/kernel page
P: plain NIXL, HiSparse disabled	128	128	64
D: NIXL + HiSparse	128	64	64

For a 160-token prompt, P must allocate two logical blocks: capacity for 256 tokens, or four physical pages. D needs only three physical pages. Let A and B each hold 64 tokens, and C hold the remaining 32:

P advertises:       [ A | B | C | unused ]
D has space for:    [   |   |   ]

Current selection:  [ B | C | unused ]  <- last three producer pages
Correct selection:  [ A | B | C ]

Logical-to-physical expansion already happens. The mistake comes afterward: the region-based suffix matcher assumes that any extra pages on P represent a prefix D has already cached. Here the extra page is unused capacity at the end, so that assumption shifts the whole transfer. If D has actually cached A, the remaining transfer should select B, C; the old logic selects C, unused.

The fix selects pages by token position. Pass the producer's existing remote_num_tokens through to the worker, combine it with each destination group's cached-prefix offset, and select the corresponding physical pages. Whole-page transfers remain: C is transferred with its unused rows, but the entirely unused fourth page does not shift the selection. This adds no prefill on D.

HiSparse also exposes a separate assumption at transfer completion. P can keep main KV and indexer KV in one merged cache group. D separates host KV and GPU indexer KV into different groups, so NIXL uses regions to match the corresponding memory buffers. The old get_transfer_results() still pairs groups by position. Looking up P's group 1 for D's group 1 can fail because P has only group 0, or pair the wrong blocks when the index exists.

Completion now uses the region read plan. The worker records region_blocks_to_zero: whole destination pages that were allocated but excluded from the read. For example, in the reverse allocation geometry, D may reserve four pages while P supplies three:

D after the read:   [ A | B | C | untouched ]
D after completion:[ A | B | C | zeroed    ]

For two example regions with destination IDs [10, 11, 12, 13] and [20, 21, 22, 23], receiving three pages into each produces a cleanup plan of [[13], [23]]. Completion clears page 13 only in its region and page 23 only in its region, then skips the old positional group processing.

This plan is local to one request on the receiving worker. None means use the existing completion path; empty per-region lists mean a region-mapped read with nothing to zero. In the original P-four-pages/D-three-pages example, D receives all its pages, so there is no zeroing. The plan does not include cached prefix pages or unused rows within the transferred partial page. It is discarded with the request's receive metadata; failed reads skip this cleanup.

Scope: pure MLA region pulls with equal physical page sizes, DCP=1, and one cache-group owner per region. Heterogeneous logical geometry without the required metadata or with shared-group regions is rejected. Existing homogeneous fallback behavior remains.

Dependency: the empty-receive fixes belong to #57049 (also present in #56946, originally from #55398). They are excluded from this PR, which will be rebased after that prerequisite lands. #57049's HiSparse tail-import policy remains separate from this transfer-alignment fix.

Duplicate checks included related issue #57019, its comments, and open PR searches for that issue, NIXL logical/physical geometry, and region padding. #57027 fixes layer binding offsets. #55374 introduces broader piecewise token-window loading in the ordinary group-based pull/push paths; its current diff does not wire that through the region-based path introduced by #53780. This PR addresses that region path without adding the piecewise-loading API.

AI assistance: OpenAI Codex assisted with investigation, implementation, and validation. The human submitter is responsible for reviewing the changed lines and validating the patched GPU deployment. GPU accuracy validation remains pending.

Test Plan

.venv/bin/python -m pytest tests/v1/kv_connector/unit/test_nixl_connector_hma.py -m cpu_test -q

uv run --no-project pre-commit run --files \
  vllm/distributed/kv_transfer/kv_connector/v1/nixl/metadata.py \
  vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py \
  vllm/distributed/kv_transfer/kv_connector/v1/nixl/pull_worker.py \
  tests/v1/kv_connector/unit/test_nixl_connector_hma.py

GPU follow-up: apply this fix alongside #57049 and #57027, restore block size 128, and repeat the same 100-question GSM8K sample with the same model, prompt/scoring settings, and concurrency 10. A full matched evaluation remains pending.

Test Result

• CPU suite: 91 passed, 1 deselected. Includes real receive-path page expansion and descriptor selection with transport mocked, producer/consumer padding, unequal per-region cache hits, and completion zeroing across strided layer views.
• Negative control: restoring main's mapping methods in memory makes the 160-token regression fail on the source descriptors: [21, 22, 23, 121, 122, 123] versus [20, 21, 22, 120, 121, 122].

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

[NickLucche]

/ci run

[github-actions]

✅ Triggered Buildkite CI #89191 for commit 6891a87ce839.

[LucasWilkinson]

should self.device_kv_caches

self.host_xfer_buffers if self.use_host_buffer else self.device_kv_caches

[LucasWilkinson]

Overall looks good except (other then maybe my comment above)