[gilby]

What this changes

When a checkpoint's shared expert has a different bit width or group size from its routed experts, the HIP MoE loads it as its own projections (each tensor at its own width) and runs it beside the routed launch, instead of refusing the checkpoint; a new tf_moe_combine_gate adds its row inside the same fp32 slot accumulation as tf_moe_combine, weighted by the router's shared-expert gate. Matching widths keep the fused stack exactly as before. Fixes #1599.

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). Checkpoints: TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP revision 2b170fa6 (fused path) and an oMLX oQ4e quantization of Qwen3.8-Flash-Next whose shared expert is 8-bit g64 over 4-bit g32 routed experts (separate path). --parallel 1 --context 131072 --prompt-cache-gib 0.
• Exactness, fused path: token_sha256 of tools/bench_openai.py, tools/bench_concurrent.py and both tools/bench_cache.py prompts identical to the base build on the same machine; bench_concurrent --alone --serial every cell alone 3/3 and 12/12 equal, serial equal, 0 failed. Separate path (oQ4e): loads (stock refuses with does not stack with the shared expert), alone 12/12 equal, 0 unequal, outputs deterministic across repeats.
• Decode speed, tools/bench_openai.py --tokens 256 --reps 3 --temperatures 0: fused path base 46.98 -> 46.59 tok/s (code), within run noise on a shared machine; separate path (oQ4e, no baseline exists since stock refuses it) 40.2 code / 37.1 chat tok/s, MTP 3 drafts.
• Prompt speed, tools/bench_cache.py --require-cold on public 7,377- and 29,364-token arrays: fused path unchanged within noise (individual cold cells -1.0% to +0.7%).
• Tests run: zig build test -Dhip=true --summary all 466/466 steps, 0 failed; new host tests in experts.zig (8-bit g64 shared over 4-bit g32 routed loads with both widths kept, rank slicing keeps the shared expert on rank 0 only, a 7-bit shared expert and a wrong shape still refuse) fail on the base with its refusal and pass with the change; zig build hip-kernel-test 15/15 steps; zig build hip-kernel-contract on gfx1151: 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 as the base (the 126 limit problems are live-versus-recorded register counts under this machine's compiler, identical on the base); python3 tools/zig/lean_check.py 425 on base and change; git diff --check clean.
• Not run: tf_moe_combine_gate is not registered with the kernel contract: none of the Flash Next MoE kernels (tf_moe_combine, tf_moe_select, tf_moe_act, tf_moe_route) have contract entries yet, and I would rather agree how to cover that family than add one record alone. 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). Metal and CUDA (HIP-only change); other AMD chips.

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. (tf_moe_combine_gate not yet registered, see above.)
• 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]

Thanks for taking this on — supporting a shared expert at its own width is useful. A few things before it can leave draft:

1. tf_moe_combine_gate is a new kernel, so it needs a kernel-contract record and gate, plus a device test against moe_combine (same bits when the shared row is y's last slot). Since act.hip is a shared image, please also attach SASS/byte-identity receipts for the existing act kernels.
2. core/quant/experts.zig is also used by qwen3_5. With this change a mixed-width qwen3_5 checkpoint would load instead of being refused, but the qwen3_5 HIP bridge never passes separate through, so the shared expert would be dropped and its pick would index past the stack. Could you either refuse there when separate is set, or wire it through the same way as Flash Next?
3. Please add token-equality runs showing existing same-width Flash Next and qwen3_5 checkpoints are unchanged, and a short description with the checkpoint this is for.

[gilby]

Thanks, all three points are addressed in three new commits; receipts on gfx1151 (the only chip I have), TensorFold 7c068895 vs this branch.

1. Contract, device test, act.hip bytes

• 974e039c0 registers tf_moe_combine_gate (contract runner zig/tests/hip/contract_moe_gate.zig, usage records for gfx1030 / gfx1100 / gfx1151 / gfx1201). zig build hip-kernel-contract: SUMMARY gfx1151 52 kernels: 52 pass gates 1-4, 0 fail, 0 known findings; 5 controls: 5 red as planted, 0 missed; 126 limit problems. The 126 are the same live-versus-recorded register differences in existing kernels as on the base (identical set of names); the new record passes its live limits. Shapes at fp32, fp16 and bf16 outputs are bit-equal to tf_moe_combine with the shared row as y's last slot (7 shapes, worst 0.000), and the sums gate is exact at k 11. The other three chips' records come from the compiled code objects (llvm-readelf on each chip's bundle); I can't run them.
• 8f17d1f64 adds the device test hip/kernel_tests/moe_gate.zig: the gate against tf_moe_combine byte for byte at all three output kinds, 1 / 3 / 16 rows, a width that does not fill a block (511) and a single routed slot beside the shared one. The gate's own y carries a poison value (1e30) in the shared slot, so a gate that summed every slot instead of reading the shared operand cannot pass. zig build hip-kernel-test 15/15 steps.
• act.hip byte identity: the gfx1151 code object built from the base and from the branch with the same flags has 19 and 20 kernels; all 19 pre-existing kernels have identical .text (sha256 per kernel), and the only addition is tf_moe_combine_gate. The source diff in act.hip is only the appended kernel.

2. qwen3_5

51e844738 refuses it at the qwen3_5 HIP host index, which both Bridge.init and fromCheckpointRank consume: if experts.routed returns a separate shared expert, it logs <mlp>: a shared expert beside the routed stack is unsupported here and returns error.UnexpectedTensor. A host test in checkpoint_test.zig builds a mixed-width qwen3_5-shaped expert set and checks the refusal (zig build test-qwen35 27/29, 2 skipped).

Your concern was real. I served an oMLX oQ4e quantization of Qwen3.6-35B-A3B (8-bit g128 shared expert over 4-bit g64 routed experts) three ways: the base refuses it with does not stack with the shared expert; this PR without the refusal loaded it and then failed every request with HSA_STATUS_ERROR_MEMORY_FAULT; with the refusal it stops at load with the message above.

3. Token equality and the checkpoint

• qwen3_5, same width: TensorFold/Qwen3.6-35B-A3B-MLX-4bit-MTP revision f84b054c (4-bit g64 experts and shared expert, 8-bit routers), --parallel 1 --context 131072 --prompt-cache-gib 0. Base and branch give the same token_sha256 set across bench_openai, bench_concurrent and both bench_cache --require-cold prompts; bench_concurrent --alone --serial is equal in every cell (3/3 and 12/12, serial equal, 0 failed) on both. Decode 112.5 -> 112.8 tok/s (code), 77.4 -> 79.9 (chat, same tokens).
• Flash Next, same width: TensorFold/Qwen3.8-Flash-Next-MLX-4bit-MTP revision 2b170fa6, identical hashes and equal alone/serial, as in the description.
• The checkpoint this is for: an oMLX oQ4e quantization of Qwen3.8-Flash-Next whose shared expert is 8-bit g64 over 4-bit g32 routed experts. The base refuses it; with this PR it loads, alone 12/12 equal, deterministic across repeats, 40.2 / 37.1 tok/s code / chat with 3 MTP drafts.

Other checks on the branch: zig build test -Dhip=true --summary all 466/466 steps, 0 failed; python3 tools/zig/lean_check.py 425 problems as on the base, none in files this touches; zig fmt --check and git diff --check clean. hip-kernel-contract-limits still stops at tiles/dot2_tiles.hip for gfx1100 on this machine's compiler, same on the base (#1521).