[AbrahamCete]

What this changes

On gfx1151 the prompt GEMM tile (affine_wmma_gemm) holds two blocks a WGP (rocprofv3: 316 resident waves). Its 45,056 bytes of LDS are the limit: LDS goes out in 1 KB granules here, so three blocks need at most 43,008 bytes (a timed probe kernel fits three at 43,008 and two at 43,520). The extra bytes are the padding of the staging rows: each row of 16 words is padded to 20, so that eight 16-byte readers hit eight bank groups. Dropping the padding alone isn't enough: an XOR-swizzled layout came out at 250 VGPRs, which the compiler holds at five waves a SIMD.

This PR adds affine_wmma_skew: the same tile (wmma_gemm_tile with SKEW = true), with unpadded rows where each pair of rows sits 16 bytes on in an 8-row group.

• The eight readers of a fragment still hit eight bank groups, and a fragment is still 32 contiguous bytes at one lane offset plus immediates.
• It takes 39,936 bytes of LDS and 230 VGPRs (217 with ROCm 10.1's hipcc) with no spill, so a WGP holds three blocks (rocprofv3 shows 467 resident waves against 316).
• Only the LDS addresses change. Every sum keeps its operands and order, so the bytes are the padded tile's.

Three blocks a WGP don't win everywhere.

• Deep products lose: a launch's band of 60 blocks outgrows the caches, and the bytes fetched rise about five times.
• On n = 5120 and 2048 rows, the skewed tile is 8-9% faster up to K = 8192 and slower from K = 10240 (+21% at the 27B's down, K = 17408).
• So the registry takes it on gfx1151 for products up to skew_k = 8192 codes deep (fitsPrefillSkew), at cost 9.5, between item64 and matrix. Deeper products keep the padded tile, and other GPUs' tables have no row for it.

Where it loses: shapes with few column blocks (the 35B's n = 2048, k = 4096 and n = 512 products) run 3% to 8% slower on the skewed tile, inside the totals below. A bound on column blocks could keep them on the padded tile; I left that for a follow-up.

It builds on #1658's gfx1151.zon.

• One of hip: a 64-row routed prompt tile for many-expert plans on gfx1151 #1658's selection tests now expects skew past 1,024 rows instead of matrix.
• The first commit is a separate fix: tuning.build's comptime branch quota now grows with a table's rows. The default stopped gfx1151.zon near 30 rows with "evaluation exceeded 1000 backwards branches".

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 (rocprofv3 from ROCm 7.13). Checkpoints: Qwen3.8-27B-MLX-4bit and Qwen3.6-35B-A3B-MLX-4bit-MTP. Base: main 70ede48. I first measured on hip: a 64-row routed prompt tile for many-expert plans on gfx1151 #1658's head e3ea8e7 and rebased once it landed; the served runs below are from that base. The branch is now rebased onto main c2aa0fe (after dots.mocr on HIP: the publisher's BF16 OCR checkpoint, token-exact against its model card #1938), where the exactness, contract and test lines below were run again. Both sides were built from git archive trees.

• tf-qwen35-test prefill, one prefill after a warm-up, alternated B C B C, seconds:

  model, tokens	main	this PR
  27B, 2048	5.269 / 5.282	5.087 / 5.075 (-3.7%)
  27B, 8192	22.149 / 22.108	21.238 / 21.396 (-3.7%)
  35B, 2048	1.001 / 1.006	0.989 / 0.984 (-1.7%)
  35B, 8192	4.639 / 4.645	4.540 / 4.557 (-2.0%)

  On hip: a 64-row routed prompt tile for many-expert plans on gfx1151 #1658's head, three rounds gave the same picture: 27B -3.7% at 2048 and 8192 and -2.5% at 32768; 35B -1.9%, -2.2% and -1.6%.

• Which tile runs (rocprofv3, 2048 tokens): every 27B prompt product except down launches affine_wmma_skew, and down keeps affine_wmma_gemm. The 35B's dense products launch affine_wmma_skew. Its routed products launch affine_wmma_item, because hip: a 64-row routed prompt tile for many-expert plans on gfx1151 #1658's chunk cap doesn't bind there (see my note on hip: a 64-row routed prompt tile for many-expert plans on gfx1151 #1658). Where the skewed tile runs them instead, they take 5.06 / 2.98 ms a call against 5.76 / 3.36.

• Served, on hip: a 64-row routed prompt tile for many-expert plans on gfx1151 #1658's head (the logits hashes are the same on both bases, and the gains above reproduce on main), one fresh server a run, alternated B1 C1 B2 C2:

  • tools/bench_cache.py --require-cold --tokens 16, --context 34816 --parallel 1 --prompt-cache-gib 0, ttft_s:

    model, tokens	hip: a 64-row routed prompt tile for many-expert plans on gfx1151 #1658	this PR
    27B, 2048	5.33 / 5.36	5.15 / 5.21 (-3.2%)
    27B, 8192	22.36 / 22.43	21.40 / 21.55 (-4.1%)
    27B, 32768	103.79 / 104.06	100.58 / 101.07 (-3.0%)
    35B, 2048	1.17 / 1.19	1.17 / 1.17 (-1.0%)
    35B, 8192	5.07 / 5.11	5.04 / 5.05 (-0.8%)
    35B, 32768	26.41 / 26.42	26.13 / 26.14 (-1.0%)

    Every cold request's token_sha is the same on both sides, with 0 cached tokens.

  • tools/bench_concurrent.py --levels 1,4 --tokens 256 --reps 3 --temperatures 0 --alone --serial, --parallel 4: token hashes are equal on every server (35B code 9f2923ac14d0, chat a8c28b8297b4; 27B code 2b28aa2108de, chat a06fab348fbe), 0 failed. Decode is unchanged: code at 4 streams 229.5 / 230.9 on hip: a 64-row routed prompt tile for many-expert plans on gfx1151 #1658 and 230.7 / 231.7 here on the 35B; 43.8 / 43.8 and 44.0 / 44.1 on the 27B.

• Exactness:

  • Prompt logits are byte-identical to main's on both checkpoints, bf16 and --f32, on 70ede48 and again on c2aa0fe (27B dcc85692582d / 232dd4cf1fcd, 35B 18bee11563f1 / 84fbaeb871d7).
  • The order test's ORDER lines: 31 of 31 equal, 15 of them the skewed tile against the dot2 tile.
  • On hip: a 64-row routed prompt tile for many-expert plans on gfx1151 #1658's head, prompt logits at 777, 1,000 and 1,500 tokens, so with partial row blocks, are byte-identical too on both checkpoints, bf16 and --f32. So are the 64-step decode logits.

• Kernel contract, built with ROCm 10.1's hipcc (HIP 7.16.26385) on c2aa0fe: 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. Main has 93 kernels and the same 6 limit problems.

  • affine_wmma_skew at groups 32 and 64 passes nonfinite, repeat, shapes (18 shapes bit-equal to affine_wmma_gemm) and sums.
  • Its records hold that hipcc: 217 VGPRs, 58 SGPRs and 39,936 LDS bytes on gfx1100 and gfx1151; gfx1030 and gfx1201 only carry the symbol, so they are marked not served.
  • With that hipcc the offline limits read FAIL kernel-contract-limits: 22 problems, the same 22 as main. The tile's lines read record current on gfx1100 and gfx1151 and N/A on the other two.

• Tests run:

  • zig build test -Dhip=true --summary all: 570/570 steps.
  • zig build hip-kernel-test --summary all: 132/136 passed, 4 skipped, 0 failed.
  • tools/zig/lean_check.py: 425 problems, the same as main. zig fmt --check is clean.
  • New tests: the skewed tile in the order test, in the routed row sweep and in the kernel contract (contract_wmma_skew.zig); a selection test for skew_k's edges on gfx1151 and for gfx1100 keeping matrix.

• Not run, and why: gfx1030, gfx1100, gfx1150 and gfx1201, because I have no card. gfx1100's table has no row for the tile, and gfx1150 (RDNA3.5, 1024 VGPRs) would read gfx1151's table but hold the same two blocks a WGP.

[ashhart]

Thanks @AbrahamCete, a careful piece of work: the same sums in the same order, one more block a WGP, byte-identical logits, and the 27B's cold first token 3 to 4% sooner, with the depth cut-off measured instead of guessed. Two things before it lands. Main gained #1938's records and contract imports in the same two lists, so it needs a rebase. And our HIP records hold the release hipcc (ROCm 10.1, HIP 7.16), while this tile's come from the system compiler. Please regenerate the four affine_wmma_skew records with the release hipcc and post the offline limits line. A bound that keeps the few-column-block shapes on the padded tile is welcome as a follow-up.

[AbrahamCete]

Thanks @ashhart. Both are done.

• Rebased onto main c2aa0fe. The conflicts were the records list and the contract registrations. hip_records.zig is at the 600-line limit, so the skew records join the list's last line. Prompt logits are still byte-identical to main's on both checkpoints, bf16 and --f32, and the ORDER lines read 31 of 31 equal.
• In de38a34 the four affine_wmma_skew records hold ROCm 10.1's hipcc (HIP 7.16.26385): 217 VGPRs, 58 SGPRs and 39,936 LDS bytes on gfx1100 and gfx1151. gfx1030 and gfx1201 only carry the symbol, so they are marked not served, as hip_records_qwen3.zig marks the chips it doesn't serve. With that hipcc:
  • hip-kernel-contract-limits: FAIL kernel-contract-limits: 22 problems, the same 22 as main. The tile's lines read record current on gfx1100 and gfx1151 and N/A on the other two.
  • hip-kernel-contract on 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. Main has 93 kernels and the same 6.

I've updated the description to match. The column-block bound will come as a follow-up.

One more thing that hipcc showed: 8 of main's 22 are the planned-walk records of tf_sum_partials and tf_sum_partials_act from #1961, which still hold the system compiler's values. I can re-derive them in a separate PR if that helps.

[ashhart]

Thanks @AbrahamCete, both done: the rebase keeps logits byte-identical, and the skew tile's records now read current against the release hipcc. This is ready to land. And yes please to the separate PR re-deriving #1961's tf_sum_partials records with that hipcc. We don't have one here, so it's very welcome.