[AbrahamCete]

What this changes

tf_plan_score_* (zig/kernels/hip/decode/plan.hip) is the planned decode walk's Q·K: one score per key for each query head of a KV group. It took one key a warp:

• eight two-byte loads a lane;
• the group's queries read again from LDS for every key;
• a ds_bpermute wave sum for every head.

On long contexts it read the cache at about 72 GB/s and was most of a full-attention layer's decode time. On gfx1151 at 64K tokens, Qwen3.6-35B-A3B spent 921 µs a layer on it (2 KV heads × 33.5 MB of keys), against 388 µs for the values' pass over the same bytes.

The changes, one commit each:

• test(hip): pin the planned decode score's bits to its lane order
  • The planned walk test compared the walk's output with the flat walk and float64, so nothing pinned the score kernel's own bits.
  • It now checks every score, for fp32 and activation-type queries, against a host emulation of the kernel's order: lane l sums q·k over l, l + 32, … with fp32 multiplies and adds, the wave folds its lanes by xor from 16 down to 1, then the sum is scaled.
• perf(hip): the planned decode score takes four keys a warp on long spans, two on short ones
  • A warp scores four contiguous keys from 16384 keys on, and two below (tf_plan_score2_*), where four would leave the card short of warps.
  • The warp's key loads go out together. The page table and descriptor are read once. Each head's query slice is read from LDS once for the warp's keys.
  • The wave sum folds by DPP row_xmask and permlanex16 in wave_sum's order, xor 16 down to 1. After the last fold every lane holds the sum, so the broadcast is gone.
  • The two- and four-key kernels are separate symbols. One kernel with both bodies behind a grid-uniform branch cost the four-key path 3-4% (32K: 335 µs alone, 343-351 µs behind the branch, in a same-session A/B).
  • The test now also runs 16384-key spans at d = 256 and d = 128, ending mid-warp, so both kernels are pinned. The planned passes write NaN-filled score buffers of their own, so a key the kernel skips can't pass for the flat walk's bits.
  • Records: tf_plan_score_* move from 17 to 48 VGPRs. The four tf_plan_score2_* records (31 VGPRs) live in hip_records_score.zig, chained from hip_records.zig as the prompt tile's are, since hip_records.zig would pass 600 lines.

Every score keeps its lane mapping, its in-lane order and its folds, so the bits are the same.

