[RFC]: NixlConnector: convert standard KV pages to ROCm AITER shuffled pages on the decode side

Motivation.

TL;DR

• On ROCm, AITER's shuffled KV pages give the fastest decode at high batch. Prefill doesn't benefit from them, and a different backend is the fastest for prefill.
• P/D disaggregation could give each side its best format, but NixlConnector requires both sides to use the same format.
• I propose letting a ROCm decoder accept standard pages and convert them to shuffled pages after the transfer.
• Decisions I'm asking for now:
  • whether the connector should compare a page-format tag and call a converter owned by the backend;
  • whether prefill and decode may use different backends once their page formats are negotiated;
  • whether a conversion on the decoder's step path is acceptable.
• The measurement that decides whether to merge the conversion comes after this discussion (see the go/no-go plan).

Depends on #59110 (bug: NixlConnector can't see the KV page format). That issue's two fixes are:

• a per-region page-format tag in NixlAgentMetadata;
• splitting the K and V slots separately in NIXL descriptors.

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. AITER's asm and gluon decode kernels read this format. The other backends write standard pages, where each head's tokens hold K and V side by side ([H, N, 2*D]). These are CUDA FlashAttention, unshuffled ROCM_AITER_FA and ROCM_AITER_UNIFIED_ATTN. Neither kernel family reads the other's pages.

Decode: shuffled pages win at high batch. I timed full decode steps on one MI300X (Qwen3-8B, BF16, TP 1, 16-token pages). The shuffled arm used the non-gluon asm/ll4mi path, with VLLM_KV_CACHE_LAYOUT=LHBNC. Each cell is the shuffled speedup over ROCM_AITER_FA, the backend ROCm selects by default with a KV connector. The number in parentheses is the speedup over the fastest unshuffled backend at that shape (TRITON_ATTN, ROCM_AITER_UNIFIED_ATTN or ROCM_AITER_FA):

A dash marks a shape that exceeded the harness's 95% KV-capacity margin. Each configuration ran once, so treat gaps under about 2% as ties.

• Against ROCM_AITER_FA, the gain rises with batch once AITER's asm decode kernel takes over. For Qwen3-8B on MI300X that happens above batch 19, when batch × query heads exceeds twice the CU count.
• Against the fastest unshuffled backend, shuffled pages are 2.3–6.0% faster at 6 of 7 shapes with batch ≥ 32. At batch 128 the margin narrows to +3.6% at 1k and −1.5% at 4k.
• Across all 15 shapes, the median gain over the fastest unshuffled backend is +1.3%.

Kernel-level evidence comes from profiler traces at batch 64 and 8k context. There, shuffled (asm) attention took 10.6% less time than ROCM_AITER_FA's ll4mi kernel and 4.4% less than AITER unified, measured as time reduction over 12 steps. Below the asm threshold there is no consistent attention gain.

Qwen3-32B's shuffled decode ran on the gluon path (128-token pages). There it was roughly a wash: a median of −0.6% against the fastest unshuffled backend, ranging from −2.3% to +3.8% at batch ≥ 32. #54805 (open) would remove gluon until a follow-up restores it.

Separately, #29887 reported +19.4% output throughput (+26.5% TPOT) against unshuffled ROCM_AITER_FA for Qwen3-30B-A3B-FP8. That was on MI308 with TP 2 + EP, in a colocated benchmark at concurrency 200 with about 3.8k input and 1k output tokens. MI308 has fewer CUs, so asm engages at a smaller batch, and a small-active-parameter MoE spends more of each step in attention. Either may explain part of the difference, but I haven't measured it.

My hypothesis is that the gain holds or grows at the larger batch × context shapes the table doesn't reach. P/D decode instances run in that regime. The go/no-go run below tests this.

Prefill: shuffling doesn't help. I measured prompt extension after prefix reuse, the prefill path that reads cached context, over 63 shapes on Qwen3-8B and Qwen3-32B:

• Shuffled pages were never the fastest.
• With 16-token pages, shuffled ROCM_AITER_FA was a median 2.8–3.5% slower than unshuffled ROCM_AITER_FA (up to 12.7%). The cost is in gathering the cached context.
• The fastest unshuffled backend was ROCM_AITER_UNIFIED_ATTN (64-token pages) at 44 of 63 shapes.
• Cold prefill was not measured.

Why not just shuffle both sides? Shuffling both sides gets shuffled decode without any conversion, but it has three costs:

1. Prefill pays the shuffled-page penalty above, about 3% of extension time.
2. Prefill must stay on ROCM_AITER_FA, the same backend as the decoder, so it can't use AITER unified. That is also true under this proposal unless prefill and decode may use different backends. attn_backend_name is a compatibility-hash factor, so a ROCM_AITER_UNIFIED_ATTN prefill can't pair with a ROCM_AITER_FA decoder today, even though both write standard pages. This is decision 2 below, and the go/no-go includes a measurement of it.
3. An NVIDIA prefill can't write shuffled pages (the AMD–NVIDIA item on #33702). There, conversion is the only way to use shuffled decode.

