Skip to content

Cache the x casts to fp8 and build ragged layout for fp8 gemm once - #3407

Open
MarioSieg wants to merge 1 commit into
mainfrom
mario/optimize-fp8
Open

Cache the x casts to fp8 and build ragged layout for fp8 gemm once#3407
MarioSieg wants to merge 1 commit into
mainfrom
mario/optimize-fp8

Conversation

@MarioSieg

@MarioSieg MarioSieg commented Aug 26, 2026

Copy link
Copy Markdown
Member

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 on w2 where x changes).

grouped_fp8_gemm and _GroupedFP8Gemm now accept optional layout and x_fp8_cache, with new helpers compute_grouped_layout and cast_grouped_input_to_fp8 for callers to precompute once. The gated and non-gated grouped-MM FP8 paths in moe.py compute layout (and x FP8 cast where w1/w3 share 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.

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.

1 participant