[AbrahamCete]

What this changes

A prompt projection on HIP wrote its fp32 sum into an arena buffer, then ran tf_cast to round it to the activation type. Nothing else read the fp32 buffer (affine() in zig/src/hip/ops/project.zig). On Qwen3.8-27B that is 496 extra launches in a 2048-token prompt pass, about 4% of its kernel time.

The changes:

• The prompt tiles store through a new epilogue, GemmOut<T>. These are affine_wmma_gemm, affine_gemm_block, affine_gemm_kp, and the affine_dot2_block reference.
  • When out16 is set, GemmOut<T> rounds into it with narrow<T>, which is float_to_bf16/__float2half_rn bit for bit, NaN to 0x7fc0 included.
  • Otherwise it writes fp32 as F32Out did.
• The plain prompt entries declare rounds = .always, so affine() asks a prompt for out16 the way decode already does. The cast launch is gone.
• Routed prompt products keep fp32, because the expert combine reads it.
• Decode entries and Kernels.run are unchanged.

The bytes are unchanged: the tile hands narrow<T> the same fp32 value that tf_cast read back. There is no new kernel symbol.

This is independent of #1648 and #1663. It conflicts trivially in two places:

• with hip: an opt-in prompt GEMM tile that stages dequantized bf16 weights (dq=on) #1650, in the prefill() entry list;
• with hip: Qwen3.8-Flash-Next on gfx1151, a reference branch for review (to be split) #1667, in the same one-line change to half in affine(). hip: Qwen3.8-Flash-Next on gfx1151, a reference branch for review (to be split) #1667 makes that change for its dense tile, and this PR makes it for every prompt tile.

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 ea3567b (before) with this change on top of it (after).
  • The branch is now rebased onto c3c3bd0. The only conflict was the import list in kernel_tests.zig. main now includes perf(hip): band the prefill WMMA GEMM's blocks so row blocks share weights #1648 and perf(hip): keep the matrix tile's inner loop free of selects and divisions #1663, which change the same prompt GEMMs, so I measured again on the rebased tip against c3c3bd0:
    • 27B prompt logits (2,048 tokens) are byte-identical to main, in bf16 and with --f32.
    • Cold ttft_s, runs alternated B1 C1 B2 C2, token_sha identical at every size: 2048 5.52 -> 5.38 s (-2.5%), 8192 23.16 -> 22.53 s (-2.8%), 32768 109.40 -> 107.22 s (-2.0%).
    • hip-kernel-test 108 passed, 1 skipped, 0 failed (main: 107 passed, 1 skipped; the skip is main's Flash-Next MTP wide-norm test). Kernel contract SUMMARY gfx1151 50 kernels: 50 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 126 limit problems, limits 220 and lean_check 424, each the same as main. check: 27B --only invariants 42/0/5, 35B-A3B 46/0/1.

• Exactness:

  • New kernel test kernel_tests/affine/round.zig. Every prompt entry that takes each of 12 shapes runs with both out and out16 valid, 63 launches in all:
    • shapes: 1, 2, 3, 4, 17, 32, 33, 100 and 300 rows; group 32 and 128; 8-bit.
    • It requires tf_cast's bytes of the same entry's fp32 sum in out16, with out untouched.
    • It fails before the change (round 1 row: mlx.project.prefill.kp0 with out16 writes other bytes than tf_cast of its fp32 sum) and passes after.
    • The existing affine test now also runs the prompt entries in rounded mode. It checks the fp32-sum error bound and that every prompt entry writes the same bytes.
  • Prompt logits (tf-qwen35-test logits) are byte-identical before and after, in bf16 and with --f32: 27B on 300 and 2048 prompt tokens, and 35B on 300.
  • Planted control: rounding by truncation in GemmOut fails the new test and two existing affine error checks, and changes the 27B logits.
  • Served bytes: the cold replies' token_sha is identical 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 and 12/12, serial 1/1, and 0 failed.
    • token_sha is the same before and after: code 2b28aa2108de, chat a06fab348fbe.
  • tf-qwen35-test check: 27B --only invariants 42/0/5 and 35B-A3B 46/0/1, the same on both trees.

• Decode speed before and after: unchanged, since decode already rounded in its tiles.

  • 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 tok/s:
    • fibonacci-raw: 13.21, 13.29 before; 13.25, 13.27 after.
    • gpu-chat-no-think: 13.52, 13.58 before; 13.49, 13.61 after.
  • bench_concurrent aggregate at 4 streams: code 43.8-43.9, chat 43.2-43.4, on both sides.

• 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.59 [8.55-8.64]	8.26 [8.24-8.28]	-3.8%
  8192	35.48 [35.33-35.66]	34.25 [33.98-34.51]	-3.5%
  32768	159.35 [158.42-159.97]	154.47 [154.05-154.91]	-3.1%
  65536	365.03 [364.57-365.56]	355.25 [353.16-357.58]	-2.7%

  • tf-qwen35-test prefill <27B>: 2048 tokens 8.50 to 8.11 s, and 8192 tokens 35.29 to 33.74 s.
  • 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.

• Tests run:

  • hip-kernel-test (binary run directly): 107 passed, 0 failed, 0 skipped (main has 106; the new one is affine.round).
  • 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. The only line that differs between runs is where the planted atomic-add control first differs.
  • Contract limits with ROCm 7.13 hipcc: the same PASS and FAIL lines on both trees. The prompt tiles have no contract records.
  • 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 fp16 path rounds with __float2half_rn on both sides, which I checked by reading the code.
  • Other HIP families' checkpoints: I have none locally. affine() is their prompt path too, and the change keeps the bytes for any tile, as the kernel test and the logits show here.
  • 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 round test, the prompt logits, 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 prompt tiles have no records; SUMMARY unchanged.)
