Repository navigation
MoE: native FP8 QAT for routed-expert GMMs (FP8 forward + FP8-weight backward dLHS) - #5969
shuningjin wants to merge 3 commits into
Conversation
Codecov Report❌ Patch coverage is 📢 Thoughts on this report? Let us know! |
…backward dLHS), and fix the dropless full-buffer gather index Adds a native-QAT path on top of fake QAT (fp8_moe_fake_quant + fp8_fake_quant_hoist=gather8): - fp8_moe_native_fwd: the wi/wo forward GMMs run FP8 x FP8 on the MXU from the gathered FP8 weights and per-channel scales; gmm_v2 quantizes the activation in the kernel (fp8_moe_native_fwd_lhs_quant, fp8_moe_native_fwd_acc_dtype). - fp8_moe_native_bwd_dlhs: the backward dLHS GMMs read the FP8 weights and dequantize each tile in VMEM, fused into the first matmul of each block (rhs_contract_dequant); bit-identical to the XLA dequant backward. Unsupported tiles fall back to the XLA dequant and log a warning. - Fp8NativeWeight(anchor, codes, scale) carries the weight through shard_map; the weight gradient is returned as the anchor's cotangent (straight-through into the master weight). - A RoutedMoE call with fp8_moe_native_fwd but no gather8 absmax (outside the layer scan) raises instead of silently running fake QAT. - ragged_sort: clamp the masked (-1) indices of the full-buffer ring gather-reduces; the unclamped rows halted the SparseCore on dropless replay. The clamp applies on every full-buffer ring gather-reduce, not only with moe_routing_index_math=count; it leaves outputs unchanged. - gmm_v2: new lhs quantization kwargs (lhs_quant_block_size, lhs_f32_scale, lhs_hoist_quant, rhs_scale_first) and GmmConfigs fields; rejects a negative lhs_quant_block_size. - Tests: fp8_moe_native_test (flags, config validator, outside-scan guard, dequantize_fp8_codes bit-identical to dequantize_weight_codes_ste, lhs_quant config values match ops._LHS_QUANT_MODES); ragged_sort_full_buffer_indices_test (cpu_only: it passes unclamped indices on purpose). All new flags default to off; with them off, the lowered program is unchanged (lowered TPU StableHLO + Mosaic kernels compared: 12 GMM-op programs and 3 model fwd+grad programs identical).
2593b30 to
0bf1938
Compare
shuningjin
changed the title
Summary of ChangesHello, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request implements a native FP8 Quantization-Aware Training (QAT) path for Mixture-of-Experts (MoE) GMMs, significantly enhancing training throughput by leveraging TPU7x hardware capabilities. By processing FP8 weights directly and performing in-kernel activation quantization, the changes reduce HBM bandwidth usage and increase computational efficiency. Additionally, the PR addresses a pre-existing bug in the dropless ring gather-reduce logic, ensuring index safety and preventing SparseCore crashes during micro-batch replays. Highlights
New Features🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here. Footnotes
|
| # gmm_v2 accumulator across K blocks for fp8_moe_native_fwd: "" (kernel default, bfloat16), "bfloat16" or "float32". | ||
| fp8_moe_native_fwd_acc_dtype: "" | ||
| # In-kernel lhs quantization scheme for fp8_moe_native_fwd ("" = default block512_bf16; see types.py). | ||
| fp8_moe_native_fwd_lhs_quant: "" |
There was a problem hiding this comment.
Could the recommended settings be the defaults here, fullk_f32 and float32?
setting only fp8_moe_native_fwd: true gives a bf16 scale per 512-column block and bf16 accumulation. Your numerics table shows that combination is further from the FP8 rollout than fake QAT (0.0637 vs. 0.0468), while fullk_f32 with f32 accumulation gets 0.0131. So the simplest config is the least accurate one, and you only get the good numerics if you know about two extra flags.
Both flags are only read when fp8_moe_native_fwd is on, so the flags-off path doesn't change. The types.py defaults and test_flags_are_defined_and_off_by_default would need the same update.
There was a problem hiding this comment.
Thanks, good point. In 10fdc4f I removed both knobs instead: fp8_moe_native_fwd: true now always uses per-row full-K f32 scales and f32 accumulation, which is both the closest to the rollout kernel (0.0131 vs 0.0637) and the fastest. The other modes remain ops-level arguments for benchmarks.
The native forward now also requires fixed tiles, since the heuristic tiling budgets a bf16 accumulator and could overflow VMEM with f32. Lowered programs are unchanged vs the old explicit config; tests updated.
| else: | ||
| tile_k, tile_n = tiling[4], tiling[5] | ||
| size_n, size_k = rhs_shape[1], rhs_shape[2] | ||
| supported = tile_n % 32 == 0 and tile_k % 128 == 0 and size_n % tile_n == 0 and size_k % tile_k == 0 |
There was a problem hiding this comment.
maybe this check also needs "one K tile", i.e. tile_k == size_k.
The kernel's VMEM cache holds one weight block, keyed by (expert, n tile, k tile), and the grid loops over K innermost (n → m tile → k). With two or more K tiles, each step evicts the block the next step needs. The cache then never hits, and every step dequantizes its weights again. The _rhs_cache_miss docstring already assumes a single K tile, but nothing enforces it.
The MLPerf recipe isn't affected, because its tuned tiles give one K tile for both wi and wo. But a config with wi_tile_dlhs_mlp_dim=1024 (the recipe's value before the tuned override) would silently take this path with 2 K tiles, since the fused wi has K = 2 × 1024. If multi-K-tile support matters, the alternative is to cache the whole [tile_n, size_k] block, keyed by (expert, n).
There was a problem hiding this comment.
Good catch, thanks! Agreed the cache never hits with more than one K tile (n tiles are fine; the grid is (n, gm, k)). I measured it on TPU7x at our shapes before gating, and the miss path is still 9-16% faster than the XLA-dequant fallback, bit-identical, since it reads FP8 instead of a bf16 copy:
| DLHS | tiles (m, k, n) | K tiles | fallback | in-kernel |
|---|---|---|---|---|
| wi | 256, 1024, 4096 | 2 | 5.01 ms | 4.20 ms |
| wi | 256, 512, 4096 | 4 | 5.29 ms | 4.84 ms |
| wo | 512, 2048, 1024 | 2 | 2.11 ms | 1.82 ms |
| wo | 512, 1024, 1024 | 4 | 2.38 ms | 2.14 ms |
So I kept the path and clarified the _rhs_cache_miss docstring in 10fdc4f. Caching the full [tile_n, size_k] block would need 16 MiB of VMEM for wi, so I left that out.
| return pos.reshape(-1) | ||
|
|
||
|
|
||
| def _clamp_masked_indices(indices, valid_rows_mask): |
There was a problem hiding this comment.
nit: suggest taking this index fix into its own PR first? It could land right away instead of waiting on the kernel review, and if the native path ever needs a revert, the crash fix stays in.
…, multi-K DLHS doc; split out the ragged_sort fix - Remove fp8_moe_native_fwd_lhs_quant and fp8_moe_native_fwd_acc_dtype: fp8_moe_native_fwd always quantizes the activation with one float32 absmax scale per row over the full K dim and accumulates in float32 (the setting that measured closest to the rollout's FP8 MoE kernel and fastest); the other modes stay ops-level arguments for benchmarks and tests. Lowered programs are unchanged vs the old explicit fullk_f32 + float32 config. - fp8_moe_native_fwd requires use_gmm_v2_heuristic_tiling=False: the heuristic tiling budgets a bf16 accumulator when the lhs is quantized in the kernel, so with the float32 accumulator its tiles could exceed VMEM. - _rhs_cache_miss: document that the DLHS dequant cache only hits with a single K tile; with more, each step dequantizes inside its matmul, unchanged outputs, still faster than the XLA dequant fallback. - Move the ragged_sort full-buffer index clamp and its test to a separate PR.
shuningjin
changed the title
| # enough to minimize compute overhead of quantization. | ||
| quant_block_size=512, | ||
| quant_block_size=( | ||
| dims.size_k if lhs_quant_block_size == 0 else (lhs_quant_block_size or _DEFAULT_LHS_QUANT_BLOCK_SIZE) |
There was a problem hiding this comment.
Here lhs_quant_block_size == 0 sets quant_block_size = dims.size_k, but k_ranges in inner_kernel (line 926) only iterates within a single tile_k (range(0, cfgs.tiles.tile_k, q_block_size)).
if someone enables fp8_moe_native_fwd=True on Qwen3.5-397B (emb_dim=4096) without overriding base.yml's default wi_tile_fwd_embed_dim: 1024 (or under eval rules with eval_wi_tile_fwd_embed_dim: 1024), wi will silently compute one scale per 1024-column tile instead of one scale over the full K dim?
There was a problem hiding this comment.
Good catch — confirmed (tile_k > size_k also fails on rhs_scale bounds). Our Qwen3.5-397B RL runs here were unaffected (the recipe explicitly sets wi_tile_fwd_embed_dim=4096 and wo_tile_fwd_mlp_dim=1024, with no separate eval_* rules in RL), though runs leaving base.yml's 1024 default would hit it.
Guarded at both layers in 22ebd03 (covered in fp8_moe_native_test.py and verified on tpu7x-8):
- Kernel (
make_gmm_configs): raisesNotImplementedErrorwhen dynamic in-kernel LHS quant useslhs_quant_block_size == 0withtiles.tile_k != dims.size_k. - Config (
_fp8_moe_native_fwd_tile_requirements): checks upfront (when K is unsharded) thatwi/wo(andeval_*when enabled) forward tiles equalemb_dim/moe_mlp_dim.
uwlei
left a comment
uwlei
left a comment
There was a problem hiding this comment.
HBM cap 84992 vs 90112: sj-c3-n-s2q bumped the HBM cap from 84992 to 90112 and peak HBM went from 85.2 to 86.6 GiB. Does fp8_moe_native_fwd + fp8_moe_native_bwd_dlhs fit within the recipe's default 84992 cap (including dropless replay), or will enabling this in the recipe need the 90112 cap too?
There was a problem hiding this comment.
HBM cap 84992 vs 90112: sj-c3-n-s2q bumped the HBM cap from 84992 to 90112 and peak HBM went from 85.2 to 86.6 GiB. Does fp8_moe_native_fwd + fp8_moe_native_bwd_dlhs fit within the recipe's default 84992 cap (including dropless replay), or will enabling this in the recipe need the 90112 cap too?
TL;DR: It fits under the default 84992 cap without OOM (including dropless replay), so enabling it in the recipe does not require 90112. Bumping to 90112 is optional (8.1 GiB free out of 94.7 GiB in sj-c3-n-s2q) and only avoids extra XLA remat on larger-buffer steps (rbf=2.75 or dropless replay).
- Why
sj-c3-n-s2qhad90112: We didn't bump it for e2e (rbf=2.0) — we added90112during our 128-chip speed screen onlewu-a1-pr(rbf=2.75), where larger ragged buffers pushed the module over84992and triggered ~120 ms/micro of extra XLA remat (0remat at90112, both arms run at90112). We kept that flag unchanged when launchingsj-c3-n-s2q. - Normal recipe steps (
rbf=2.0):84992 MiB(83.0 GiB) caps one XLA module (excluding the ~6 GiBpersistent grad accumulators shown in XProf peak). Our earlier 640-chip run at default84992(sj-c2-n-p1,fp8_moe_native_fwdonly) peaked higher (88.4 GiB) with0XLA remat ops, and addingfp8_moe_native_bwd_dlhsdrops peak to86.6 GiB. - Dropless replay: Lowering the XLA cap triggers remat rather than OOM. Our earlier
84992run compiled and launched dropless replay without OOM (before hitting the pre-#6068 index fault); full replay timing was verified at90112(sj-a4edl-n,86.5 GiBpeak), so under84992an overflowed micro-batch may pay ~0.1–0.2 sof extra remat.
Description
This PR adds a native-QAT path for the routed-expert GMMs, on top of the existing fake QAT (
fp8_moe_fake_quant+fp8_fake_quant_hoist=gather8).What fake QAT does today:
What the new flags change:
fp8_moe_native_fwd=true: thewi/woforward GMMs run FP8×FP8 on the TPU7x MXU. They read the gathered codes and per-channel scale directly, and quantize the activation inside the kernel.fp8_moe_native_bwd_dlhs=true: the backward dLHS GMM also reads the codes and dequantizes each tile in VMEM. No BF16 copy of the weight is written to HBM.All new flags default to off. With them off, the MoE GMM programs are byte-identical to base.
Related: the dropless full-buffer index clamp that was part of this PR is now #6068. The end-to-end results below were measured with that clamp in the tree; land #6068 first if you run dropless replay with
moe_routing_index_math=count.Why:
Data flow
flowchart LR subgraph F["Fake QAT"] F1["FP8 weights<br/>+ scales"] --> F2["dequant<br/>to BF16"] F2 --> F3["fwd GMM<br/>BF16 x BF16"] F2 --> F4["bwd dLHS<br/>BF16 x BF16"] F2 --> F5["bwd dW<br/>BF16"] end subgraph N["Native QAT (this PR)"] N1["FP8 weights<br/>+ scales"] --> N3["fwd GMM<br/>FP8 x FP8"] N1 --> N4["bwd dLHS<br/>VMEM dequant"] N1 --> N2["dequant<br/>to BF16"] --> N5["bwd dW<br/>BF16"] endfp8_moe_native_fwd=truefp8_moe_native_bwd_dlhs=truedequantize_weight_codes_ste)[g,1,n]f32 scale, passed asrhs_scalefloat8_e4m3fnin the kernel: one f32 absmax scale per row over the full K, f32 accumulationGradient anchor: an FP8 primal would get an FP8 cotangent. So each weight also enters the op as its BF16 dequantized value, an extra input whose value the forward never reads; XLA drops that unread dequantization. The weight gradient is returned as that input's cotangent. It then flows through
dequantize_weight_codes_steto the FSDP shard exactly as it does today.Implementation
kernels/megablox/ops.pyadds four functions:gmm_silu_fp8: wi with fused SwiGLU.gmm_dswiglu_fp8: wo.gmm_fp8_forward: forward only, forring_ragged_unsort_down_projection.dequantize_fp8_codes.Each
custom_vjpforward callsgmm_v2(rhs=codes, rhs_scale=…, maybe_quantize_lhs=True). Withfp8_moe_native_bwd_dlhs, the dLHS backward callsgmm_v2with a codes RHS (rhs_contract_scale). Otherwise the backward dequantizes and reuses the existing BF16 backward.pallas_mosaic_tpu_v2_gmm_kernel.py:lhs_quant_block_size(512 or full K),lhs_f32_scale(max(amax, 1e-12)/448in f32),lhs_hoist_quantandrhs_scale_first.fp8_moe_native_fwdalways uses full K with an f32 scale and f32 accumulation; the other options are ops-level arguments for benchmarks and tests only.rhs_contract_dequantpath: an FP8 double buffer with per-tile dequant into a VMEM cache, fused into the first matmul of each block by default.tile_n % 32 != 0,tile_k % 128 != 0, non-dividing tiles) fall back to the XLA dequant. Production tiles take the native path.layers/moe.py:Fp8NativeWeight(anchor, codes, scale)pytree, built from the gather8 codes after the prefetch barrier. The shard-mapin_specsgive the scaleP(e, None, n).gmm_up_or_silu(wi);_run_wo_gmm_dswiglu(wo);down_projection/down_projection_bwd_dswigluinside_ring_down_project_and_unpermute(ring-of-experts wo).RoutedMoEcall withfp8_moe_native_fwd=truebut no gather8 absmax now raises instead of silently running fake QAT. This happens for unscanned layers, KV-cache decode, or a Qwix-intercepted call.configs/types.py,configs/base.yml: two new flags, both default off.fp8_moe_native_fwdtruefp8_moe_native_bwd_dlhstruefp8_moe_native_fwdrequiresfp8_moe_fake_quant,fp8_fake_quant_hoist=gather8,sparse_matmul,use_tokamax_gmm,use_gmm_v2,moe_wi_gmm_fused_swiglu,moe_wo_gmm_fused_dswiglu, fixed tiles (use_gmm_v2_heuristic_tiling=false: the heuristic tiling budgets a bf16 accumulator when the LHS is quantized in the kernel, but the native forward accumulates in f32), and single-K-tile forward tiles (wi_tile_fwd_embed_dim= embed dim andwo_tile_fwd_mlp_dim= MoE MLP dim, pluseval_*when separate eval axis rules are set;make_gmm_configsalso rejectstile_k != size_kwhenlhs_quant_block_size=0).fp8_moe_native_bwd_dlhsadditionally requiresfp8_moe_native_fwdandmoe_gmm_v2_dlhs_transpose_rhs.Follow-ups:
fp8_moe_native_fwdimplyfp8_fake_quant_hoist=gather8and the fake-QAT flags; today it requires them, because it reuses gather8's FP8 weights.Tests
Unit and correctness
tests/unit/fp8_fake_quant_hoist_test.py(CPU)scan_layers=false+ native raises the new error.scan_layers=true+ native passes the guard (and stops only at the TPU-only gmm_v2 kernel on CPU). gather8 without native still lowers.tests/unit/fp8_moe_native_test.py(CPU)eval_*);make_gmm_configsrejectslhs_quant_block_size=0withtile_k != size_k(also verified unmocked ontpu7x-8);fp8_moe_native_fwdpasses full-K f32 activation quantization and f32 accumulation; a model withfp8_moe_native_fwdoutside the layer scan raises instead of running fake QAT. 19/19 pass.fp8_moe_native_bwd_dlhsvs dequant backwardfullk_f32+float32config.fullk_f32+ f32 accumulation vs the rollout's MoE kernel: relative error 0.0131–0.0133 (per-512 BF16-scale variant, the old default: 0.0637; fake QAT: 0.0468).f8e4m3fnweights with no BF16 weight upstream; backward kernels are the same 4 as basePerformance: Qwen3.5-397B-A17B, DeepSWE RL, TPU7x. Fake = base (
caafa5bc6). Native = base +fp8_moe_native_fwd=true fp8_moe_native_bwd_dlhs=true, with full-K f32 activation scales and f32 accumulation (set then by two flags that are now the fixed behavior). Both arms use the same image and--xla_tpu_max_hbm_size_mib=90112.Grouped matmul (GMM) kernel microbenchmark (one TPU7x TensorCore, the trainer's per-device shapes and production tiles):
Trainer bench (128 chips, 4x4x8, mean of steps 1–4; tok/s/chip = 37,748,736 tokens per step ÷ (128 × s/step)):
lewu-a1-prrepro (+ HBM cap 90112): uniform routing,ragged_buffer_factor=2.75sj-a4s-f-10251647/sj-a4s-n-10251647pb-r2c9b2-s2's trainer):ragged_buffer_factor=2.0, synthetic production-like routingsj-a4edl-f-10251718/sj-a4edl-n-10251718Row 2 overflows the expert buffer on one micro-batch per step, and the dropless path reruns it (+16 s/step in both arms). That rerun needs the index clamp from #6068, which these runs had.
640-chip e2e (seed 677013886):
sj-c2-f-p1pb-r2c9b2-s2(fake QAT)sj-c3-n-s2qfp8_moe_native_fwd(fullk_f32, f32 acc) +fp8_moe_native_bwd_dlhsXProf (one trainer micro-batch, TPU:0):
Total device busy time isn't compared here: the two runs sampled different micro-batches (attention time alone differs by 247 ms). The step-time gain is measured in the trainer bench above.
Convergence:
at-1015b-s1pb-r2c9b2-s2)sj-c2-f-p1)sj-c3-n-s2q)Checklist
Before submitting this PR, please make sure (put X in square brackets):
gemini-reviewlabel.