Score kernel time by keys a warp (Qwen3.6-35B-A3B, rocprofv3, the launches at the context's span):

span	1 key (main)	2 keys	4 keys
2K	42.0 µs	39.3 µs	48.2 µs
8K	—	110.9 µs	123.7 µs
16K	—	200.4 µs	200.3 µs
32K	467.6 µs	375.1 µs	349.6 µs
64K	921.3 µs	727.0 µs	651.9 µs

The 2- and 4-key columns are from the one-kernel build with both bodies; the shipped four-key kernel alone runs 330 µs at 32K and 630 µs at 64K.

Overlap with open pull requests: none touch plan.hip or launches/plan.zig.

• hip: a 64-row routed prompt tile for many-expert plans on gfx1151 #1658 and hip flashnext: sparse attention with a lane a key over several heads, the scalar kernel's bits #1877 add entries to hip_records.zig. hip: a 64-row routed prompt tile for many-expert plans on gfx1151 #1658 chains its own records file from hip_records_pair.zig, not from the line this PR changes.
• perf(hip): narrowest stream columns at RDNA3.5's one- and two-row launches #1855 and hip flashnext: sparse attention with a lane a key over several heads, the scalar kernel's bits #1877 add lines to launches.zig.
• This PR changes four records in place, adds the hip_records_score.zig chain, and adds two lines to launches.zig.

Receipt

• Environment: Ryzen AI Max+ 395 (gfx1151), Ubuntu 26.04, Zig 0.17.0, hipcc with HIP 7.1.52801 and Ubuntu clang 21.1.8 (ROCm 7.13 from TheRock for the contract limits and rocprofv3); checkpoints Qwen3.6-35B-A3B-MLX-4bit-MTP and Qwen3.8-27B-MLX-4bit.

• Base: the branch is two commits on main 21c8a0f. The receipts below ran on 9e8c368; none of the 105 commits between them touches these files. On 21c8a0f:

  • 64-step decode logits on both checkpoints and the 35B's replies after a 32K prompt are byte-identical to main's.
  • hip-kernel-test passes 109 with 1 skipped. hip-host-test, test-qwen35-hip-host and check-native-hip-host exit 0.
  • lean_check and the ROCm 7.13 contract limits give the same counts as main (424 and 220).

• Decode speed after a long prompt, served. tools/bench_cache.py --require-cold --tokens 128 on a 32K and a 64K prompt; one fresh server a run (--context 69632 --parallel 1 --prompt-cache-gib 0), alternated B1 C1 B2 C2 with B = main and C = this PR. The 35B serves with its MTP drafts.

  checkpoint, prompt	main	this PR
  35B-A3B, 32K	40.24 / 40.25	45.38 / 45.68 (+13.1%)
  35B-A3B, 64K	25.81 / 26.47	31.12 / 31.04 (+18.9%)
  27B, 32K	9.90 / 9.88	10.64 / 10.63 (+7.5%)
  27B, 64K	7.36 / 7.35	8.08 / 8.08 (+9.9%)

  • Short prompts, tools/bench_openai.py --tokens 256 --reps 3 --temperatures 0: 35B code 112.6 / 113.3 vs 113.0 / 113.0, chat 79.6 / 79.5 vs 81.6 / 81.4; 27B code 13.4 / 13.3 vs 13.3 / 13.3, chat 13.6 / 13.7 vs 13.7 / 13.6.
  • tf-qwen35-test check --only speed, twice each on 35B: one stream drafts off 70.4 / 70.4 vs 70.3 / 70.2, drafts on 95.8 / 95.6 vs 95.8 / 95.9; four streams drafts off 138.9 / 139.1 vs 139.1 / 138.8, drafts on 127.3 / 126.7 vs 126.1 / 127.4. 27B the same within 0.1.
  • Score kernel, tf-qwen35-test lanes with drafts off: 35B 42.0 → 39.8 µs at 2K, 467.6 → 330.1 at 32K, 921.3 → 630.1 at 64K; 27B 64.2 → 58.2, 822.4 → 562.9, 1634.5 → 1099.9.

• Exactness:

  • Token hashes are equal on every server of every run: 35B 61da022c at 32K, aaa410bc at 64K; 27B 28bb6ebb and 63bf1db1.
  • lanes replies are equal to main at 2K, 32K and 64K on both checkpoints.
  • 64-step decode logits are byte-identical on both.
  • tf-qwen35-test check --only invariants: 35B 46/0/1 and 27B 42/0/5, the same as main.

• Prompt speed: ttft_s of the same cold prompts is unchanged: 35B 27.5 / 27.4 vs 27.4 / 27.4 s at 32K and 71.2 / 71.2 vs 71.1 / 71.3 s at 64K; 27B 103.5 / 103.2 vs 104.2 / 103.7 s and 246.1 / 245.4 vs 247.1 / 245.6 s. The prompt path doesn't use these kernels.

• Kernel contract:

  • hip-kernel-contract on gfx1151: SUMMARY gfx1151 52 kernels: 52 pass gates 1-4, 0 fail, 2 known findings; 5 controls: 5 red as planted, 0 missed; 134 limit problems. main has 130.
  • The 4 extra are the new records' live counts under this machine's hipcc (30 VGPRs against 31 recorded from ROCm 7.13), the same one-VGPR gap as the existing score records and the rest of main's 130.
  • The planned walk's planCausal and planCausalAs cases pass every gate.
  • hip-kernel-contract-limits with ROCm 7.13: 220 problems on both, with every planned-walk score record current.

• Tests run:

  • The planned walk test with the score check: on the unchanged kernel it passes. It fails when the folds run 1 to 16, when only the four-key kernel's partial sums are scaled, and when the host launches one block too few (Nemotron's attention test fails there too).
  • hip-kernel-test: 109 passed, 1 skipped, 0 failed.
  • hip-host-test, test-qwen35-hip-host and check-native-hip-host exit 0.
  • zig build test on a git archive of each side: 465/465 steps, 6,185 of 6,230 passed and 45 skipped, the same as main.
  • tools/zig/lean_check.py: the same count as main. zig fmt --check is clean.

• Not run, and why:

  • gfx1030, gfx1100 and gfx1201: I have no card. Their kernels change the same way, and plan.hip builds for all three with the build's flags; the DPP row_xmask and permlanex16 folds exist on gfx10 and later. The contract records for them come from ROCm 7.13's code objects.
  • Nemotron and Flash-Next: I have no HIP checkpoint locally. They use the same planned walk, and the contract's cases cover its d = 128 shapes.
  • Metal and CUDA: untouched.

Checklist

• Every token still goes through the lane rounds. No serial path, and nothing that needs drafts off.
• Drafted output equals "draft": false, a resumed prompt equals a fresh one, and concurrent equals solo.
• No precision traded for speed. If the bits change, the description says which sums change and why.
• A new or changed GPU kernel is registered with the kernel contract (docs/kernel-contract.md); each backend's SUMMARY line is in the receipt.
• Prompt processing is no slower than the last release.
• The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
• New tests fail before the change, pass after it, and skip cleanly without their dependency.
• Comments and docstrings are one line. No measurements or history in the source.
• No personal data, machine names, internal hosts or local paths. No AI attribution lines.

[ashhart]

Thanks @AbrahamCete. We gave this a full independent review. The four-key and two-key kernels keep the lane mapping, in-lane order and the xor-16-to-1 folds, the host-emulated test pins every score for both query types including 16,384-key tails ending mid-warp, and your served receipt has identical tokens and logits. The four extra contract limit problems are the same one-VGPR live-against-recorded gap, so they don't block. It lands now: it's gated and queued, and we've written the changelog line.