perf(moe): route granite-4.0-h single-token MoE through the fused kernel #2245

Part of #1801

Problem / Background

The fused decode MoE kernels (Metal/CUDA since #268, ROCm port in #2065 / PR #2098) never run for granite-4.0-h (granitemoehybrid). #2061's profile (docs/benchmark_results/rocm-decode-profile-gfx1151-2026-09-30.md:98, :114, :124, :147) measured granite-4.0-h-tiny's fallback MoE at 25.4% of decode GPU time on gfx1151 and named wiring granite to forward_fused_kernel as the step that would reach it.

Current Behavior

Verified on origin/main 1b4e3657, after the epic #2166 refactor (last touched by d54b209b, #2205). GraniteMoeHybridMoE::forward (src/models/granitemoehybrid.rs:775-810) always calls self.switch_mlp.forward(&x_flat, &topk_indices) followed by moe_weighted_sum, which is the gather_qmm path. forward_fused_kernel has 19 call sites in other MoE models (for example mixtral.rs:126-150, qwen3_moe.rs:219-240) and none in granitemoehybrid.rs. The same holds on Metal and CUDA, since this is model code.

Proposed Solution

In GraniteMoeHybridMoE::forward, follow the Mixtral pattern exactly: when array_shape(&x_flat)[0] == 1 && crate::models::switch_layers::fused_moe_enabled() (switch_layers.rs:88), call self.switch_mlp.forward_fused_kernel(&x_flat, &topk_indices, &scores) (switch_layers.rs:1202) and reshape the result to [1, hidden_dim]; on None, or for more than one token, keep the current forward + moe_weighted_sum path. scores stays as computed today (f32 softmax over the top-k logits, which equals renormalized softmax); the launcher casts scores to f32 itself (src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp:3223), so no cast is needed. The shared expert (GraniteMoeHybridSharedMLP) and the router are unchanged. forward_fused_kernel already declines unsupported configs (non-affine, unsupported bits, Dff above the backend cap, missing biases), so no extra gating is needed. Update the // Used by: list on forward_fused_kernel and the statement at rocm-decode-profile-gfx1151-2026-09-30.md:98 that granite is not reached.

Acceptance Criteria

• A test in src/models/granitemoehybrid_tests.rs, modeled on the fused-geometry test in src/models/qwen3_moe_tests.rs:360-410, builds a small quantized GraniteMoeHybridMoE, checks one-token output on the fused path against forward + moe_weighted_sum within the tolerance that test uses, and skips with a message when fused_moe_kernels_available() is false.
• Greedy decode of granite-4.0-h-tiny-4bit on gfx1151 and on Metal: logit_trace with MLXCEL_FUSED_MOE=0 vs default gives 0 decided mismatches (compare_logit_traces.py --decided 2.0).
• Guarded decode throughput on gfx1151 (scripts/rocm_gpu_guard.sh -- scripts/bench_decode.sh, granite-4.0-h-tiny-4bit, at least three interleaved runs per arm) is reported for both arms; the default stays on only if the fused arm is not slower beyond run-to-run spread. Report Metal decode the same way.
• A rocprofv3 trace shows the fused MoE kernels and no gather_qmv_* MoE calls on granite decode.

Verification

cargo test --profile test-fast --features rocm granitemoehybrid -- --test-threads=1
make verify-rocm
make verify   # on a Mac, for the Metal arm

Related: #2061, #2065, PR #2098, #1814, #1685, #2166.