Skip to content

perf(hip): prompt attention tile keeps its probabilities in registers and fits 6 waves, same bits - #1788

Merged
ashhart merged 7 commits into
ashhart:mainfrom
AbrahamCete:perf/hip-prefill-attention
Nov 1, 2026
Merged

ashhart merged 7 commits into
ashhart:mainfrom
AbrahamCete:perf/hip-prefill-attention

Conversation

@AbrahamCete

Copy link
Copy Markdown
Contributor

What this changes

fa_dot2<T, paged> (zig/kernels/hip/attention/prefill.hip, symbols tf_fa_paged_bf16, tf_fa_paged_f16, tf_fa_wide_bf16 and tf_fa_wide_f16) is the prompt path's attention tile on HIP. Qwen3.8-27B runs it in its 16 full-attention layers, and it grows with the prompt: on gfx1151 at v1.2.0 it was 1.7% of a 2,048-token prompt's kernel time and 6.3% at 8,192. Each block takes 64 query rows through 64-key tiles with 256 threads, and it held 50,176 B of LDS, so gfx1151 ran 4 waves per SIMD.

The changes, one commit each:

  • test(hip): pin the paged prompt tile's output bits on 27B shapes. Fixed integer-derived inputs at 24/4/256, with partial last pages, a short last row block, diagonal tiles and whole tiles before them, and an FNV-1a digest of the output per architecture (gfx1151 recorded, with the hipcc named below). The existing test compares the paged and flat forms of the same template, so a reordered sum passes it; a planted reversed P×V order fails the new one.
  • perf(hip): skip the masks and range guards on whole prompt tiles. A tile entirely inside the causal edge and the span takes a branch without per-element guards. The -inf selects stay, since removing them changes bits for non-finite rows.
  • perf(hip): reduce the prompt tile's softmax rows with DPP butterflies. The 16-lane row max and row sum read their partners with DPP instead of ds_bpermute, in the same xor order, so every add keeps its operands. ds_bpermute per tile goes from 32 to 0.
  • perf(hip): prefetch the prompt tile's next key and value chunks. The next chunk's K/V loads are issued before the current chunk's products.
  • perf(hip): keep the prompt tile's probabilities in registers. Each lane keeps its rows' probability pairs and rescale factors in registers, and the P×V products take a key pair's probabilities from the lane that holds them with v_permlane16 instead of an LDS round trip. Without the probability, rescale and sum arrays a block uses 40,960 B of LDS, so gfx1151 runs 6 waves per SIMD instead of 4.
  • test(hip): register the prompt attention tiles with the kernel contract. None of the four symbols had a spec or record. This adds a HIP spec on CUDA's prompt-attention pattern (a float64 reference with attention_ref's tolerance, all-equal keys for gate 4), with the served shapes of Qwen3.8-27B, Qwen3.6-35B-A3B and Nemotron. The wide tiles have no served caller, so their spec sets replaces to the paged tile: gate 3 requires their bytes to equal it at every edge shape. Four records in hip_records_prompt.zig, chained into hip_records.zig, from hip-kernel-contract-limits.

Every sum keeps its operands and its order (dims ascending in dot2, the 16-lane key layout, fmaf(alpha, o, pv), expf, 64-key tiles), so the output bits are the same.

Overlap with open pull requests (atc overlap, plus the file lists): no open PR changes prefill.hip. #1612, #1614, #1622, #1650, #1658 and #1667 share registration lists with this one (hip_records.zig, contract_main.zig, the kernel_tests.zig imports, docs/kernel-contract.md). This PR adds one line to hip_records.zig, one import to kernel_tests.zig and one entry block to contract_main.zig.

Known finding (pre-existing)

Registering the tile surfaced a gate 1 finding that v1.2.0 already has, recorded as known in the spec and in docs/kernel-contract.md:

  • a row whose every score is NaN writes 0 where float64 gives NaN, as the running max is fmaxf and a row at max -inf takes p = 0;
  • an infinite value at a key whose fp32 weight is 0 (past the row's causal edge in its 64-key tile, or under fp32's least) gives NaN where float64 gives inf or a finite value, as P×V takes every key of the tile.

