Cache the x casts to fp8 and build ragged layout for fp8 gemm once - #3407
Open
MarioSieg wants to merge 1 commit into
Open
Cache the x casts to fp8 and build ragged layout for fp8 gemm once#3407MarioSieg wants to merge 1 commit into
MarioSieg wants to merge 1 commit into
Conversation
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
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.
Cache the ragged-group layout and x's fp8 cast in the grouped MoE FP8 GEMM path (fp8_grouped_gemm.py, moe.py) instead of rebuilding them from scratch on every w1/w2/w3 call — the layout and x are identical across w1/w3 (and the layout across w2 too), so this was pure duplicated work, verified bit-identical to the old behavior and measured ~15% faster on an isolated MoE-FFN forward at Qwen3-30B-A3B-ish dims. Companion fix to #3388's mxfp8 weight-quant caching, but for the plain (Hopper-compatible) FP8 path.
Note
Medium Risk
Changes the hot MoE FP8 forward path; correctness depends on callers only reusing layout/cache for matching
(offsets, total_m)and tensor pairs, though the PR describes bit-identical outputs when used correctly.Overview
Adds optional caching to the Hopper-compatible FP8 grouped GEMM so MoE forwards do not rebuild the ragged-group layout or re-cast the same activations on every
w1/w3(and layout-only reuse onw2wherexchanges).grouped_fp8_gemmand_GroupedFP8Gemmnow accept optionallayoutandx_fp8_cache, with new helperscompute_grouped_layoutandcast_grouped_input_to_fp8for callers to precompute once. The gated and non-gated grouped-MM FP8 paths inmoe.pycompute layout (andxFP8 cast wherew1/w3share the same input) before the chained GEMMs and pass those into sibling calls.Backward compatibility is preserved when the new arguments are omitted (same behavior as before).
Reviewed by Cursor Bugbot for commit 06b9a89. Bugbot is set up for automated code reviews on this repo. Configure here.