[AbrahamCete]

What this changes

Past the one-block walk (spans of more than 4096 keys), the planned and flat decode walks run tf_softmax_stats and tf_sum_partials (attention.hip) and tf_sum_partials_act (plan.hip) as launches of their own, one block a head. Every load in them waited on the one before:

• tf_softmax_stats: a thread reads keys i, i + 256, … once for the max and again for the sum of exps.
• tf_sum_partials: a lane adds its column's tiles one after another, 512 of them at 64K.

On gfx1151 at 64K keys, Qwen3.6-35B-A3B spent 145 µs and 172 µs a full-attention layer on them, for a few MB of data.

The changes, one commit each:

• test(hip): pin the decode softmax stats' max and the partial sums' tile order
  • Nothing pinned these kernels' bits past 4096 keys. The new kernel_tests/stats.zig checks two things.
  • tf_sum_partials and tf_sum_partials_act, bit for bit, against a host sum of each lane's tiles in tile order, rounded as __float2half_rn and float_to_bf16 round for the activation kinds. It covers 1 to 1024 tiles at d = 64, 128 and 256.
  • tf_softmax_stats' max, bit for bit. Each row's max is planted at the last visible key, the first, past 4096 or mid-row. The sum is checked within 1e-5 of float64, since the device's expf is not the host's.
• perf(hip): the decode softmax stats and partial sums load sixteen ahead
  • Each loop issues sixteen loads together, then folds or adds them in the same order. tf_sum_partials_act takes the same loop.
  • Records: tf_softmax_stats goes from 9 to 61 (gfx1030) and 70 VGPRs, tf_sum_partials from 5 to 23, and tf_sum_partials_act from 6-7 to 24, from ROCm 7.13's code objects. One block runs a CU, so occupancy does not bind.

Every max, sum and exp keeps its order, so the bits are the same. I also checked the softmax sum's bits against the previous kernel directly. With a debug print of the stats' digest, the old and new kernels wrote the same bytes on all 10 long-span cases of the test, and the digests differ from case to case.

Kernel time a full-attention layer (rocprofv3, tf-qwen35-test lanes with drafts off, the launches at the context's span):

kernel	35B, 32K	35B, 64K	27B, 32K	27B, 64K
tf_softmax_stats	48.3 → 10.2 µs	145.4 → 24.4 µs	74.6 → 16.1 µs	152.3 → 36.2 µs
tf_sum_partials	50.2 → 9.1 µs	172.0 → 28.8 µs	89.4 → 19.6 µs	183.8 → 49.2 µs

In a separate sweep on the 35B at 64K, eight loads ahead left 76 µs for the two kernels together, and sixteen left 54.

Overlap with open pull requests:

• perf(hip): the planned decode score takes four keys a warp on long spans, two on short ones #1932 (mine) changes the score kernel in plan.hip and adds score records to hip_records.zig. This PR changes tf_sum_partials_act in plan.hip and three other records. The two merge cleanly in either order.
• hip: a 64-row routed prompt tile for many-expert plans on gfx1151 #1658, hip flashnext: sparse attention with a lane a key over several heads, the scalar kernel's bits #1877, dots.mocr on HIP: the publisher's BF16 OCR checkpoint, token-exact against its model card #1938, GLM-OCR on HIP: the publisher's BF16 checkpoint, exact against its six model-card tasks #1940 and Nemotron Parse on HIP: NVIDIA's checkpoint, token-exact against its HF example #1941 add entries to hip_records.zig. hip flashnext: sparse attention with a lane a key over several heads, the scalar kernel's bits #1877 also adds to kernel_tests.zig, where this PR adds one import line.

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 45fcab6. Both sides were built from git archive trees.

• 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	39.96 / 39.85	41.34 / 41.52 (+3.8%)
  35B-A3B, 64K	25.59 / 25.95	27.03 / 26.99 (+4.8%)
  27B, 32K	9.84 / 9.86	10.03 / 10.07 (+2.1%)
  27B, 64K	7.31 / 7.29	7.50 / 7.48 (+2.6%)

  • Short prompts, tools/bench_openai.py --tokens 256 --reps 3 --temperatures 0. 35B: code 112.2 / 112.9 against 112.6 / 112.5, chat 82.4 / 81.0 against 83.4 / 83.7. 27B: code 13.3 and chat 13.6 on both.
  • tf-qwen35-test check --only speed, twice each on the 35B:
    • one stream, drafts off: 70.1 / 70.5 against 70.2 / 70.2;
    • one stream, drafts on: 95.7 / 95.5 against 95.5 / 95.1;
    • four streams, drafts off: 138.7 / 138.4 against 138.7 / 138.5;
    • four streams, drafts on: 125.8 / 125.1 against 126.5 / 125.2.
    • The 27B is the same within 0.2.

• Exactness:

  • Token hashes are equal on every server of every run: 35B 61da022c at 32K and aaa410bc at 64K; 27B 28bb6ebb and 63bf1db1.
  • lanes replies are byte-identical 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.7 / 27.7 against 27.6 / 27.5 s at 32K, and 72.0 / 71.7 against 71.7 / 71.6 s at 64K.
  • 27B: 104.0 / 103.9 against 103.9 / 103.9 s, and 246.3 / 246.2 against 247.1 / 246.4 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; 131 limit problems. main has 130.
  • The extra one is tf_softmax_stats' live usage under this machine's hipcc. It differs from the ROCm 7.13 record outside its VGPRs and LDS, which match (70 and 1024). That is the same toolchain gap as the rest of main's 130.
  • hip-kernel-contract-limits with ROCm 7.13: 220 problems on both, with the three changed records current.

• Tests run:

  • The new tests pass on main and on this PR. They fail when the partial sums add their tiles in reverse, where Nemotron's activation-type walk test fails too. They also fail when the max skips the last visible key.
  • hip-kernel-test: 111 passed, 1 skipped, 0 failed. main has 109/1/0.
  • hip-host-test, test-qwen35-hip-host and check-native-hip-host exit 0.
  • zig build test on a git archive of each side: 470/470 steps, 6,207 of 6,252 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, attention.hip and plan.hip build for all three, and their records come from ROCm 7.13's code objects.
  • Nemotron and Flash-Next: I have no HIP checkpoint locally. They run the same kernels past the one-block walk. The new test covers d = 64, 128 and 256, and the kernel tests' Nemotron walk cases pass.
  • 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 reviewed this independently: tf_softmax_stats and tf_sum_partials/_act issue sixteen loads ahead and fold them in the same index order with the same operations, so the bits are unchanged. It lands now: it's gated and queued, and we've written the changelog line.