[zixi-qi]

Purpose

Enable attention HMA with pipeline-parallel prefill in NixlPushConnector,
following #45880. Producer PP stages and a PP1 decoder may pool the same layers
into different physical KV-cache regions. This PR matches layers by name and
builds descriptors using each side's region geometry.

This updates the existing push-mode PR. Pull-mode work in #43368/#48263 and
the earlier #47981 prototype use different implementations. Packed MLA is the
stacked follow-up #50499, which should land after this PR.

Design and rebase

• NixlAgentMetadata.region_members advertises layer names; wire version 11 → 12.
• Registration derives _transfer_layer_names, _transfer_layer_region_indices,
  and _transfer_layer_group_ids per worker. _supports_pp_hma expresses the
  connector capability; _requires_layer_name_routing() applies runtime gates.
• Remote layer alignment selects this PP stage's layers and reorders addresses,
  lengths, strides, capacities, group IDs, names, and memory types. Missing or
  duplicate layers and inconsistent metadata fail explicitly.
• Descriptors use cumulative per-region capacities; physical allocations are
  registered once. PP1 and non-HMA layouts retain region-index routing.
• Unequal block sizes and unsupported non-MLA TP expansion are rejected before
  topology/NIXL registration.

Rebased onto main@df42d112ee88dd4a9b64efbad55621af6a66a44b on September 15.
The rebase preserves upstream PCP/DCP transfer ranks, HiSparse geometry,
KVConnectorTransferResults, and the latest NIXL multi-handle failure handling.

The conflict was in receive-lifecycle tests. Main's fixtures and regressions
supersede the earlier PR versions, so that file now matches main. The PR's
production-code additions/removals are unchanged from the reviewed version.
The three PP startup restrictions remain grouped under one PP-size guard.

Main now completes producer sends on consumer acknowledgement or lease expiry.
The follow-up replaces the obsolete local-WRITE completion expectation with
PP1 and PP2/HMA cases in upstream's notification test, preserving partial-handle
polling and exactly-once lease/eviction checks. No extra production changes.

Test plan

Run from the repository's .venv on one node with four GB200s:

CUDA_VISIBLE_DEVICES=3 VLLM_NIXL_SIDE_CHANNEL_PORT=5700 \
  .venv/bin/python -m pytest -q tests/v1/kv_connector/unit/test_nixl*.py
CUDA_VISIBLE_DEVICES=3 VLLM_NIXL_SIDE_CHANNEL_PORT=5700 \
  .venv/bin/python -m pytest -q \
  tests/v1/kv_connector/unit/test_kv_load_failure_recovery.py \
  tests/v1/kv_connector/unit/test_output_aggregator.py \
  tests/v1/kv_connector/unit/test_tp_mapping.py
.venv/bin/pre-commit run --files $(git diff --name-only origin/main...HEAD)
.venv/bin/pre-commit run mypy-3.12 --hook-stage manual \
  --files $(git diff --name-only origin/main...HEAD)

Fresh Gemma-4-26B-A4B-it E2E uses PP2/TP1 → PP1/TP1, HMA enabled, prefix
caching disabled, strict handshake compatibility, eager execution, explicit
--attention-backend TRITON_ATTN on both peers, bf16 KV, V2 runner, cumem
allocation, max length 8192, and UCX_TLS=tcp,cuda_copy.

GSM8K: 200 questions per route, 5-shot chat, concurrency 8, temperature 0,
2048 generation tokens. Direct producer baseline and toy-proxy disaggregation
use the same checkpoint 4d7ae4984b7db7de8f8457170b3f1a419ee76d52.

lm_eval --model local-chat-completions --tasks gsm8k --num_fewshot 5 \
  --apply_chat_template --fewshot_as_multiturn --limit 200 --log_samples \
  --gen_kwargs max_gen_toks=2048 --output_path "$OUTPUT_DIR" \
  --model_args "model=$MODEL,base_url=http://127.0.0.1:$PORT/v1/chat/completions,num_concurrent=8,max_retries=3,tokenized_requests=False"

Use port 8100 for the baseline and 8000 for the proxy. Probes are recorded
separately; transfer/cache metrics and health are checked after 35 seconds idle.

Test results

September 15 validation on b74b0866cea9d3969a8f34a67caca0c68c9af028:

• 562 applicable unit cases passed; 2 ROCm-only skips. Validation combines
  the full NIXL/supporting-suite run with a rerun of the entire push-connector
  test file after the test-only update: 484 unaffected cases passed in the
  full run, and all 78 final push tests passed (19.66 seconds). The initial
  run exposed the two obsolete WRITE-completion assertions described above.
• Changed-file pre-commit and manual mypy Python 3.12 passed, including the
  test follow-up. No production files changed after the full test run.
• E2E was not rerun for this rebase. The full model results below belong to
  the earlier tested commit, rather than the new head.

Validated September 14, 2026 on 5c12ce79ada8246f341be97414b3df60d147b818:

• 502 NIXL tests passed, 2 ROCm-only skips, 448.02 seconds, 17 warnings.
• 49 scheduler/failure-recovery/aggregation/TP-mapping tests passed,
  23.94 seconds. Pre-commit and manual mypy Python 3.12 passed.
• All 400 GSM8K evaluation requests completed. All 200 document, prompt,
  and target hashes match between routes.

Route	Strict match	Flexible extraction
Direct producer baseline	95.0%	95.0%
PP2 push-disaggregated	96.0%	95.5%

Exact generated-text parity and code-caused accuracy improvements are not
claimed. The flexible regression on question 174 is extraction of 35 from the
correct 95-minute answer expressed as “1 hour and 35 minutes.”

• 402 completed WRITEs / 45,439,385,600 bytes, including the probe.
• Zero failed transfers, failed notifications, expired leases, or serving
  errors, including the idle audit. All completion HTTP responses were 200.
• Decoder evaluation-only counters: 217,085 external hits / 217,085 queries,
  zero locally computed prefill tokens. Its separate probe contributes one
  computed token.
• Strict metadata audit passed: 15+15 producer layers, 30 decoder layers,
  block size 16; capacities 383,323 per producer region and 191,976 per decoder
  region. Runtime source hashes match the published source.

Environment: torch 2.13.0+cu130, NIXL 1.1.0, FlashInfer Python/cubin/JIT cache
0.6.18.post1. Compiled artifacts are from dc36fcce902a63eab06c1b93a5c4a5ee178a0c56,
the immediate parent of the selected main; their only difference is a Python
frontend relocation, with identical compiled sources and build dependencies.

Limitations

• Gemma validation explicitly uses TRITON_ATTN. The new auto-selected
  TRITON_FLASHINFER composite failed cache registration before any inference
  requests; this rebase does not add support for that layout.
• Equal P/D block sizes are required. Non-MLA decode TP greater than prefill TP,
  pull-mode PP/HMA, decode-side PP, and PP with Mamba/SSM remain unsupported.
• Packed MLA belongs to [KVConnector][NIXL] Support packed MLA KV layouts in pipeline-parallel push prefill #50499. These runs do not establish DP/EP/CP E2E or
  RDMA performance coverage, and do not include live fault injection.
• Local passes do not establish required upstream GPU CI or human approval.

AI assistance

AI assistance (Claude and Codex) was used for implementation, review, and
validation. The human submitter retains responsibility for reviewing and
understanding every changed line; maintainers retain merge responsibility.

Co-authored-by: Yiliu Dong https://github.com/qianlihuang

[mergify]

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

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

[mergify]

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

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

[mergify]

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

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

[mergify]

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

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

[github-actions]

✅ Queued 7 failed job(s) for retry in Buildkite CI #88958.

[mergify]

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

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

[zixi-qi]

/ci run

[github-actions]

❌ This PR is 4 commits behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[zixi-qi]

/ci run

[github-actions]

❌ This PR is 1 commit behind upstream main. Your branch must contain every commit currently on upstream main. No new CI build was started. Merge or rebase onto the latest main, then rerun /ci run. To test this branch at your own risk, use /ci run --allow-stale.

[zixi-qi]

/ci run --allow-stale

[github-actions]

✅ Triggered Buildkite CI #89186 for commit 443768d04cdc.

⚠️ This PR is 1 commit behind upstream main. Running CI at your own risk because --allow-stale was requested; outdated CI configuration may cause failures. Before merging, merge or rebase onto the latest main, then rerun /ci run on the latest PR commit.