Fix attention_block_tkg d_head>128 decode Q/K layout corruption for batch>1

[jimburtoft]

PR: Fix attention_block_tkg d_head>128 (n_d_tiles>1) decode Q/K layout corruption for batch > 1

Summary

The token-generation decode path in attention_block_tkg produces numerically wrong output for
d_head > 128 (n_d_tiles > 1) whenever the decode batch size B >= 2. B=1 is correct, and
d_head <= 128 is correct at any batch. The bug is a layout-ordering mismatch between the d-tiled
Q/K producer and the QK-matmul consumer that is masked at B=1 (singleton batch dim) and therefore
invisible to every existing d256 unit-test config (all batch=1).

Root cause

_process_head_group (in attention_block_tkg.py), n_d_tiles > 1 branch, writes the tiled Q/K
SBUF buffer's free dimension head-major:

dst_offset = i_d * (B * n_heads * S) + head_idx * (B * S)   # [n_d_tiles][n_heads][B][S]
nisa.tensor_copy(out[:, nl.ds(dst_offset, B * S)], psum)     # psum: [pmax, B*S]

But attention_tkg._compute_qk_matmul reads the same buffer batch-major for the d>128 QK^T
moving operand:

q_sb_d_offset = i_d * atp.bs_full * atp.s_active_qh          # [n_d_tiles][bs_full][s_active_qh]
q_sb_view = q_sb.slice(1, q_sb_d_offset, +bs_full*s_active_qh)
                 .reshape_dim(1, [bs_full, s_active_qh]).select(1, global_batch_offset + i_b)

i.e. [n_d_tiles][B][q_heads][S]. The d_head <= 128 path writes dst_4d = [d, B, n_heads, S]
(batch-major), matching the consumer — hence d128 is always correct. At B=1 head-major and batch-major
coincide (the batch dim is 1), so the mismatch is invisible; at B >= 2 the Q heads are permuted relative
to the batch dimension, corrupting every batch row's attention output by ~8% (cos ~0.92; both d-tiles
equally affected because QK^T sums over d).

Fix

Write the d-tiled producer batch-major so it matches both the consumer and the d128 path:

for i_batch in range(B):
    dst_offset = i_d * (B * n_heads * S) + i_batch * (n_heads * S) + head_idx * S
    nisa.tensor_copy(out[:, nl.ds(dst_offset, S)], psum[:, nl.ds(i_batch * S, S)])

B * S <= pmax on this path, so B and S are small and the extra per-batch copies are cheap.

Validation (trn2, SDK 2.32, neuronx-cc 2.27, vllm-neuron 0.24)

Standalone NF.attention_decode d256 (q=8, kv=4, update_cache=False, W_out=None), per-batch dense-fp32
CPU golden:

Config	Before	After
d256 B=1	cos 0.999997 (pass)	cos 0.999997 (pass)
d256 B=2	cos 0.92 (FAIL)	cos 0.999997 (pass)
d256 B=3/4/8	FAIL	pass
d128 B=2 (control)	pass	pass

End-to-end (Gemma3-4B / Sarvam-Translate, TP1): greedy-exact 0/12 → 9/12; TP4: 1/12 → 9/12.

Test coverage added

Three batch >= 2 d256 configs added to TestRangeAttnBlk.FAST_ATTN_BLK_CFGS in
test/integration/nkilib/experimental/transformer/test_attention_block_tkg.py (batch=2 q=2/kv=1;
batch=2 q=8/kv=4; batch=4 q=8/kv=4). They fail on the unpatched kernel and pass after the fix. The
pre-existing d256 configs were all batch=1, which was the coverage gap.

Note on the K producer

The same _process_head_group branch also produces the tiled K (K_sb); the fix applies to it
uniformly. The returned-K store (K_tkg_hbm reassembly) reads whole per-d-tile slices, so the
batch-major reordering is consistent through the K path as well.