What blocks it today. NixlConnector's receive-side conversions are:

• the HND→NHD and block-size permutes (#26556, #30275);
• a CPU_ATTN repack (#38935) keyed on the backend name, which re-runs the CPU cache writer on received standard pages.

None is negotiated at handshake. NixlConnector can't tell a convertible pair from an incompatible one: shuffled LBHNC pages even pass the handshake as standard pages (#59110).

Proposed Change.

Scope of the first version:

• standard → shuffled only, on a ROCm decoder with 16- or 32-token pages at LHBNC (the non-gluon path);
• homogeneous TP, BF16/FP16, non-HMA, non-MLA, equal block sizes, pull mode.

Everything else fails closed:

• the reverse direction;
• gluon-configured decoders;
• FP8, where the shuffle writer's scale handling and MI300's e4m3fnuz vs OCP e4m3fn storage need their own treatment;
• push mode, for now.

I'm happy to implement it.

1. Negotiate the pair. Make the decision in the kv_cache_layout check of _validate_remote_agent_handshake, per remote engine. It uses the page-format tag from #59110. A remote with standard pages at LBHNC and a local shuffled decoder at LHBNC is accepted only when:

• conversion is enabled;
• the pair is in scope.

Any other mismatch is rejected. The existing enable_permute_local_kv branch will also require the local layout to be LBNHC, so that it can never permute shuffled caches. A missing tag follows the #59110 rule: it's rejected whenever the local pages aren't standard.

2. Layout. The non-gluon shuffled path already requires LHBNC on main: it asserts that the K and V sides are dense planes. NixlConnector, however, requests LBHNC, and that request wins layout resolution. So the decoder needs VLLM_KV_CACHE_LAYOUT=LHBNC today. I'd change get_required_kvcache_layout to defer to the backend when the shuffle layout is on.

3. Descriptors.

• Under LHBNC, the decoder registers a K plane and a V plane per layer, each half a page. A standard LBHNC prefill registers one full-page region per layer.
• Present the remote side as two virtual regions per layer, at offsets 0 and block_len/2, each of length block_len/2. They are ordered like the local K-plane and V-plane regions.
• The contract with the converter: after the read, the local K-plane block followed by the V-plane block holds the remote standard block's bytes. These are the block's first and second byte halves, not K and V.
• Code changes:
  • _build_fa_remote;
  • the remote region metadata recorded in add_remote_agent;
  • the remote _compute_desc_ids call in pull_worker.read_blocks;
  • the region count, region names and block_len checks;
  • tests in test_nixl_desc_geometry.py.
• This is the same descriptor machinery as the K/V-slot split in #59110.

4. Convert, in a hook owned by the backend.

• The connector only compares tags and calls a backend hook with the caches, the received block ids and the remote tag. The hook could be a classmethod on the AITER FA backend or live in _aiter_ops. This follows #38935, where Platform.pack_kv_cache does the CPU_ATTN repack.
• The simplest correct implementation gathers the received K/V and re-runs the backend's own writer, reshape_and_cache_shuffle_triton, with a synthetic slot_mapping for those blocks. That matches the writer by construction and would carry over to a later FP8 path. A per-block view/permute version serves as the test reference.
• It runs in the receive post-process in get_transfer_results, one layer at a time. That is inside the decoder's execute_model, on the main stream, so each arriving request delays every running request's next step. If that shows up in ITL, the mitigation is a side stream plus an event that must complete before the request is scheduled.

# One block, standard (H, N, 2*D) -> shuffled; x = 16 // element_size
K, V = std[..., :D], std[..., D:]
k_shuf = K.reshape(H, N, D // x, x).permute(0, 2, 1, 3)   # (H, D/x, N, x)
v_shuf = V.reshape(H, N // x, x, D).permute(0, 1, 3, 2)   # (H, N/x, D, x)

• Data per request. A Qwen3-8B request carries 0.6 GB of BF16 KV at 4k context and 2.4 GB at 16k.
• Conversion time. Generic torch ops move at least about 4× those bytes. At 1.5–3 TB/s effective, that is roughly 1–2.5 ms per 4k request and 3–10 ms per 16k request.
• Decode savings. For comparison, shuffled decode saves 0.6–2.7 ms per step at batch 32–64 against ROCM_AITER_FA, and 0.4–1.3 ms against the fastest unshuffled backend.
• Break-even. That is roughly 40–150 output tokens against ROCM_AITER_FA, and 70–360 against the fastest unshuffled backend (4k–16k context, batch 32–64).

5. Knob. Off by default. Because it is NIXL-specific, I'd make it a kv_connector_extra_config key like enforce_handshake_compat. (enable_permute_local_kv, by contrast, is a top-level KVTransferConfig field.) The key would be named after the #59110 tag and read by the decoder. It would come with an entry in the NIXL compatibility docs.

Go/no-go measurement (before the conversion PR is reviewed). Conversion adds decode-side work, so it can't beat "shuffle both sides" on TPOT. Its possible wins are on the prefill side (TTFT and prefill throughput) and in pairings that shuffling both sides can't provide.

• Setups:
  • (a1) unshuffled FA → FA;
  • (a2) unshuffled UNIFIED → UNIFIED;
  • (b) shuffled FA → FA;
  • (c) unshuffled FA → shuffled FA, with conversion;
  • (d) UNIFIED (16-token pages) → shuffled FA, with conversion. This one is prototyped with "enforce_handshake_compat": false because of the backend-name hash.
• Configuration:
  • MI300X, decoder at TP 2 so that 32k context fits.
  • Batches above the asm threshold for that TP (48, 64, 128), with a profiler check of which decode kernel ran.
  • Models: Qwen3-8B, and Qwen3-30B-A3B-FP8 weights with BF16 KV.
  • Contexts 4k–32k; output lengths 128, 512 and 1024.
  • Both cold-prefill and prefix-hit prefill workloads.
  • At least 3 repeats per configuration.
• Metrics:
  • TTFT, split into prefill, transfer and conversion time;
  • prefill throughput;
  • TPOT;
  • p50/p99 ITL of other running requests;
  • goodput at fixed TTFT/TPOT SLOs.
• Decision rules:
  • (c) and (d) vs (a1)/(a2): better TPOT or output throughput at equal TTFT.
  • (c) and (d) vs (b): better TTFT or prefill throughput, with TPOT no more than 2% worse.
  • A result counts only if it is ≥ 3% with non-overlapping repeat ranges.
• An NVIDIA → AMD run would cover the other pairing that shuffling both sides can't. I'll add it if a mixed setup is available.

Test plan.

• CPU round-trip tests of the converter against a reference that reproduces the shuffle writer's offsets. These include the two byte-halves LHBNC read and single-head and odd-head ranks.
• Descriptor-geometry tests for the split read.
• Handshake tests: the convertible pair is accepted only when enabled, and every other mismatch, including the reverse direction and the enable_permute_local_kv branch, is rejected.
• A ROCm 1P1D accuracy run (tests/v1/kv_connector/nixl_integration, DECODER_KV_LAYOUT=LHBNC) with shuffle on the decoder only.

Compatibility. Pairs with matching page formats see no change. Conversion runs only when enabled and the pair is in scope. Older peers already fail the compatibility hash unless that check is disabled.

Decisions for reviewers.

1. Where the format and converter live. I'd propose an opaque page-format tag from an AttentionBackend hook, which the connector only compares, plus a converter hook owned by the backend. This keeps NIXL free of AITER's tile layout, in line with #42082 keeping the page content opaque (#47095 discusses layouts that aren't permutations). @LucasWilkinson @heheda12345
2. Different backends on each side. Once page formats are negotiated, should prefill and decode be allowed to use different backends when their formats match or are convertible? Examples are AITER unified → AITER FA, or FLASH_ATTN → CPU_ATTN, which today needs enforce_handshake_compat: false. Setup (d) measures what this is worth. NIXL owners.
3. Converter implementation. Should the converter re-run the backend writer in torch, or use a dedicated AITER kernel? AITER owners.

Feedback Period.

About one week.

CC List.

@NickLucche @xuechendi (NIXL), @tjtanaa @ganyi1996ppo @ukannika (AITER, shuffle layout), @LucasWilkinson @heheda12345 (KV cache layout)

Any Other Things.

Pilot setup.

• Hardware and software: one MI300X (gfx942), vLLM 0.30.0 ROCm image (AITER 0.1.21.post2, ROCm 7.2.3), VLLM_ROCM_USE_AITER=1 in every compared configuration, BF16 weights and KV, TP 1.
• Decode timing: median of 96 steps at a constant batch, CUDA graphs FULL_AND_PIECEWISE.
• Excluded baselines:
  • ROCM_ATTN, because ROCm doesn't select it automatically with a KV connector and it is affected by #59110. Including it narrows the batch ≥ 32 range to +2.3% to +5.1%.
  • Patched, unsupported NHD configurations of AITER.
• Restricting the baselines to LBHNC-capable backends (NIXL's default) leaves both ranges unchanged.

Related: #59110, #54805, #33702, #38935, #42082, #47095, #26556, #30275.

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.