Skip to content

hip: decode router over the bf16 rows the checkpoint holds - #1614

Draft
gilby wants to merge 2 commits into
ashhart:mainfrom
gilby:strix/hip-router-bf16
Draft

gilby wants to merge 2 commits into
ashhart:mainfrom
gilby:strix/hip-router-bf16

Conversation

@gilby

@gilby gilby commented Oct 29, 2026

Copy link
Copy Markdown
Contributor

What this changes

The HIP decode router reads the bf16 router rows the checkpoint already holds instead of the fp32 copy widened at load, keeping each lane's multiply-add order, so its logits are byte-identical; it is registered as a replacement of tf_router_decode_bf16 / tf_router_decode_f16. Part of #1603 (the router was the largest of the small per-round decode kernels; the hyper-connection and select/combine fusions measured smaller and are not in this pull request).

Receipt

  • Environment: TensorFold 7c068895 plus this change, tensorfold-native 1.1.8 built with zig build native-hip -Dhip-arch=gfx1151, Ubuntu archive HIP 7.1 (clang 21), AMD Ryzen AI MAX+ 395 with Radeon 8060S gfx1151 (128 GB unified memory). Checkpoint TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP revision 2b170fa6, MTP 3 drafts, --parallel 1 --context 131072 --prompt-cache-gib 0, base and candidate on the same machine.
  • Exactness: token_sha256 of tools/bench_openai.py, tools/bench_concurrent.py and both tools/bench_cache.py prompts identical to the base build in two candidate runs; bench_concurrent --alone --serial every cell alone 3/3 and 12/12 equal, serial equal, 0 failed.
  • Decode speed, tools/bench_openai.py --tokens 256 --reps 3 --temperatures 0, median tok/s code / chat: base 50.03 / 43.93; this change 51.25 / 45.08 and 51.25 / 45.09 (two server starts), about +2.5%. Router kernel alone, 48 layers of 513 x 2560 rows at 4 rows: 49.8-50.2 us -> 16.1-16.5 us.
  • Prompt speed: the change only takes the decode-window router path (1 to 16 rows); tools/bench_cache.py --require-cold first-request token hashes are identical and cold timings unchanged within noise.
  • Tests run: zig build hip-kernel-contract on gfx1151: SUMMARY gfx1151 51 kernels: 51 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 126 limit problems (base: 49 kernels, the same 126 limit problems, which are live-versus-recorded register counts for existing kernels under this machine's compiler); the router lines: shapes bit-equal to the replaced kernel at 8 shapes (1, 4 and 16 rows x 513 x 2560 and others), exact sums at k 2560, live limits 35/36 VGPRs and 0 LDS as recorded. zig build hip-kernel-test 104/104 pass, including a new test comparing the bf16-rows and fp32-rows routers byte for byte for both activation kinds at 1-16 rows with planted -0, inf and NaN inputs (fails against a deliberately reordered sum). zig build test -Dhip=true --summary all 466/466 steps, 0 failed. python3 tools/zig/lean_check.py 425 on base and change. git diff --check clean.
  • Overlap: it adds entries to the same HIP launcher and record lists as hip: load a shared expert wider than the routed experts beside the stack #1612 (shared expert) and the sparse prompt attention pull request; those are list additions that rebase trivially.
  • Not run: hip-kernel-contract-limits stops on this machine's compiler at tiles/dot2_tiles.hip for gfx1100 (branch size exceeds simm16, same on the base, see HIP: dot2_tiles.hip fails to compile for gfx1100 with Ubuntu clang 21.1.8 ("branch size exceeds simm16") #1521); the new records for gfx1030, gfx1100 and gfx1201 come from each chip's compiled code object metadata. Metal and CUDA (HIP-only change); other AMD chips at runtime.

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.
  • No precision traded for speed. If the bits change, the description says which sums change and why.
  • 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.
  • The description says which platforms I ran, Metal M1 to M5 and CUDA, and which I could not.
  • 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

ashhart commented Oct 30, 2026

Copy link
Copy Markdown
Owner

Thanks — this is well put together: matching the fp32 router's lane order and gating it bit-for-bit is exactly the right approach. To leave draft, could you attach the device receipts: the contract run for the two new router records and the decode kernel test on each arch the records list, plus SASS/byte-identity for the existing decode-image kernels and a Flash Next token-equality run with the new path on? A one-line note on how much extra memory the bf16 router copy adds per layer would help too. Heads-up: #1622 touches the same lines in hip_records.zig and contract_main.zig, so whichever lands second will need a small rebase.

@gilby

gilby commented Oct 30, 2026

Copy link
Copy Markdown
Contributor Author

Thanks. Receipts, all on gfx1151 (the only chip I have), TensorFold 7c068895 vs this branch:

  • zig build hip-kernel-contract: SUMMARY gfx1151 51 kernels: 51 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 126 limit problems (base: 49 kernels, the same 126 live-versus-recorded register differences for existing kernels under this machine's compiler). Both new records pass every gate: bit-equal to the router they replace at 8 shapes, five repeats byte-identical, 35 / 36 VGPRs live as recorded.
  • zig build hip-kernel-test: all 104 pass, including the new test that the router over the stored bf16 rows writes the fp32-rows router's bits for both activation kinds at 1-16 rows.
  • Byte identity of the existing decode image: all 7 pre-existing kernels in the decode code object are identical in bytes and llvm-objdump -d between base and branch, including the two replaced routers (which stay selectable); the branch adds exactly tf_router_bf16 and tf_router_bf16_f16.
  • The records' other chips (gfx1030, gfx1100, gfx1201): compile-only here. decode.hip builds for all four with each chip's caps, and each code object carries both new kernels; I have no device to run them on.
  • Flash Next token equality with the new path on: base twice and branch twice give identical token_sha256 sets across bench_openai, bench_concurrent and bench_cache, and bench_concurrent --alone --serial is equal in every cell. Decode 49.85 -> 51.25 tok/s (code) and 43.85 -> 45.08 (chat), prompt timing unchanged.
  • Memory: the bf16 router rows are a new device upload (base only uploaded the widened fp32 copy): 513 x 2560 x 2 B = 2.5 MiB a layer, about 120 MiB over 48 layers. If that matters I can drop the fp32 copy where only the decode window reads it, but the prompt path still uses it today, so I left both.

#1622 touches the same lines in hip_records.zig and contract_main.zig; happy to rebase whichever lands second.

@AbrahamCete

Copy link
Copy Markdown
Contributor

I ran this on Qwen3.6-35B-A3B-MLX-4bit-MTP (gfx1151, the same Ubuntu HIP 7.1 and clang 21), cherry-picked onto main 9e8c368. The only conflict was the paragraph #1788 added at the same spot in docs/kernel-contract.md; both paragraphs stay.

As it stands, the PR wires rows16 only in the Flash-Next loader, so Qwen3.5/3.6 still runs the fp32-rows router. One line in zig/src/families/qwen3_5/backend/hip/model/bridge.zig, next to .rows32, turns it on:

            .rows16 = r.router.ptr,

Qwen already uploads the bf16 router rows, so this adds no memory. core/quant/experts.zig routed() widens rows32 from those same bf16 values. With that line, against main:

  • Exactness:

    • 64-step decode logits (tf-qwen35-test logits --decode) are byte-identical.
    • tf-qwen35-test lanes replies with MTP are identical, both for 4 concurrent streams and solo.
  • Router kernel, from rocprofv3:

    • One row: 21.8 → 17.1 µs (120 → 40 VGPRs, 24 B of scratch → 0).
    • 4 streams with MTP, up to 16 rows: 43.2 → 29.1 µs.
  • tf-qwen35-test check --only speed, alternated twice, tok/s:

    case main this PR + the line
    1 stream, drafts off 70.5 / 70.5 71.5 / 71.4
    1 stream, drafts on 95.8 / 95.9 97.2 / 96.8
    4 streams, drafts off 140.0 / 138.9 140.4 / 140.1
    4 streams, drafts on 128.7 / 126.0 128.6 / 128.1
  • Tests:

    • hip-kernel-contract: SUMMARY gfx1151 54 kernels: 54 pass gates 1-4, 0 fail, 2 known findings; 5 controls: 5 red as planted, 0 missed; 130 limit problems. That is main's 52 plus your two router records; main has the same 130 limit problems on this toolchain. All eight router lines pass.
    • hip-kernel-test: 110 passed, 1 skipped, 0 failed. hip-host-test: 55/55.

One small thing in zig/tests/hip/contract_router.zig plant(): it writes the second expert's row (self.w[1 * k + k / 2]) without an e > 1 guard. The next plant has e > 2, and prepare accepts n == 1. It's harmless as long as gate 1 never draws a one-expert shape.

Happy to send the Qwen line as a separate PR after this lands if you'd rather keep this one to Flash-Next.

The decode-window router read the router weights as fp32 rows widened at
load. tf_router_bf16 (and its fp16-activation form) reads the bf16 rows
the checkpoint already stores, prefetching four groups, with each lane's
multiply-add order unchanged, so its outputs equal the fp32-rows router's
bytes. Registered with the HIP kernel contract as that router's
replacement.

Refs ashhart#1603
@gilby
gilby force-pushed the strix/hip-router-bf16 branch from df85c75 to 1558122 Compare November 2, 2026 06:05
@gilby

gilby commented Nov 2, 2026

Copy link
Copy Markdown
Contributor Author

Thanks @AbrahamCete, that's a thorough run, and good to see the same bits and the router time on Qwen3.6-35B-A3B too. I'd keep this PR to Flash Next, so please do send the Qwen line as its own PR once this lands; your measurements there make the case on their own.

You're right about plant(): 1af9dab35 writes the second expert's row only when there is one (e > 1), like the third plant. On the rebased head (9e8c368b, gfx1151): SUMMARY gfx1151 54 kernels: 54 pass gates 1-4, 0 fail, 2 known findings; 5 controls: 5 red as planted, 0 missed; 130 limit problems (main: 52 and the same 130), both router records current on all four chips, hip-kernel-test 110 passed, 1 skipped, host tests 485/485, lean 424 as on main; tokens identical to main and decode 51.0 / 44.9 tok/s code / chat against main's 49.9-50.1 / 43.9-44.0.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants