[AbrahamCete]

What this changes

The HIP matrix tile (affine_wmma_gemm, the prompt GEMM on gfx11) spilled to scratch inside its K loop. Two things caused it:

• A group's first product took first ? zero : d[mt][nt], which became a v_cndmask per accumulator element (145 in the loop body).
• The group arithmetic divided by per = group / 32 at run time, about 570 scalar instructions a loop body.

Together with the 64 + 64 accumulators, that put the kernel at 256 VGPRs with 48 bytes of scratch. A group is 32, 64 or 128 codes (groupOk), so per is 1, 2 or 4. Its index and bounds are now a shift and a mask (group_of, group_start, group_end), and the first product calls mma(af, bf, zero) on a uniform branch.

Every sum is the same in the same order, so the bytes are unchanged. The kernel no longer spills: 248 VGPRs, 0 scratch. On Qwen3.8-27B on gfx1151, cold prefill is 24.5-32% faster from 2k to 64k tokens. Decode does not use this tile.

This is independent of #1648 (banded block order). It applies to main alone and on top of #1648 with no conflicts. Together, the seven 27B prompt GEMMs take 60.4 ms, against 63.7 ms with this change alone and 105.8 ms on main. It touches only the default tile's source, so it does not conflict with the opt-in dq tile in #1650.

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), checkpoint Qwen3.8-27B-MLX-4bit. Before is main f350cd3; after is this branch on top of it.

• Exactness:

  • Kernel bytes: I ran the order test from perf(hip): band the prefill WMMA GEMM's blocks so row blocks share weights #1648 (kernel_tests/affine/order.zig) locally on both trees. It requires the matrix tile's output digest to equal the dot2 GEMM tile's on 16 shapes, and the digests are equal on both trees. The shapes are the 27B's prompt projections at 2048 rows, rows off the 128 tile, one row block, one column block, a 35B routed product, and 8-bit.
  • Served bytes: the cold replies' token_sha is identical between 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, run twice on each side:
    • Every cell reports alone equal 3/3 at 1 stream and 12/12 at 4 streams, serial 1/1, and 0 failed.
    • token_sha is the same before and after: code 2b28aa2108de, chat a06fab348fbe.
  • tf-qwen35-test check on this branch: 27B --only invariants 42 passed, 0 failed, 5 skipped; 35B-A3B 46/0/1. Both are the same as main.

• Decode speed before and after: unchanged, because the matrix tile does not run inside lane rounds.

  • 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 decode in tok/s:
    • fibonacci-raw: 13.31, 13.31 before; 13.34, 13.26 after.
    • gpu-chat-no-think: 13.56, 13.62 before; 13.65, 13.65 after.
  • bench_concurrent, aggregate tok/s:
    • code: 13.3 before and after at 1 stream; 43.8-43.9 before and 44.0-44.1 after at 4 streams.
    • chat: 13.6 before and after at 1 stream; 43.1-43.5 before and 43.6 after at 4 streams.

• Prompt speed before and after: cold first request 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, with a fresh server for each run.
  • Runs alternated B1 C1 B2 C2, giving 6 cold samples per cell. Every request reported cached_tokens = 0, and every run qualified.

  prompt	main ttft (s)	branch ttft (s)	change
  2048	8.60 [8.56-8.63]	5.83 [5.81-5.86]	-32.2%
  8192	35.53 [35.35-35.78]	24.39 [24.24-24.48]	-31.4%
  32768	159.40 [158.93-159.72]	114.64 [114.15-114.85]	-28.1%
  65536	364.33 [363.51-364.95]	274.95 [274.17-275.76]	-24.5%

  • The gain shrinks with length because attention, which this change does not touch, takes a growing share of the prompt pass.
  • The prompts are fixed token arrays, made from a corpus of TensorFold docs and server sources and tokenized by the checkpoint. I can attach the arrays, the request bodies and the raw bench_cache JSON.