• Prompt processing is no slower than the last release. (2.7-3.8% faster than main ea3567b; 2.0-2.8% faster than main c3c3bd0 after the rebase.)
• 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.
• 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.

[ashhart]

Thanks @AbrahamCete, the round test and its truncating control are what we'd want. Two things before it lands:

1. Could you rebase onto current main? zig/src/hip/kernel_tests.zig conflicts on one line with your perf(hip): band the prefill WMMA GEMM's blocks so row blocks share weights #1648's affine/order.zig import; both lines stay.
2. Your prefill table is against ea3567bb, before the banding landed. Could you post the round and order tests, the 27B prompt logits (byte-identical against main), and a cold-prefill pair at 2K and 32K on the rebased head against current main? We'd like the gain with both in. A gfx1100 or gfx1201 run would cover the fp16 path if you have one, but it isn't required.

[AbrahamCete]

Thanks @ashhart. The branch is rebased onto c3c3bd0 (both import lines stay), and the description's receipt now has the numbers against current main. On the rebased head 5986bf6 against c3c3bd0, gfx1151:

• affine.order: 16/16 digests equal (ORDER 27b gate_up m2048 n17408 k5120: 92c5bbcddf208589 92c5bbcddf208589, and the same for down, linear qkv, z, out_proj, q gate, kv, b a, the edge shapes, the 35B routed gate_up and 8-bit).
• affine.round: ROUND 63 prompt launches write tf_cast's bytes into out16.
• 27B prompt logits, 2,048 tokens: byte-identical to main in bf16 and with --f32.
• Cold prefill, tools/bench_cache.py --require-cold --tokens 16, runs alternated B1 C1 B2 C2, the same token_sha on both:
  • 2K: 5.52 -> 5.38 s (-2.5%)
  • 32K: 109.40 -> 107.22 s (-2.0%)
  • (8K: 23.16 -> 22.53 s, -2.8%)
• hip-kernel-test: 108 passed, 1 skipped, 0 failed; the skip is main's Flash-Next MTP wide-norm test, which skips on main too. Kernel contract, limits and lean_check are the same as main.

I don't have a gfx1100 or gfx1201 card. I can open an ATC exchange test request for the fp16 path if that helps.

[ashhart]

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.

[AbrahamCete]

Thank you @ashhart!

[ashhart]

Thanks @AbrahamCete, that answers both: the rebase keeps #1648's import, and the round and order tests, the 27B logits and the cold-prefill pair on the rebased head are all against current main. 2-3% off cold prompts with the same bytes is a clean win. We'll land it, with gfx1100 and gfx1201 listed as not run. No need for an exchange request for the fp16 path. If someone with one of those cards runs hip-kernel-test later, we'll add it to the notes.

[ashhart]

Thanks @AbrahamCete! HIP prompt projections that round to the activation type in their tile are in 1.2.1: https://github.com/ashhart/TensorFold/releases/tag/v1.2.1