Repository navigation
perf(hip): prompt attention tile keeps its probabilities in registers and fits 6 waves, same bits - #1788
Conversation
|
/atc exchange a8515ed93e2765eb7d841dc1fa9cd71fc650e650db5058b6c35022cdfe9d2540 |
|
Thanks @AbrahamCete, the receipts answer everything: the pin and the 1,816-case digest equal to v1.2.0, served |
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.
c195321 to
7217ce7
Compare
|
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. |
|
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 |
What this changes
fa_dot2<T, paged>(zig/kernels/hip/attention/prefill.hip, symbolstf_fa_paged_bf16,tf_fa_paged_f16,tf_fa_wide_bf16andtf_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 ofds_bpermute, in the same xor order, so every add keeps its operands.ds_bpermuteper 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 withv_permlane16instead 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 withattention_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 setsreplacesto the paged tile: gate 3 requires their bytes to equal it at every edge shape. Four records inhip_records_prompt.zig, chained intohip_records.zig, fromhip-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 changesprefill.hip. #1612, #1614, #1622, #1650, #1658 and #1667 share registration lists with this one (hip_records.zig,contract_main.zig, thekernel_tests.zigimports,docs/kernel-contract.md). This PR adds one line tohip_records.zig, one import tokernel_tests.zigand one entry block tocontract_main.zig.Known finding (pre-existing)
Registering the tile surfaced a gate 1 finding that v1.2.0 already has, recorded as
knownin the spec and in docs/kernel-contract.md:fmaxfand a row at max -inf takes p = 0;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
mainc3c3bd0. It started from v1.2.0 and rebased with no conflicts;main's changes since v1.2.0 don't touchprefill.hip. Withmain's new recordship_records.zigwould reach 607 lines, so the four prompt-tile records live inhip_records_prompt.zig. On the rebased tip, 2,048-token prompt logits for Qwen3.8-27B are byte-identical tomain's, bf16 and--f32. The timing receipts below compare v1.2.0 with the branch on v1.2.0.Kernel time,
rocprofv3 --kernel-traceovertf-qwen35-test prefill <27B> 8192(256 calls a run):Prompt speed: the cold first request's
ttft_sfromtools/bench_cache.py --require-cold --tokens 16.tensorfold-native serve --context 66560 --parallel 1 --prompt-cache-gib 0 --snapshot-dir none --max-snapshots 0, a fresh server for each run.cached_tokens = 0.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:
tf-qwen35-test logits, bf16 and--f32, 2,048 tokens) are byte-identical to v1.2.0, andtf-qwen35-test check <27B> --only invariantsgives 42/0/5 as on v1.2.0.token_shais the same before and after at every size (2kfaceb02af600, 8k1ebfd15d62c7, 32kc8dfa0175e40, 64kd92b23432e24).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_shais the same before and after: code2b28aa2108de, chata06fab348fbe.--f32.ds_bpermute(32 on v1.2.0).v_permlane16reads 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-raw13.29 and 13.18 before, 13.29 and 13.32 after;gpu-chat-no-think13.60 and 13.61 before, 13.63 and 13.55 after.bench_concurrentaggregate 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 ismain's Flash-Next MTP wide-norm test). The new test iskernel_tests/attention_pin.zig.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 asmain's 126. On v1.2.0 with only the spec added, gate 1 shows the same 345 class differences.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.* scalefails 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 asmain.zig fmt --checkis clean on the changed Zig files.test-qwen35-hip-host,hip-host-test,check-native-hip-hostandtf-qwen35-test: exit 0.Not run, and why:
Checklist
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (Bytes unchanged: the pin digest, prompt logits, servedtoken_sha, andbench_concurrentalone and serial equal.)SUMMARYline is in the receipt. (Registered in this PR; one pre-existing gate 1 finding is recorded as known.)