• GEMM time alone: a local timing test runs each of the 27B's seven prompt GEMMs at 2048 rows, with the weights rotated past 512 MiB so they come from DRAM, and takes the median of 7 batches.

  shape (n, k)	main (us)	branch (us)
  gate_up (17408, 5120)	23929 (15.3 TFLOPS)	14947 (24.4 TFLOPS)
  down (5120, 17408)	28059	15893
  linear qkv (10240, 5120)	14768	8675
  z (6144, 5120)	9102	5351
  out_proj (5120, 6144)	10525	6547
  q gate (12288, 5120)	17866	11215
  kv (1024, 5120)	1572	1047
  total	105820	63675 (-39.8%)

• How I found it:

  • With perf(hip): band the prefill WMMA GEMM's blocks so row blocks share weights #1648's banding, rocprofv3 counters showed the tile was not DRAM bound: every L2-to-memory read was 128 B, the L2 hit rate was 80-92%, and there were no LDS bank conflicts.
  • The kernel metadata and a disassembly of the gfx1151 code object showed the spills and the instructions listed above.
  • On RDNA3 the WMMA and the VALU share the SIMD, so this overhead takes WMMA cycles. Measured with perf(hip): band the prefill WMMA GEMM's blocks so row blocks share weights #1648 applied, the change cuts the VALU instructions per wave on gate_up from 44912 to 27290, and the texture-addresser busy time from 76% to 36%.

• Tests run:

  • hip-kernel-test (binary run directly): 106 passed, 0 failed, 0 skipped, on both trees. With the two local tests (perf(hip): band the prefill WMMA GEMM's blocks so row blocks share weights #1648's order test and the timing test), it is 108/108.
  • test-qwen35-hip-host, hip-host-test, check-native-hip-host: exit 0.
  • Kernel contract: 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, the same on main and on this branch. The only line that differs between runs is where the planted atomic-add control first differs, which changes from run to run.
  • Contract limits with ROCm 7.13 hipcc: 220 problems and the same PASS and FAIL lines on both trees. There is no affine_wmma_gemm record.
  • tools/zig/lean_check.py: 426 problems on main f350cd3 and on this branch, none new.

• Not run, and why:

  • gfx1100: I have no card. The matrix tile runs on every gfx11 WMMA chip (caps.matrix == .wmma11). The 256-VGPR wave limit that caused the spill is the same there, but I have not measured it.
  • 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: matrix-vs-dot2 digests, 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.
• Prompt processing is no slower than the last release. (24.5-32% faster than main.)
• The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not. (gfx1151 only; no Metal, CUDA or gfx1100.)
• New tests fail before the change, pass after it, and skip cleanly without their dependency. (No new test: the change keeps the bytes, which perf(hip): band the prefill WMMA GEMM's blocks so row blocks share weights #1648's order test pins; here it was run locally.)
• 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]

An independent measurement of this PR on another checkpoint. @gilby ran it on Qwen3.8-Flash-Next-MLX-4bit-MTP (uniform 4-bit) on two gfx1151 machines, on a tree with #1650 underneath and dq off, so #1650's tile was not used: #1650 (comment)

• Steady cold prefill against the main tile:
  • 7,377 tokens: 20.76 -> 18.09 and 18.06 s.
  • 29,364 tokens: 91.82 -> 81.81 and 80.96 s.
  • That is about -11% to -13%.
• Token hashes equal main's, bench_concurrent --alone --serial is equal in every cell, and decode is unchanged.
• Gates on that tree: hip-kernel-test 108/108, the kernel contract with the same 126 limit problems as main, and zig build test -Dhip=true 468/468.

Thanks @gilby.

[ashhart]

Thanks @AbrahamCete. This shipped in TensorFold 1.2.0 today, with your credit in the release notes and the changelog: https://github.com/ashhart/TensorFold/releases/tag/v1.2.0

[ashhart]

Thank you @AbrahamCete. This is in TensorFold 1.2.0, released today: https://github.com/ashhart/TensorFold/releases/tag/v1.2.0

It landed with your commits under your own name, so you're listed as a TensorFold contributor, and it's credited to you in the release notes and the changelog.

[AbrahamCete]

Thanks @ashhart, glad it landed. More HIP work for gfx1151 is on the way: #1694 and #1695.