Repository navigation
perf(hip): the decode tail issues a row's loads before its branches - #2247
Conversation
|
Thanks @AbrahamCete. Pinning main's tail bytes before changing the kernel, and planting two wrong orders against it, makes byte-equal easy to trust, and registering |
|
@ashhart, done: the receipt now has main-against-branch decode tok/s from |
|
Thanks @AbrahamCete, the decode numbers are what we wanted beside the kernel timing. Ready, queued for after the release cut. |
1cbe0df to
d38c100
Compare
|
@AbrahamCete, #2247 conflicted with tonight's main in |
|
/atc exchange 34d76bc8277b45b07c725da75e6d4b83841c17084e635a6344c638b6e95e7141 |
The tail's test bounds its error against float64, so a change in its sum order or roundings could pass it. The new test hashes tf_tail's residual and norm in fp16 and bf16 over plain and slot-summed rows, a verify window, ragged, unweighted and the widest rows, and cases whose slots cancel, with three and with nine slots, so that the order of the slot sum shows; gfx1151's digest is pinned.
tf_tail chose fp16 or bf16 and checked the row's end at every load, so each load sat in its own basic block and waited before the next one issued; a MoE tail also waited once a slot. The tail now runs at one activation kind and with or without slots as template instances, a thread loads four elements at once with the indices past the row clamped to its last, and a MoE tail loads eight slots' rows before it sums them. Every element's slots are summed in slot order with the same roundings and the sum of squares is taken in the same order, so every output byte is unchanged for inputs without NaN. With the kind fixed, the fp16 instance adds in fp16: the same values, but an fp16 NaN may come out with another sign or payload.
tf_tail, the decode tail that adds the residual (a row of the activation kind, or a MoE layer's slots' fp32 rows times their weights, summed in slot order) and writes the RMS norm after it, had no contract entry or gate-5 record. CONTRIBUTING requires one for a changed kernel, and this branch changes how it loads its rows. contract_tail.zig registers both forms, plain and nine slots, against a float64 reference at the served widths (Qwen3.6-35B-A3B's 2,048, the 27B's 5,120 and the widest row of 8,192), decode rows and windows, and the ragged edges of a pass over the row. It is not bit-equal to the launches it merges (their sums reduce in another order), so it declares no replacement and meets the float64 bound on its own; its non-finite plants and gate 4 follow tf_add_rms's. Gate 5 records its usage from ROCm 10.1's hipcc on gfx1030, gfx1100, gfx1151 and gfx1201, and the coverage list names it.
d38c100 to
88b720b
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. |
What this changes
tf_tail(zig/kernels/hip/decode/decode.hip) is the decode round's residual add and RMS norm in one launch, and after a MoE layer it also sums the slots' fp32 rows. It chose fp16 or bf16 and checked the row's end at every load, so each load sat in its own basic block and waited before the next one issued. A MoE tail also waited once per slot.The tail now runs as template instances, one per activation kind and one each with and without slots. A thread loads four elements at once, with any index past the row clamped to the row's last element; only the stores and the sum of squares check the row's end. A MoE tail loads eight slots' rows before it sums them. On gfx1151 the kernel takes half the time and the served gain is small (receipt below).
Each element's slots are summed in slot order with the same roundings, and the sum of squares and the block reduction run in the same order. So every output byte is unchanged for inputs without NaN. One exception: with the kind fixed at compile time, the fp16 instance adds in fp16. It gives the same values, but an fp16 NaN may come out with another sign or payload. Kind 1 (fp16) is served only on gfx1030, and only NaN activations reach this.
Two test changes come with it:
kernel_tests/tail_pin.zigpins the tail's residual and norm bytes in fp16 and bf16. Before, the tail's only test bounded its error against float64.tf_tailwas not registered with the kernel contract, so it is now (zig/tests/hip/contract_tail.zig,hip_records_tail.zig).Issue: n/a. There is no issue: this is a byte-identical speedup found while profiling decode on gfx1151, with no model family or arithmetic change.
Stacked on: main.
atc overlap --json(2026-11-05 11:16Z, against main 610c67b): no conflict this PR causes except one line inhip_records.zigwith Clef-Flash: joint-schema decisions on /v1/systemone (HIP) #2048, where both append a records file to the same list. Its other conflicts there are with main, and Nemotron Parse on HIP: NVIDIA's checkpoint, token-exact against its HF example #1941's conflicts are all with main. perf(hip): a skewed-row matrix tile, three blocks a WGP on gfx1151 for prompts up to 8192 codes deep #2087 sets the same record-count test inhip_records_families.zigto 183/171; whichever lands second moves it to 184/172.Checks on this head
python3 tools/zig/lean_check.py: 425 on the base and 425 on this head, 0 new (tools/zig/lean_new.py).zig build hip-kernel-test --summary all: 136/140 passed, 4 skipped, 0 failed.tail_pinholds the digest main's kernel writes, and the branch's kernel writes the same one.zig build test -Dhip=true --summary allin a cleangit archivetree: 587/587 steps, 7446/7508 tests passed, 62 skipped.zig build test-golden --summary all: 5/5 steps.tf-qwen35-test check --only invariants: 42/0/5 on the 27B, 46/0/1 on the 35B-A3B.zig fmt --checkandgit diff --check: clean.zig,zig-windows,zig-macos-dev-tools,fmt,release-paths, andbenchmark-clientson Ubuntu and Windows).tail_pin, which passes kind 1 to the launcher.Receipt
tf_tailis launched only by the qwen3_5 family's HIP lane rounds (forward/window.zig):addRmsafter attention andmoeTailafter a MoE layer. So it reaches the decode and verify rounds of that family.tf_add_rmsand is untouched. Metal and CUDA are untouched.token_shais equal in every cell below: 27B code2b28aa2108de, chata06fab348fbe; 35B-A3B code9f2923ac14d0, chata8c28b8297b4.bench_concurrent --alone --serial: alone and serial equal in every cell, 0 failed.tf_tailunder rocprofv3,tf-qwen35-test logits --decode, 64 steps:On the 35B-A3B, two alternated rounds put the tail after attention at 4.68 -> 2.78 us and the MoE tail (9 slots) at 8.57 -> 3.73 us.
Decode tok/s,
tf-qwen35-test check --only speed, which times the last of three runs: one stream generating 256 tokens, and four streams generating 128 each. Main 0311c92 then this head, built from clean trees and alternated B1 C1 B2 C2:A 35B-A3B decode round at one stream drops from 12.4 to 12.1 ms, in line with the kernel timing above. The 27B checkpoint has no draft head, so with drafts on it still decodes one token a round and reads within 0.3 tok/s of these.
Served: one fresh server a run,
--parallel 4 --context 4096 --prompt-cache-gib 0, alternated B1 C1 B2 C2,bench_concurrent --tokens 256 --reps 3 --temperatures 0 --levels 1,4 --alone --serial. Aggregate tok/s, main then this PR. These ran on main 249e33b before the rebase; the kernel is the same commit's.tf_add_rms, nottf_tail. The samecheck --only speedruns give the same prefill tok/s at 13, 2048 and 8192 tokens on both builds, within run-to-run noise.Kernel contract, only when a GPU kernel is new or changed
tf_tail, in two modes: a row of the activation kind, and nine fp32 slots (zig/tests/hip/contract_tail.zig).tf_addandtf_rmsreduce in another order), so it states noreplacesand meets the float64 gates on its own.hip_records_tail.zig, filled from ROCm 10.1'ship-kernel-contract-limitsRECORD line.contract_add_rms.zigdoes not. The harness's 11 controls go red as planted.SUMMARYline and exit status:hip-kernel-contracton gfx1151:SUMMARY gfx1151 95 kernels: 95 pass gates 1-4, 0 fail, 2 known findings; 11 controls: 11 red as planted, 0 missed; 6 limit problems(exit 1 for the limit problems). Main 0311c92 readsSUMMARY gfx1151 93 kernels: 93 pass gates 1-4, 0 fail, 2 known findings; 11 controls: 11 red as planted, 0 missed; 6 limit problems.hip-kernel-contract-limits:FAIL kernel-contract-limits: 22 problems, the same 22 as main.tf_tailreadsrecord currenton gfx1030, gfx1100, gfx1151 and gfx1201: 125 VGPRs on gfx1030 and 120 on the others, 58 SGPRs, 32,832 LDS bytes.Checklist
The title and description match the head commit, including the receipt and the stack.
New tests fail before the change, pass after it, and skip cleanly without their dependency. (
tail_pinpins bytes the change must keep. Planted changes to the slot order and to the norm-weight pairing fail it, and it skips with the digest printed on a chip without one.)Comments and docstrings are one line. No measurements or history in the source.
No personal data, machine names, internal hosts or local paths in the diff or its commits. No AI attribution lines.
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. (The bytes are unchanged: decode logits, servedtoken_sha, alone and serial equal.)No precision traded for speed. If the bits change, the description says which sums change and why. (No sum changes. The only bits that can change are an fp16 NaN's sign or payload, as above.)
Decode and prompt processing are no slower than main at one stream and the widest admitted; with
--drafter, no slower than without it, greedy and sampled. (Served cells above, with the 35B-A3B's MTP drafts on; the prompt path is untouched.)