Repository navigation
perf(moe): route granite-4.0-h single-token MoE through the fused kernel - #2283
Merged
Merged
Conversation
GraniteMoeHybridMoE::forward always called SwitchGLU::forward plus moe_weighted_sum, the gather_qmm path, so the fused single-token decode MoE kernel (#268, ROCm port #2065) never ran for granite-4.0-h and every decode step paid the fallback (25.4% of decode GPU time on gfx1151 per the #2061 profile). Follow the Mixtral pattern: at one token with MLXCEL_FUSED_MOE on, call SwitchGLU::forward_fused_kernel with the router's top-k indices and f32 softmax scores (which already equal renormalized top-k softmax), and reshape the result to [1, hidden_dim]; on None or more than one token keep the gather path. The shared expert and router are unchanged; forward_fused_kernel already declines unsupported configs. Also expose build_feed_forward/FeedForward/GraniteMoeHybridMoE.switch_mlp at pub(crate) so the granite test module can load a real quantized block, add a Used by list on forward_fused_kernel, and update the #2065 reach statement in the gfx1151 decode profile doc. Validated on gfx1151: the new granitemoehybrid fused-kernel parity test passes, and logit_trace of granite-4.0-h-tiny-4bit with MLXCEL_FUSED_MOE=0 vs default agrees at 0 of 96 positions (0 decided mismatches at gap 2.0). Refs #2245
The fused-kernel parity test added with the routing fix calls SwitchGLU::forward_fused_kernel directly, so it passes without the fix too: the method pre-existed and the routing in GraniteMoeHybridMoE::forward is what issue #2245 changed. Add a reproduction test that drives the production forward at one token and asserts its output is bit-exact the fused-kernel computation and not the gather computation, with both candidates built from the block own router outputs; without the routing it fails at the bit-exact assert, verified by forcing the fused branch off on gfx1151. Also expose GraniteMoeHybridMoE.router and ::forward and GraniteMoeHybridMLP at pub(crate) for the test module; the FeedForward exposure from the routing commit made pub(crate) FeedForward::Dense reach the private GraniteMoeHybridMLP, which cargo build flagged as private_interfaces. Refs #2245
Clippy 1.97.0 denies err().expect() under the workspace -D warnings gate, so verify-clippy-rocm failed on main since fc45284 shipped the yaml parser replacement; the three test call sites move to expect_err, a mechanical rewrite with no behavior change. Carried on this branch because the local gate is the merge gate while hosted CI is down.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
GraniteMoeHybridMoE::forward always ran SwitchGLU::forward plus moe_weighted_sum (the gather_qmm path), so the fused single-token decode MoE kernel (#268, ROCm port #2065) never ran for granite-4.0-h; the #2061 profile measured the fallback MoE at 25.4% of decode GPU time on gfx1151.
What changed
At one routed token with MLXCEL_FUSED_MOE on, GraniteMoeHybridMoE::forward calls SwitchGLU::forward_fused_kernel with the router top-k indices and f32 softmax scores and reshapes the result to [1, hidden_dim]; on None or more than one token the gather path is kept.
The shared expert and router are unchanged; forward_fused_kernel declines unsupported configs itself, so no extra gating.
The Used by list on forward_fused_kernel and the reach statement in the gfx1151 decode profile doc are updated.
A reproduction test drives the production forward at one token and asserts its output is bit-exact the fused-kernel computation and not the gather computation; it fails without the routing and passes with it (verified by forcing the fused branch off on gfx1151).
Test plan
cargo test --profile test-fast --features rocm granitemoehybrid -- --test-threads=1: 12 passed, 0 failed, including the new routing test and the fused-kernel parity test.
logit_trace of granite-4.0-h-tiny-4bit with MLXCEL_FUSED_MOE=0 vs default: 0 of 96 positions mismatched, 0 decided mismatches at gap 2.0.
Guarded decode throughput on gfx1151 (three interleaved runs per arm): fused 105.0/105.7/105.2 tok/s vs fallback 95.9/96.2/95.3 tok/s (1211 to 1220 ms vs 1331 to 1344 ms per 128 generated tokens), about 10% faster and well beyond run-to-run spread, so the default stays on.
make verify-rocm passes locally; hosted CI is down and the local gate is the merge gate. The branch also carries a mechanical expect_err fix for the clippy 1.97 err_expect lint that has been failing verify-clippy-rocm on main since fc45284.
Closes #2245