[AbrahamCete]

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 weights
  • The row blocks of one column block now run next to each other on the grid's x, so the later ones find its weights in cache.
  • launches/stream_layout.zig is the one place that lays a round's row blocks on the grid, and StreamSides moves there with it. Kernels.stream_interleave = false switches back to the reference layout on y, which the tests compare against.
  • Rounds of 1 to 4 rows are one row block and launch as before.
• perf(hip): stream tile makes a chunk's code pairs once and takes 4-bit BF16 pairs from bytes
  • Each column's 16 code pairs are made once a chunk. Every row of the block then dots them in turn, in the order a row alone takes.
  • On the BF16 build a 4-bit word becomes two byte-per-code words, v_cvt_f32_ubyteN makes each code's float in one op, and v_perm_b32 packs the top halves. The 4x4 tile's inner loop drops from 1149 to 852 instructions on gfx1151.
• Each row's arithmetic and summation order are unchanged, so a row writes the same bytes alone and in a round.
• Tried and dropped: naming the layout in a StreamSides field read by stream_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):

• hip: Qwen3.8-Flash-Next on gfx1151, a reference branch for review (to be split) #1667 (draft) changes the same tile. Its row_blocks field puts row blocks on the grid's x, and its HELD variant makes a round's code pairs once. The two do not merge cleanly, in stream.hpp and launches/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.
• hip: an opt-in prompt GEMM tile that stages dequantized bf16 weights (dq=on) #1650 and hip: a 64-row routed prompt tile for many-expert plans on gfx1151 #1658: a one-line conflict each in launches/affine.zig (the StreamSides alias and the stream_layout import).

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.

  • The receipts compare main 22f3d28 (before) with this change on top of it (after).
  • The branch is now rebased onto 01e25d3 with no conflicts; the 27 commits in between touch no HIP source. On the rebased tip the tests and checks below pass, and a smoke run matches the receipts (27B at 16 streams: code 75.2, chat 71.8 tok/s, the same token_sha).

• Exactness:

  • New GPU test 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.
  • New GPU test 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 from stream_pair, before the second change.
  • Served bytes: token_sha is the same before and after in every run. 27B: code 2b28aa2108de, chat a06fab348fbe. 35B-A3B MTP: code 9f2923ac14d0, chat a8c28b8297b4.
  • 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 invariants 42/0/5 and 35B-A3B 46/0/1, the same as main.

• 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:

  streams	code main	code branch	change	chat main	chat branch	change
  1	13.2, 13.3	13.7, 13.7	+3.4%	13.6, 13.6	14.0, 14.1	+3.3%
  4	43.7, 44.0	45.8, 45.7	+4.3%	43.2, 43.4	45.4, 45.4	+4.8%
  8	47.6, 47.6	62.0, 62.3	+30.6%	46.4, 46.4	60.2, 60.0	+29.5%
  16	50.8, 50.9	74.8, 75.0	+47.3%	49.1, 49.3	71.7, 71.7	+45.7%

  • tools/bench_openai.py --tokens 256 --reps 3 --temperatures 0, median tok/s: fibonacci-raw 13.30, 13.22 before and 13.72, 13.73 after; gpu-chat-no-think 13.62, 13.60 before and 14.04, 13.98 after.

• Speed, Qwen3.6-35B-A3B MTP. Server --parallel 4, same alternation, --levels 1,4.

  • One stream: code 112.7, 112.5 -> 120.5, 120.5 (+7.1%); chat 83.2, 76.6 -> 88.6, 89.0 (+11%). bench_openai: fibonacci-raw 112.6, 112.2 -> 120.4, 120.2; gpu-chat-no-think 83.2, 76.6 -> 86.7, 88.2.
  • Four streams: slower in code on average, the same in chat, and both sides are bimodal. Over 11 servers on main and 12 on the branch (--levels 4, 3 or 5 reps each):
    • code: mean 225.0 -> 214.8 (-4.5%), median 227.0 -> 235.1 (+3.6%);
    • chat: mean 191.3 -> 191.4;
    • servers whose code median fell under 215 tok/s: 1 of 11 on main, 6 of 12 on the branch.
  • The slow mode is the lane core drafting deeper with lower acceptance. In the logs of one pair of servers, a fast request took 115 rounds with 140 of 177-184 drafts accepted. A slow one took 86-106 rounds with 150-169 of 236-262 accepted, and 5.0-5.6 s instead of 4.4-4.5 s.
  • The bytes are the same in both modes, and a mode usually holds for a whole server process. main falls into it too, and two local experiments, neither part of this PR, did not remove it:
    • with the lanes' wall-clock feedback off (Config.timed = false), each process stays in one mode;
    • with allocate_shared_drafts on in HIP's lane facts, it gets more frequent on both sides, with 10 of 15 code reps slow.
  • So the kernel is faster in either mode, but its cheaper wide rounds tip the lane core's depth choice into the slow mode more often. I'll open a separate issue with the data. I'd rather not tune the lane core's pricing in a kernel PR, but I can if you prefer.

• 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.
  • Branch: the same with 129 limit problems. The change moves the VGPR count of all 112 stream-tile launches against hip_records_stream.zig. 109 of them already differed from their records on main with 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.
  • With ROCm 7.13 hipcc: 220 limit problems on main and 473 on the branch. The 253 new ones are the stream-tile launches, about 111 a chip on gfx1030, gfx1100, gfx1151 and gfx1201, now reported record stale.
  • Toolchain: hipcc with HIP 7.1.52801 and Ubuntu clang 21.1.8 for the gates above. Following the answer on Proposal: DFlash2 on HIP for Qwen3.8-27B (gfx1151: 2.0-2.5x one stream, same tokens) #1563, I left every existing record untouched. This PR adds no kernel symbol, so it needs no new record. Re-deriving the records with your pinned hipcc will pick up the new stream-tile counts. cc @BlivionIaG, who recorded these launches.

• Tests run:

  • hip-kernel-test (binary run directly): 109 passed, 0 failed, 0 skipped (main has 106; the new ones are affine.layout, affine.golden and the grid test in launches/stream_layout.zig).
  • test-qwen35-hip-host, hip-host-test, check-native-hip-host: exit 0.
  • tools/zig/lean_check.py: 425 problems on main and on this branch, none new.
  • zig fmt --check clean on the changed Zig files.

• Not run, and why:

  • gfx1100, RDNA2 and gfx1201: I have no card. The layout change is the same on every chip. The byte path for code pairs is on the 4-bit BF16 build only; other widths and the FP16 build keep stream_pair, now once a chunk.
  • Other HIP families' checkpoints: I have none locally. Their stream launches go through the same tile, and the layout and golden tests cover plain, pair, routed and grouped products.
  • 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 layout and golden tests, served token_sha, and bench_concurrent alone and serial equal; tf-qwen35-test check matches main.)
• 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. (No new symbol; the changed launches' register counts differ from their records, as above.)
• Prompt processing is no slower than the last release. (The prompt path is untouched.)
• 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. (Both new tests pin bytes the change must keep rather than a behaviour it adds; the layout test fails on a planted swap of the column and row block indices.)
• 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]

/atc exchange a4436cc5af37c0bc75f72d0c8ffe6b6378e346fad48e70287735aae0dd0fa826