fix: pad MoE expert GEMMs to the chip's smallest MMA shape - #2
fix: pad MoE expert GEMMs to the chip's smallest MMA shape#2travispchen wants to merge 3 commits into
Conversation
|
@travispchen is attempting to deploy a commit to the Sail Research Team on Vercel. A member of the Team first needs to authorize it. |
Codex Review SummaryThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3b447e54cc
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
…ding
- A100, RTX 4090, RTX 5090, RTX PRO 6000 get [{m: 16, rate: 1}]: mma.sync is
their only tensor-core path and runs at full rate, as on CDNA.
- groupedRows rounds rows up to the chosen shape, so a one-group GEMM
(EP = experts) never prices below the identical dense GEMM.
- Doc: rate is relative to the full-rate shape (stacks on
realizableFlopsFrac); matmulSatRows/realizableFlopsFrac docs mention
mmaShapes.
- Test that every chip's mmaShapes is non-empty with rates in (0, 1].
- Drop the H200 test assertion that pinned the pre-fix model's wrongness.
|
Good work, def really bad if we're reporting "upper bounds" which are below real engine perf. I took this a bit further and made it more comprehensive across chips + not just apply to group gemms in #4 (sorry would have done it on this PR but couldn't edit it, credited it you as an author) |
2 participants
Summary
ChipSpec.mmaShapes: the chip's MMA row shapes and the fraction of peak each sustains.m / rate) instead of matmulSatRows (unset on every NVIDIA/AMD chip, soDEFAULT_MATMUL_SAT_ROWS = 128). Dense GEMMs (groups = 1), N/K padding and chips without mmaShapes are unchanged.- H100/H200: m64 @ 1.0, m16 @ 0.67 (Luo et al. 2024). B200/B300/GB200: m128 @ 1.0, m64 @ 0.5 (PTXtcgen05), so no change. MI300X/MI325X/MI355X: m16 @ 1.0.Why
Compared against SemiAnalysis InferenceX's public single-node rows, HTDYM is a ceiling for LLaMA, R1, Kimi and GLM-5 as intended, but engines beat it by up to 2x on gpt-oss-120b decode. gpt-oss has 128 small experts (top-4), so at 64 concurrent requests each active expert sees ~2 rows, and
tileUtilbills each one a 128-row tile: 256 real rows become 111 x 128 = 14k and decode flips to compute-bound. Real kernels size the tile from the token count: vLLM's Marlin MoE (what InferenceX runs on H200) picks block 8 here from[8, 16, 32, 48, 64](marlin_moe.py#L334-L336) on anmma.sync m16kernel (marlin_mma.h#L39); its tuned Triton config for this shape usesBLOCK_SIZE_M=16up to 256 tokens and 128 only from 2048 (E=128,N=2880 config). 128 is right for dense GEMMs and prefill, wrong as the per-expert floor at decode.H200 TP=1, 1k/1k, conc 64: HTDYM 37.4 ms/step before (compute-bound), 14.0 ms after (memory-bound); InferenceX TRT-LLM measures 22.5 ms and I got 23.7 ms rerunning their
gptoss_fp4_h200.sh(vLLM v0.22.0) on a rented Vast.ai H200. gpt-oss violations go 114 -> 48 (H100/H200/MI300X/MI325X to 0), Kimi 4 -> 0, others unchanged. The 48 left are B200 (preset prices bf16 activations/KV, InferenceX runs MXFP8 + fp8 KV) and 3 memory-bound MI355X rows at <= 1.11x, both separate issues. The rates are from the literature, not fitted.Tests
In
tests/run.test.ts:a grouped GEMM pads each group to the cheapest MMA instruction: Hopper, Blackwell and the no-mmaShapesfallback.gpt-oss-120b decode on one H200 stays under the measured step time: 1k/1k conc 64 decodes under the measured 22.5 ms and memory-bound, and fails both withoutmmaShapes.