Repository navigation
perf(hip): wide stream rounds reuse cached weights, and a chunk's code pairs are made once - #1694
AbrahamCete wants to merge 7 commits into
Conversation
|
/atc exchange a4436cc5af37c0bc75f72d0c8ffe6b6378e346fad48e70287735aae0dd0fa826 |
… cached weights A lane round of more than four rows ran the stream tile as ceil(rows / 4) row blocks on the grid's y, so each later row block read every weight from DRAM again. The row blocks of one column block now run next to each other on x and the later ones find its weights in cache. Each row's arithmetic is unchanged, so a row still writes the same bytes alone and in a round. On gfx1151 (Radeon 8060S), Qwen3.8-27B MLX 4-bit, alternated A/B against main: verify of 8 rows 154 -> 124 ms, 16 rows 283 -> 208 ms; 1 to 4 rows and plain decode unchanged. Qwen3.6-35B-A3B MLX 4-bit MTP community-v1 unchanged.
…nce layout The interleave was an implicit contract between the stream launch and stream_tile. stream_layout.zig is now the one place that lays a round's row blocks on the grid (StreamSides moves there with it), Kernels.stream_interleave switches back to the reference layout on y, and a GPU test runs plain, relu2, swiglu pair, routed and grouped products at rounds of 5 to 32 rows in both layouts and requires equal bytes. A planted swap of the column and row block indices fails it. Naming the layout in a StreamSides field read by stream_tile was tried and dropped: on gfx1151 it doubled the stream kernels' VGPRs (4x4: 140 -> 256) and cost about 10% on wide verify rounds.
…erleave Each row goes to both experts (x_div 2), so every expert holds the whole round and the routed case reaches the interleaved grid from 5 rows; each round now asserts it spans more than one row block.
…t BF16 pairs from bytes Each column's 16 code pairs are made once a chunk, then every row of the block dots them in turn, in the order a row alone takes, so the bytes stay. On the BF16 build a 4-bit word becomes two byte-per-code words and v_cvt_f32_ubyteN makes each code's float in one op; v_perm_b32 packs the top halves. The 4x4 tile's inner loop drops from 1149 to 852 instructions on gfx1151. gfx1151, A/B alternated twice of the previous commit's build against the same stream.hpp change on a local experiment build (1, 2 and 4 rows a block): Qwen3.8-27B verify 9-16 rows 204 -> 179 ms, decode 13.2 -> 13.6 tok/s; Qwen3.6-35B-A3B decode 71.4 -> 76.1 tok/s, MTP 95.9 -> 101.9 tok/s.
…tream_pair made Plain, gate|up pair and top-2 routed launches of 1 to 16 rows hash to the digest the stream tile wrote before the previous commit, when its code pairs came from stream_pair.
…'s rows The second expert takes only the first half of the rows, so its row blocks past them return early in both the interleaved and the reference layout.
The golden test held one digest for every chip with BF16 activations. Its bytes also depend on each chip's BF16 dot and activation rounding, so another architecture could fail it with a correct byte path. The digest is now recorded per architecture, with gfx1151's as the first entry. On an architecture without one, the test prints the digest it computed and skips.
675c33b to
f69c292
Compare
|
/atc exchange 64237b5932fbc98e1afc246183edcbe674e9b656736e7d9077294b9084dfa330 |
|
Thanks @AbrahamCete, the 27B's +30% to +47% at 8 and 16 streams with the same bytes is a real win, and the layout and golden tests are what we'd want. One thing blocks landing: on the 35B-A3B at 4 streams, code's mean falls 4.5%, because the cheaper wide rounds tip the lane core into its deeper, lower-acceptance mode more often (6 of 12 servers against 1 of 11). We can't land a default that's slower on average for a served model. Please open the issue you mention with the data. Then either a fix where the lane core prices the cheaper rounds so it doesn't choose the slow depth, measured on the 35B-A3B at 4 streams, or a way to keep 4-stream MoE rounds at main's speed. Happy for the pricing fix to be a separate PR that this one waits on. We'll re-derive the stream-tile records with our pinned hipcc when it lands. |
|
Thanks @ashhart. The issue with the data is #1695: the The fix there is the pricing change you describe. HIP timed its shared rounds at one stream and scaled totals past 16 rows linearly, so the cheaper wide rounds made deeper drafts look cheap. The change times them at the served stream count and at the row buckets real rounds pad to. A preliminary run on I'm measuring it now on current |
|
Thanks for landing #1830, @ashhart. GitHub closed this one through #1830, because #1830's body says "the pricing fix #1694 waits on" and GitHub reads "fix #1694" as a closing keyword. That was my wording, sorry. This PR was waiting on #1830, so could you reopen it? I'll then rebase it onto |
|
Reopening: GitHub closed this by mistake through #1830's wording, as you pointed out. Sorry about that. |
What this changes
A lane round of more than four rows ran the stream tile (
zig/kernels/hip/tiles/stream.hpp) as ceil(rows / 4) row blocks on the grid's y. Each later row block read every weight from DRAM again, and each row of a block made the column's code pairs again.The changes:
perf(hip): interleave a stream tile's row blocks so wide rounds reuse cached weightslaunches/stream_layout.zigis the one place that lays a round's row blocks on the grid, andStreamSidesmoves there with it.Kernels.stream_interleave = falseswitches back to the reference layout on y, which the tests compare against.perf(hip): stream tile makes a chunk's code pairs once and takes 4-bit BF16 pairs from bytesv_cvt_f32_ubyteNmakes each code's float in one op, andv_perm_b32packs the top halves. The 4x4 tile's inner loop drops from 1149 to 852 instructions on gfx1151.StreamSidesfield read bystream_tile. On gfx1151 it doubled the stream kernels' VGPRs (4x4: 140 -> 256) and cost about 10% on wide verify rounds.This is the stream-tile work from question 4 of #1563, proposed on its own because it does not depend on DFlash2.
Overlap with open pull requests (
atc overlap):row_blocksfield puts row blocks on the grid's x, and itsHELDvariant makes a round's code pairs once. The two do not merge cleanly, instream.hppandlaunches/affine.zig. I'm happy to align with whichever shape the maintainers prefer; the note above is what I measured for the field variant here.launches/affine.zig(theStreamSidesalias and thestream_layoutimport).Receipt
Environment: Ryzen AI Max+ 395 (gfx1151), Ubuntu 26.04, Zig 0.17.0, Ubuntu clang 21 hipcc (ROCm 7.13 from TheRock for the contract limits), checkpoints Qwen3.8-27B-MLX-4bit and Qwen3.6-35B-A3B-MLX-4bit-MTP.
main22f3d28 (before) with this change on top of it (after).kernel_tests.zig, and main's HIP changes since 01e25d3 (perf(hip): band the prefill WMMA GEMM's blocks so row blocks share weights #1648, perf(hip): keep the matrix tile's inner loop free of selects and divisions #1663 and Flash-Next) don't touch the stream tile. On the rebased tip the tests and checks below pass, and a smoke run matches the receipts: 27B at 16 streams, code 75.2 and chat 72.2 tok/s; 35B-A3B at 4 streams, code 210.0 and chat 196.3 tok/s.token_shais the same as above in every cell, and alone and serial are equal.Exactness:
kernel_tests/affine/layout.zig. Plain, relu2, SwiGLU pair, top-2 routed, routed with an item shorter than the launch's rows, and grouped products run at rounds of 5 to 32 rows in both layouts, and it requires equal bytes. A planted swap of the column and row block indices fails it.kernel_tests/affine/golden.zig. Plain, gate|up pair and top-2 routed launches of 1 to 16 rows must hash to the digest the stream tile wrote when its code pairs came fromstream_pair, before the second change. The digest is recorded per architecture, gfx1151 for now. On another chip the test prints the digest it computed and skips, so a run there can report it.token_shais the same before and after in every run. 27B: code2b28aa2108de, chata06fab348fbe. 35B-A3B MTP: code9f2923ac14d0, chata8c28b8297b4.tools/bench_concurrent.py --alone --serial: all 48 cells of the speed runs below, over both checkpoints, report alone and serial equal, and 0 failed.tf-qwen35-test check: 27B--only invariants42/0/5 and 35B-A3B 46/0/1, the same asmain.Speed, Qwen3.8-27B. Server
--context 4096 --parallel 16 --prompt-cache-gib 0, runs alternated B1 C1 B2 C2.bench_concurrent --tokens 256 --reps 3 --temperatures 0 --levels 1,4,8,16 --alone --serial, aggregate tok/s:tools/bench_openai.py --tokens 256 --reps 3 --temperatures 0, median tok/s:fibonacci-raw13.30, 13.22 before and 13.72, 13.73 after;gpu-chat-no-think13.62, 13.60 before and 14.04, 13.98 after.Speed, Qwen3.6-35B-A3B MTP. Server
--parallel 4, same alternation,--levels 1,4.bench_openai:fibonacci-raw112.6, 112.2 -> 120.4, 120.2;gpu-chat-no-think83.2, 76.6 -> 86.7, 88.2.mainand 12 on the branch (--levels 4, 3 or 5 reps each):main, 6 of 12 on the branch.mainfalls into it too, and two local experiments, neither part of this PR, did not remove it:Config.timed = false), each process stays in one mode;allocate_shared_draftson in HIP's lane facts, it gets more frequent on both sides, with 10 of 15 code reps slow.Kernel contract:
main:SUMMARY gfx1151 49 kernels: 49 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 126 limit problems.hip_records_stream.zig. 109 of them already differed from their records onmainwith this toolchain; the other three (plain 5-bit 2x8, relu2 pair 4-bit 2x8, SwiGLU 8-bit 4x2) matched and now differ. None is a limit violation: 128 threads, as at run time, and no scratch.mainand 473 on the branch. The 253 new ones are the stream-tile launches, about 111 a chip on gfx1030, gfx1100, gfx1151 and gfx1201, now reportedrecord stale.Tests run:
hip-kernel-test(binary run directly): 109 passed, 0 failed, 0 skipped (mainhas 106; the new ones areaffine.layout,affine.goldenand the grid test inlaunches/stream_layout.zig). On the rebased tip: 110 passed, 1 skipped, 0 failed; the skip is main's Flash-Next MTP wide-norm test.test-qwen35-hip-host,hip-host-test,check-native-hip-host: exit 0.tools/zig/lean_check.py: 425 problems onmainand on this branch, none new.zig fmt --checkclean on the changed Zig files.Not run, and why:
stream_pair, now once a chunk.Checklist
"draft": false, a resumed prompt equals a fresh one, and concurrent equals solo. (Bytes unchanged: the layout and golden tests, servedtoken_sha, andbench_concurrentalone and serial equal;tf-qwen35-test checkmatchesmain.)SUMMARYline is in the receipt. (No new symbol; the changed launches' register counts differ from their records, as above.)