On the plant set, 345 of 49,152 outputs differ in class, on this branch and on v1.2.0 with only the spec added (254 from the NaN query row, 91 from infinite values at zero weights). I left the kernel as it is, since fixing it changes bytes for non-finite inputs. I can open an issue or a follow-up if you want it fixed.

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.8-27B-MLX-4bit and Qwen3.6-35B-A3B-MLX-4bit-MTP.

  • Base: the branch is on main c3c3bd0. It started from v1.2.0 and rebased with no conflicts; main's changes since v1.2.0 don't touch prefill.hip. With main's new records hip_records.zig would reach 607 lines, so the four prompt-tile records live in hip_records_prompt.zig. On the rebased tip, 2,048-token prompt logits for Qwen3.8-27B are byte-identical to main's, bf16 and --f32. The timing receipts below compare v1.2.0 with the branch on v1.2.0.

  • Kernel time, rocprofv3 --kernel-trace over tf-qwen35-test prefill <27B> 8192 (256 calls a run):

    • v1.2.0: 22.46 ms a call (6.3% of kernel time);
    • whole-tile fast path: 22.25-22.33 ms;
    • DPP butterflies: 21.69 ms;
    • prefetch: 21.54 ms;
    • probabilities in registers, 6 waves: 18.48 ms (-18% overall). At 2,048 tokens: 5.82 -> 4.79 ms.
  • Prompt speed: the cold first request's ttft_s from tools/bench_cache.py --require-cold --tokens 16.

    • Server: tensorfold-native serve --context 66560 --parallel 1 --prompt-cache-gib 0 --snapshot-dir none --max-snapshots 0, a fresh server for each run.
    • Runs alternated B1 C1 B2 C2 B3 C3. Every request reported cached_tokens = 0.
    prompt v1.2.0 ttft (s) branch ttft (s) change
    2048 5.55 [5.53-5.57] 5.50 [5.49-5.51] -0.9%
    8192 23.16 [23.11-23.37] 22.92 [22.91-22.97] -1.0%
    32768 110.07 [109.55-110.30] 106.44 [106.27-106.52] -3.3%
    65536 266.68 [265.77-266.82] 252.55 [251.91-252.65] -5.3%

    The gain grows with the prompt because attention's share does. The prompts are the fixed token arrays used for perf(hip): band the prefill WMMA GEMM's blocks so row blocks share weights #1648 and perf(hip): keep the matrix tile's inner loop free of selects and divisions #1663.

  • Exactness:

    • After each kernel commit: the pin digest matches v1.2.0's, prompt logits (tf-qwen35-test logits, bf16 and --f32, 2,048 tokens) are byte-identical to v1.2.0, and tf-qwen35-test check <27B> --only invariants gives 42/0/5 as on v1.2.0.
    • Served bytes: the cold replies' token_sha is the same before and after at every size (2k faceb02af600, 8k 1ebfd15d62c7, 32k c8dfa0175e40, 64k d92b23432e24).
    • tools/bench_concurrent.py --levels 1,4 --tokens 256 --reps 3 --temperatures 0 --alone --serial: every cell reports alone and serial equal, and 0 failed. token_sha is the same before and after: code 2b28aa2108de, chat a06fab348fbe.
    • A wider check against v1.2.0, on gfx1151:
      • Digests of 1,816 cases are identical: qlen 1-1024 by q_pos0 0-1000 at 24/4/256, 16/2/256 and 32/2/128, both cache types, paged and flat, plus wider value ranges, batch 2, spans the engine never passes, and NaN and inf in the cache. Unused page slots, rows past the span and the output start as NaN, so a stray read or a missing write would show.
      • The paged and flat tiles give the same bits in all 896 pairs, and 20 repeated runs give the same digests.
      • 2,048-token prompt logits are byte-identical for Qwen3.8-27B and Qwen3.6-35B-A3B, in bf16 and --f32.
      • The code objects: 40,960 B of LDS, no scratch, 0 ds_bpermute (32 on v1.2.0). v_permlane16 reads its partner by lane within 16-lane rows; with wave32 and 1D blocks of 256 threads, lane % 16 = tid % 16, and v1.2.0's __shfl_xor(x, mask, 16) already relied on it.
  • Decode speed (decode does not use this tile): server --context 4096 --parallel 4 --prompt-cache-gib 0, runs alternated B1 C1 B2 C2.

    • tools/bench_openai.py --tokens 256 --reps 3 --temperatures 0, median tok/s: fibonacci-raw 13.29 and 13.18 before, 13.29 and 13.32 after; gpu-chat-no-think 13.60 and 13.61 before, 13.63 and 13.55 after.
    • bench_concurrent aggregate at 4 streams: code 43.9 and 43.9 before, 43.8 and 44.0 after; chat 43.4 and 43.4 before, 43.4 and 43.3 after.
  • Tests run:

    • hip-kernel-test (binary run directly) on the rebased tip: 108 passed, 1 skipped, 0 failed (main: 107 passed, 1 skipped; the skip is main's Flash-Next MTP wide-norm test). The new test is kernel_tests/attention_pin.zig.
    • Kernel contract on the rebased tip: SUMMARY gfx1151 52 kernels: 52 pass gates 1-4, 0 fail, 2 known findings; 5 controls: 5 red as planted, 0 missed; 130 limit problems. main: SUMMARY gfx1151 50 kernels: 50 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 126 limit problems. The two new entries are the paged and wide tiles of the chip's cache type, and the two known findings are the gate 1 finding above, on each. The 4 more limit problems are the new records checked on the device: system hipcc gives 234-235 VGPRs on gfx1151 where the ROCm 7.13 records hold 231-232, the same kind as main's 126. On v1.2.0 with only the spec added, gate 1 shows the same 345 class differences.
    • Limits with ROCm 7.13 hipcc: 220 problems on the rebased tip and on main; the 16 new lines (4 symbols x 4 chips) are current. Every tile uses 40,960 B of LDS and no scratch on gfx1030, gfx1100, gfx1151 and gfx1201.
    • Planted faults: dropping * scale fails the paged tile's gate 3; a paged-only fault fails its gates 3 and 4 and the wide tile's gate 3.
    • tools/zig/lean_check.py: 424 problems, the same as main. zig fmt --check is clean on the changed Zig files.
    • test-qwen35-hip-host, hip-host-test, check-native-hip-host and tf-qwen35-test: exit 0.
  • Not run, and why:

    • gfx1030, gfx1100 and gfx1201: I have no card. The tile compiles for all three with no scratch.
    • The f16 symbols' gates 1-4 run only on an fp16 build (gfx1030); here they are covered by the limits audit and the shared template.
    • Nemotron: no checkpoint locally. The contract spec covers its 32/2/128 shape.
    • 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. (Bytes unchanged: the pin digest, prompt logits, served token_sha, and bench_concurrent alone and serial equal.)
  • No precision traded for speed. If the bits change, the description says which sums change and why. (The bits do not change.)
  • 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. (Registered in this PR; one pre-existing gate 1 finding is recorded as known.)
  • Prompt processing is no slower than the last release. (0.9-5.3% faster than v1.2.0.)
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not. (gfx1151 only.)
  • New tests fail before the change, pass after it, and skip cleanly without their dependency. (The pin test fails on a planted reorder; the contract spec fails on planted faults.)
  • 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.

@AbrahamCete

Copy link
Copy Markdown
Contributor Author

/atc exchange a8515ed93e2765eb7d841dc1fa9cd71fc650e650db5058b6c35022cdfe9d2540

@ashhart

ashhart commented Nov 1, 2026

Copy link
Copy Markdown
Owner

Thanks @AbrahamCete, the receipts answer everything: the pin and the 1,816-case digest equal to v1.2.0, served token_sha unchanged, and 18% off the tile's kernel time at 8K. We'll land it with gfx1030, gfx1100 and gfx1201 listed as compiled but not run. A separate issue for the two known non-finite findings is welcome, as a fix there would change bytes for non-finite inputs and wants its own review.

AbrahamCete and others added 7 commits November 1, 2026 06:37
The paged tile was only compared with the flat tile of the same template, so a change to both left the test green. The new test digests its output on 24 heads over 4 kv heads at 256 dims, with partial pages, a short last row block, diagonal tiles and later chunks, and pins the digest on gfx1151.
A key tile that ends inside the span and before the first row of the 64-row block is seen in full by every stored row. Such tiles now take a path without the per-score mask and without the guarded key and value loads. The tiles are visited in the same order and every stored value is the same, so the output bits do not change.
The 16-lane max and sum of each softmax row now pair lanes with DPP row_xmask moves instead of ds_bpermute shuffles. The pairing over lane distances 8, 4, 2 and 1 and each lane's own-then-partner operand order are unchanged, so the output bits are the same.
Each staged chunk's global loads now issue right after the previous chunk reaches LDS, before that chunk's products; the first value chunk and the next tile's first key chunk issue the same way, and the next tile's page id loads one chunk ahead. A row past the span reads the span's last row and is zeroed when staged, instead of a guarded load. The arithmetic and its order are unchanged, so the output bits are the same.
Each lane keeps its rows' bf16 probability pairs and rescale factors in registers, and the PV products take a key pair's probabilities from the lane that holds them with v_permlane16 instead of an LDS round trip. The PV work covers rows sr + 16i at dims 4sk .. 4sk + 3, the rows each lane already holds, and the epilogue divides by the lane's own running sum. Without the probability, rescale and sum arrays a block uses 40,960 bytes of LDS instead of 50,176, so three blocks fit a workgroup processor and gfx1151 runs 6 waves per SIMD instead of 4. Every sum keeps its operands and order, so the output bits are the same.
The four symbols of prefill.hip's 64-row tile now have contract specs and launch records. tf_fa_paged_bf16 and tf_fa_paged_f16 run over a permuted page table at the 27B's, the 35B-A3B's and Nemotron's head counts, each output within attention_ref's float64 bound, and gate 4's equal keys prove fp32 sums. tf_fa_wide_bf16 and tf_fa_wide_f16, which no model calls, must equal the paged tile's bits on the same keys. Gate 1 is a known finding on both, and v1.2.0 shows the same one: a row whose every score is NaN writes 0, and an infinite value at a key whose fp32 weight is 0 gives NaN, as P x V takes every key of its tile. The records hold what ROCm 7.13's hipcc compiles for gfx1030, gfx1100, gfx1151 and gfx1201.
@ashhart
ashhart force-pushed the perf/hip-prefill-attention branch from c195321 to 7217ce7 Compare November 1, 2026 06:37
@ashhart
ashhart merged commit 4446638 into ashhart:main Nov 1, 2026
9 checks passed
@ashhart

ashhart commented Nov 1, 2026

Copy link
Copy Markdown
Owner

Thank you @AbrahamCete. This is merged into main with your commits under your own name, so you're listed as a TensorFold contributor. It ships in the next release, credited to you in the changelog and the release notes.

@AbrahamCete

Copy link
Copy Markdown
Contributor Author

Thank you @ashhart, glad it's in. I've opened #1821 for the two non-finite findings.

@ashhart

ashhart commented Nov 1, 2026

Copy link
Copy Markdown
Owner

Thanks @AbrahamCete! The faster HIP prompt attention tile on gfx1151, now under the kernel contract, is in 1.2.1: https://github.com/ashhart/TensorFold/releases/tag/v1.2.1

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants