[Bug]: NixlConnector silently misreads KV pages with split K/V slots (ROCm AITER shuffle, ROCM_ATTN, B12X)

Your current environment

Environment

I haven't reproduced this on GPUs yet, so there's no collect_env.py output to attach. I'll add it with the reproduction.

These findings come from reading main at 24c9772d19 (2026-09-27). The affected setups all use NixlConnector: ROCm with --attention-backend ROCM_AITER_FA and VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1, and the ROCM_ATTN and B12X backends.

Describe the bug

Summary. For heterogeneous TP, NixlConnector assumes each block of an attention KV page is laid out head-first, as [H, N, 2*D]. Its handshake also has no way to see how bytes are arranged inside a page. Some backends instead store K and V as two separate slots in the page. This causes two silent misreads:

1. A ROCm AITER engine using the shuffled layout can be paired with a standard engine. The decoder then reads scrambled K/V with no error.
2. For any backend with split K/V slots, heterogeneous TP splits each page into K and V rather than by head.

Background. With VLLM_ROCM_SHUFFLE_KV_CACHE_LAYOUT=1 (#29887), ROCM_AITER_FA writes each block's K as [H, D/x, N, x] and its V as [H, N/x, D, x], where x = 16 // element_size (writer). Since #52849, the page is published as two head slots: K for all heads, then V for all heads (customize_spec). ROCM_ATTN (spec) and B12X (spec) publish two slots the same way.

1. NIXL treats shuffled pages as standard pages. Nothing NIXL compares tells them apart:

• get_name() returns "FLASH_ATTN" for shuffled AITER, unshuffled AITER and CUDA FlashAttention (AITER, CUDA).
• All three resolve to the LBHNC layout that NIXL requests. They also have the same page size, block_lens and block_strides.
• No compatibility-hash factor reflects the shuffle setting.
• vllm_version is the same for CUDA and ROCm builds of one commit, because vllm/_version.py is written before the +cu/+rocm suffix is added (setup.py).

Whether a mixed pair fails today depends on details unrelated to the page format:

Pair	Result on main
A shuffled engine on the gluon decode path (head size 128, GQA group 8 or 16, 128-token pages) with a CUDA FlashAttention engine using 128-token pages, in either direction	The handshake passes, and the decoder silently reads scrambled K/V
A shuffled decoder with 16- or 32-token pages at LBHNC	Fails: the paged_attention_common assert requires LHBNC
A shuffled engine at LHBNC with a standard engine at LBHNC	Fails: layout mismatch at handshake
A shuffled prefill with 128-token pages and an unshuffled AITER decoder (which supports only 16- and 32-token pages)	Fails: the prefill block is larger than the decode block

The failures come from other checks; none of them compares page formats. The silent row is the AMD–NVIDIA P/D case on the #33702 roadmap. One case I haven't checked: a shuffled prefill with 16- or 32-token pages that never runs a decode batch (for example, with CUDA graphs disabled), paired with a standard decoder.

2. Heterogeneous TP splits pages with split K/V slots by K/V, not by head. NIXL gives each rank an equal contiguous byte range of every block (remote, local). For an [H, N, 2*D] block, that range is a set of heads. A block with split slots at LBHNC is [K of all heads | V of all heads], so the ranges don't line up with heads:

• Decode TP = 2 × prefill TP: one decode rank receives only K, and the other receives only V. Half of each rank's KV never arrives.
• Prefill TP = 2 × decode TP: all the bytes arrive, but only half of them land where the decoder reads them.

The handshake accepts both cases. The block-contiguity check passes, and the per-rank block_len scales with the head ratio as expected. Exposure on main:

• shuffled AITER in the gluon configuration;
• ROCM_ATTN or B12X when selected explicitly. ROCm already leaves ROCM_ATTN out of automatic backend selection when a KV connector is used (rocm.py).

Proposed fix. I'm happy to send PRs.

1. Check the page format at handshake.
   • Have the backend supply a page-format tag, advertise it per region in NixlAgentMetadata, and check it in _validate_remote_agent_handshake next to kv_cache_layout.
   • Use a metadata field rather than a hash factor. Properties validated at runtime are deliberately kept out of the hash (metadata.py). A field also lets the error name the cause, and leaves room for conversion later.
   • Treat a missing tag as unknown, and reject it when the local pages aren't standard. Bump NIXL_CONNECTOR_VERSION.
   • In push mode, the prefill runs this check. I'd also check in _add_notif_only_remote_agent so the decoder fails fast.
2. Reject heterogeneous TP for split K/V slots, then fix it. Reject non-replicated heterogeneous TP when either side has num_head_slots == 2, for all backends. Then fix it by giving each rank its share of the K slot and its share of the V slot separately. This works when each slot is head-major, as in AITER shuffle and ROCM_ATTN. In simulation, it fixes both TP directions.

Related: #54805. It makes the shuffled layout prefer LHBNC but keeps LBHNC supported. NixlConnector's LBHNC request wins layout resolution whenever the backend supports it (utils.py). So after that PR, shuffled engines using NixlConnector would still hit the assert above unless VLLM_KV_CACHE_LAYOUT=LHBNC is set.

Follow-up. #59111 proposes converting instead of rejecting for one pairing: a standard prefill feeding a shuffled ROCm decoder. It builds on both fixes above.

How this was found. I read the code at 24c9772d19 and ran pure-Python simulations of the shuffle writer's offsets and NIXL's descriptor math. I haven't reproduced it on GPUs yet and will add the results here.

Related issues: #42082 and #47095 (KV layout RFCs; #47095 notes that AITER's layouts aren't permutations), #33702.

cc @NickLucche @xuechendi @tjtanaa @ganyi1996ppo @ukannika @LucasWilkinson

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.