[Kernel] RPA: do the q token/kv-head swap on the 4D view to avoid a VMEM round trip

[shungcp]

What

prepare_inputs in ragged_paged_attention/v3/kernel.py reshapes q to 5D and
then swaps tokens with kv-heads. Doing the swap one step earlier — while q is
still 4D — produces a byte-identical operand for far less work.

# before
q.reshape(T, K, G // p, p, H).swapaxes(0, 1)
# after
q.swapaxes(0, 1).reshape(K, T, G // p, p, H)

One function; the kernel body, q_hbm_ref, input_output_aliases and
prepare_outputs are untouched.

Why it is faster

Splitting the q-head axis first fixes the two minor dims at
(q_packing, head_dim) = (2, 128). XLA cannot reach that layout from a
transpose directly, so it stages the whole tensor through VMEM: copy in,
bitcast, copy back out to HBM. On the 4D view the minor dims stay
(num_q_heads_per_kv_head, head_dim), which one straight HBM→HBM copy reaches,
and the trailing reshape is then a pure bitcast onto exactly the operand the
kernel already asks for.

The retile itself was never the cost: bf16[T,N,H]{T(8,128)} and
bf16[T,K,G/p,p,H]{T(2,128)} are byte-identical. What cost 130.8 µs/layer was
the permutation being forced through VMEM.

Numbers

One 2048-token prefill step on v6e-1, Qwen3-4B, total device self time:

dtype	before	after
int8 W8A8	30,202 µs	28,929 µs
bf16	36,406 µs	36,116 µs

The attention-path layout ops drop from 5.91 ms to 4.03 ms per int8 step. bf16
is not regressed.

End to end on a real server (Qwix online int8, ISL 2300 / OSL 10, prefix caching
off, fixed seed, concurrency 64):

model	TP	K per device	req/s	throughput
Qwen3-4B-Instruct-2507	1	8	25.28 → 26.23	+3.8%
Llama-3.1-8B-Instruct	1	8	21.68 → 22.31	+2.9%
Qwen3-32B	4	2	12.27 → 12.52	+2.1%
Llama-3.1-8B-Instruct	4	2	37.10 → 37.76	+1.8%

All five reported metrics move together in every row (e.g. Llama-3.1-8B at TP=1:
throughput +2.9%, TTFT −2.8%, P99 TTFT −2.8%, TPOT −2.8%). The spread tracks how
attention-heavy the model is and how many KV heads land on each chip. It is a
prefill-path change, so decode-heavy traffic sees proportionally less.

Who it helps

The gain needs G = num_q_heads // num_kv_heads to be a power of two and >= 4
(>= 8 when q itself is 8-bit, since q_packing goes 2→4) and
head_dim % 128 == 0. Both are per device — a K=8 model at TP=8 sees one KV
head per chip and the swap degenerates. Anything outside the set is an exact
no-op, measured inside a 0.3% noise floor: num_q_heads_per_kv_head is
align_to(..., q_packing), so the 4D minor dims are never smaller than the 5D
ones and a miss cannot regress.

Qualifying families:

• Llama 3.1 / 3.3, every size — 8B (G=4), 70B (G=8), 405B (G=16)
• the Mistral dense line — 7B v0.3, Nemo, Small 24B, Mixtral 8x7B,
  Magistral, Devstral (all G=4)
• Qwen3 from 4B up except the 14B — 4B/8B (G=4), 32B/30B-A3B (G=8),
  235B-A22B (G=16) — plus Qwen2.5-72B, phi-4 and Granite-3.3-8B

kernel_hd64.py has the same construct but is head_dim=64 by definition, so
it never meets the criteria; I left it alone rather than churn it for no gain.

Risk and testing

• Output is byte-identical by construction; greedy generations were verified
  token-for-token on 5 prompts, including one crossing the chunked-prefill
  boundary and one tile-unaligned.
• tests/kernels/ragged_paged_attention_kernel_v3_test.py: 67 passed.
• A pure reordering of two existing ops: no layout constraint, no new import,
  no configuration.

Follow-up, not in this PR

Pinning the producer's layout with
with_layout_constraint(q, Layout(major_to_minor=(0,1,2,3))) is worth a further
+4.8 points on Qwen3-4B when activations are int8: the fused
rope+activation-quant emits [K][G][T][H], so even the 4D copy interleaves G
through T, and pinning [T][K][G][H] makes it contiguous.

Excluded here because it only pays when there is such a producer to pin. End to
end on Qwen3-4B it is +4.5% with int8 activations, −1.2% on bf16 and −1.3% on
weight-only int8 — two independent lines agreeing, and matching the +449 µs
per step the bf16 single-step profile predicted. Where XLA has already found a
good global layout, overriding it just relocates copies. So it needs a
quantization-aware condition, which seemed worth separating from a change that
is unconditionally free.