[mxfp8 training] Add fused grouped GEMM + SwiGLU + MXFP8 quantization ops for MoE routed experts by wolfcomos · Pull Request #4820 · pytorch/ao

wolfcomos · 2026-08-24T10:54:27Z

Summary

This PR adds four torch custom ops that run a routed-experts (grouped) MLP entirely in MXFP8 with the GEMM↔activation boundaries fused, so the intermediate activations are quantized in-kernel and the BF16 hidden tensor is never materialized:

mxfp8_grouped_gemm_swiglu_fwd — FC1 grouped GEMM on MXFP8 operands + SwiGLU + dual MXFP8 RCEIL requantization (rowwise 1x32 and columnwise 32x1) in one launch; also emits the BF16 pre-activation z consumed unchanged by the backward op.
mxfp8_grouped_gemm — grouped GEMM on prequantized MXFP8 operands → BF16 (serves FC2 forward and FC1 dgrad).
mxfp8_grouped_gemm_dswiglu_bwd — FC2-dgrad grouped GEMM + dSwiGLU + dual MXFP8 requantization.
mxfp8_grouped_gemm_wgrad — ragged-K grouped weight-gradient GEMM (serves both FC1 and FC2 wgrads).

The device kernels are the grouped-GEMM wrapper family shipped by the cudnn python package (>= 1.27, cudnn-frontend); this module contributes the op registration, fakes, validation, and layout plumbing. The package is imported lazily inside the op bodies — torchao gains no hard import dependency; on machines without it, an availability probe reports an actionable reason and is_supported()/the ops raise cleanly. The consuming composite lives in torchtitan (pytorch/torchtitan#4257's override module); a follow-up torchtitan PR wires these ops in as a selectable fusion plan. This torchao PR is ops-only.

Contracts

The kernels have hard alignment contracts, enforced at the op boundary:

Per-expert row counts must be multiples of 256 (ROW_GROUP_ALIGNMENT, the kernels' fixed pad size). This is the one contract that fails silently if violated — sub-256 groups produce nondeterministic corruption, not an error — so the callers' padded token dispatcher is the load-bearing guarantee, and an opt-in per-call check exists (below).
Feature dims (D, F) multiples of 128 (DIM_ALIGNMENT); R * max(dim) < 2**31 (int32 element indexing); 16-byte pointer alignment.
Layouts: activations carry rowwise 1x32 blocked scales; columnwise operands carry per-group blocked scales (logical [cols, rows_g/32] per expert); columnwise qdata is accepted in both majors (no transpose copies on the hot path); offsets is an int32 [G] exclusive-end cumsum. Rows past offsets[-1] of kernel-allocated outputs are garbage and read-forbidden (poisoned-tail immunity is tested).

Validation is two-tier: an always-on metadata tier memoized per (shapes, strides, dtype, device) signature (~3.4 µs/op-call after the first), and an opt-in offsets-VALUES tier (TORCHAO_MXFP8_VALIDATE_OFFSETS=1, costs a D2H sync per call) that checks nondecreasing offsets, the %256 group contract, and the allocation bound — the debug switch for the silent-corruption contract above.

All four ops are bitwise deterministic across runs, including with the kernels' dynamic scheduling enabled (byte-checked in the tests).

Performance

End-to-end in torchtitan (DeepSeek-V3 16B, 4x GB200 at 2062 MHz application clocks, EP=4, TP=1, seq 4096, MXFP8 MoE recipe, 50 steps, interleaved A/B rounds, pooled medians of steps 11-50). Arms: the stock MXFP8 path (pad 128), the same path pad-matched to these kernels' 256-row groups, and the fused composite backed by this PR's ops:

arm                        bs4 tokens/s   bs8 tokens/s
-------------------------  -------------  -------------
stock MXFP8 (pad 128)      14,289         18,706
stock MXFP8 (pad 256)      13,258         17,882
fused (this PR's ops)      13,563 (+2.3%) 18,783 (+5.0%)

The fused gain vs the pad-matched baseline grows with batch (fixed pad tax dilutes; fusion win scales with expert work); at bs8 the fused arm is already net-positive against the pad-128 stock path. Loss trajectories passed paired-seed parity gates (26-step and 50-step budgets with sign-mixing criteria) and a 3-seed x 500-step convergence protocol.

API

from torchao.prototype.moe_training.kernels.mxfp8.cutedsl_grouped_mlp import (
    mxfp8_grouped_gemm_swiglu_fwd,   # (x_q, x_sf, w13_q, w13_sf, offsets)
    mxfp8_grouped_gemm,              # (a_q, a_sf, b_q, b_sf, offsets)
    mxfp8_grouped_gemm_dswiglu_bwd,  # (dy_q, dy_sf, w2_col_q, w2_col_sf, z_bf16, offsets)
    mxfp8_grouped_gemm_wgrad,        # (dy_col_q, dy_col_sf, x_col_q, x_col_sf, offsets)
    is_supported,                    # (model_dim, hidden_dim) -> bool
    DIM_ALIGNMENT, ROW_GROUP_ALIGNMENT, SCALE_BLOCK_SIZE,
)

w13_q is [G, 2F, D] in the kernel's 32-block interleaved gate/up order (the layout a fused w13 projection pack produces). Registered as torch custom ops with fake impls; torch.compile(fullgraph=True) works; zero-row inputs (R == 0) return empty outputs without touching the backend.

Numerics and tests

30 tests, ~36 s on GB200 (skipped cleanly, no collection errors, without SM100 or the cudnn package — CI-safe):

Reference-chain SQNR gates with calibrated bands: a dequantized-operand reference with a measured reduction-order band, and an independent quantize-chain reference the ops must match within 0.1 dB of the pure-quantization band (ops land 93–188 dB at the tested shapes).
Negative controls proven discriminating (whole-matrix instead of per-group columnwise scales < 25 dB, gate/up swap, single scale-byte flip, RCEIL-vs-FLOOR identity) — the gates have teeth.
Both production wgrad stride mixes; NaN-poisoned inactive tails never contaminate active rows; zero-token experts and R == 0; fake/meta parity with eager metadata; compile fullgraph bitwise vs eager; a 9-row malformed-input rejection matrix; the opt-in offsets validator.

Reproducing

pytest test/prototype/moe_training/test_mxfp8_grouped_mlp.py -q

Follow-ups

The torchtitan fusion-plan wiring PR (stacked on Add MXFP8 fused-MLP overrides in torchtitan torchtitan#4257) that consumes these ops end-to-end.
A fixed-capacity (static-shape) dispatcher mode would make the %256 contract satisfiable without per-step host-synced offsets, enabling CUDA-graph capture; today the padded dynamic dispatcher is the validated seam.

🤖 Generated with Claude Code

Adds the shared infrastructure for three fused MXFP8 kernels on the routed expert path (SM100a), plus the first complete kernel. Shared: grouped_mlp_validation.py host preconditions at the custom-op boundary grouped_mlp_epilogue.py device primitives: RCEIL E8M0, NaN-propagating packed amax, tcgen05 blocked-scale indexing, SwiGLU/dSwiGLU policy grouped_mlp_ops.py the three custom ops and their fake/meta impls grouped_gemm_config.py frozen configs, support predicate, the published T2R and epilogue protocols grouped_gemm_core.py ragged blockscaled GEMM: TMA descriptors, tcgen05 mainloop, TMEM, accumulator-to-register handoff epilogue_quant.py rowwise 1x32 and columnwise 32x1 quantization Kernel C (mxfp8_grouped_gemm_wgrad): kernel_wgrad.py BF16 store epilogue and cached launcher cutedsl_grouped_mlp.py launcher facade (A and B land here next) Per-expert row counts are multiples of 128, so no tile straddles an expert boundary. That removes the per-group tensormap updates and the persistent tile scheduler the general grouped-GEMM path needs: an expert is selected by an integer tile-coordinate base. The blockscaled scale layout's K stride was measured to be a constant 512 bytes for every row block, which is what makes that indexing correct. Columnwise blocked scales use whole-matrix to_blocked rather than the per-group K-groups form. The two encodings have identical byte counts and coincide when N <= 128, so the difference is invisible to any length check and to small-shape tests; mixing them corrupts the weight gradient (cosine 0.82). Validation: the quantize chain is bitwise-identical to torchao's RCEIL quantizers over 23.7M qdata bytes and all special values; the GEMM core is bitwise exact on both ragged orientations including zero-token experts, strict inactive tails and the pipeline stage wrap; Kernel C is bitwise exact against a float64 oracle on a production shape (23,068,672 elements) and clean at the A/B -> C seam. Note: cute.testing.assert_ is compiled out unless CUTE_DSL_ENABLE_ASSERTIONS=1, so device-side offset checks are a debugging aid only and the host validation is the sole enforcement of the alignment precondition. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

…tests, bench Replaces the foundation commit's kernel layer with a correctness-first implementation of all three physically fused kernels on public CuTe DSL 4.7.0 API only, and completes the operator surface: mxfp8_grouped_gemm_swiglu_fwd FC1 ragged grouped GEMM + SwiGLU + dual (1x32 rowwise, 32x1 columnwise) MXFP8 RCEIL quantization + BF16 preactivation save, one kernel launch mxfp8_grouped_gemm_dswiglu_bwd FC2 dgrad grouped GEMM + dSwiGLU + dual quantization, one launch mxfp8_grouped_gemm_wgrad generic ragged-K grouped wgrad (FC1 and FC2), BF16 out, one launch per call No inline PTX, no llvm.inline_asm, no private cutlass._mlir/nvvm interfaces, no dsl_user_op asm wrappers remain. The replacements are public API and were verified against the prior implementations: Float8E8M0FNU conversion is natively round-upward (RCEIL), the public f32->E4M3 conversion is byte-identical to torch's cast, fmax(abs=True, nan=True) provides the NaN-propagating amax, and sigmoid composed as 1/(1+exp(-x)) matches torch.sigmoid bitwise. Epilogues stage through shared memory and derive every index from the T2R partitioner's coordinate tensor, removing the previous thread-to-row-ownership and physical-layout assumptions. The config/protocol framework, the published epilogue plug-in docs, and the decorative device-side offset assertions are gone; offset values are documented caller invariants with an opt-in synchronized host validator. Structure now follows moe_training conventions: public functional wrappers with availability detection in torchao/prototype/moe_training/mxfp8_grouped_mlp.py (exported from the package __init__, so a normal import registers the ops), custom ops + output allocation + fakes in kernels/mxfp8/grouped_mlp_ops.py, host validation in kernels/mxfp8/grouped_mlp_validation.py, and all three kernels + launchers in kernels/mxfp8/cutedsl_grouped_mlp.py. Launcher compile caches key on device index, compute capability, dtypes, shapes, and DSL version; streams come from the input tensor's device; G == 0 is rejected and R == 0 short-circuits without a launch. Checked-in tests (test/prototype/moe_training/test_mxfp8_grouped_mlp.py) cover numerics against pure-torch to_mx/to_blocked references (bitwise quantization via saturated-gate exact-product constructions and the shared MXFP8 semantic cases; SQNR gates for the BF16 GEMM stages), inactive-tail and zero-token semantics, FakeTensor/compile contracts, validation negatives, and a warmed torch.profiler launch-count test proving one kernel per op. The benchmark (benchmarks/prototype/moe_training/mxfp8/bench_grouped_mlp.py) compares the decomposed torchao path, TransformerEngine's modular and fused grouped-MLP lanes, and these kernels on identical inputs. The kernels, tests, and benchmark do not import cute_utils and run on the public nvidia-cutlass-dsl 4.7.0 wheel with no local patches. Validation (GB200 SM100, driver 580.173.02, CUDA 13.4 fwd-compat, torch 2.14.0a0, nvidia-cutlass-dsl 4.7.0, pristine cute_utils.py): test suite 46 passed + 1 skipped on GPU (the cross-device negative passes with two visible devices) and 7 passed CPU-only. The rewritten wgrad is bitwise-identical to the previous kernel on zero-token, strict-tail, stage-wrap and 16B FC1-shape configs; kernels A/B are bitwise against the to_mx(RCEIL)+to_blocked references at a production shape (0 mismatches over 2.16M and 4.33M qdata bytes), with kernel A bitwise even on random inputs. torch.profiler confirms exactly one CUDA kernel per warmed call. Capped-clock (1200 MHz) relative timings vs the decomposed torchao lane at R=2048/D=2048/F=1408/G=8: A 194 vs 339 us, B 192 vs 434 us, wgrad ~118 vs ~271 us per call. TE's tuned fused kernels remain substantially faster; tuning is the follow-up phase and the operator contracts are frozen for it. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

- Delete the three _*_entry wrappers; the launchers cute.compile the shared _launch_grouped_gemm directly, with the output tensors passed as a tuple argument (probe-verified: no compile-cache aliasing across the three kernels, outputs bitwise-identical to the entry path, still one launch). - Inline the single/dual-use trace-time builders at their call sites: _t2r_partition, _make_tiled_mma, _make_sf_gemm_tensor, _sigmoid_f32. - Drop dead code: TileCoords.tile_m/tile_n, the unused subtile_idx epilogue parameter, the constant-folded kernel-local l_a, and _KernelConfig's three single-consumer properties (inlined as expressions). - Compress narration comments to contracts; the offset contract and the quantization numeric bullets are kept verbatim. 1723 -> 1606 lines; module-level defs 29 -> 22. Public API (__all__, the three launch_* signatures) and the ops/validation/wrapper modules are byte-identical. Generated cubins and SASS are sha256-identical pre/post at both tested shapes, so no machine code changed. Full suite 46 passed + 1 skipped (GPU) / 7 passed (CPU), green after every fold stage; one-launch property preserved; MXFP8_BENCH_VALIDATE=1 bench clean and within run-to-run noise at the 16B-class shape (1200 MHz-capped host, ratios only). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Four torchao:: custom ops for the routed-expert grouped MLP, each one launch of a cudnn.grouped_gemm_*_wrapper_sm100 CuTe DSL kernel from the standalone cudnn-frontend python package (>= 1.27, SM 10.x; no TransformerEngine dependency): mxfp8_cudnn_grouped_mlp_fwd FC1 GEMM + SwiGLU + dual MXFP8 RCEIL quant + BF16 pre-GLU (32-block GLU row order) mxfp8_cudnn_grouped_mm grouped GEMM on prequantized operands -> BF16 (FC2 forward and FC1 dgrad; b [G,N,K] along K) mxfp8_cudnn_grouped_mlp_bwd FC2 dgrad + dSwiGLU + dual quant of dz mxfp8_cudnn_grouped_mlp_wgrad ragged-reduction weight gradient, dense mode Contract highlights (probe-derived, see agent_scratch/cudnn_fe_torchao): - per-expert row counts and R must be multiples of 256 (cuDNN FE FIX_PAD_SIZE; 128-only groups corrupt results silently and NONDETERMINISTICALLY). Two-tier validation: always-on metadata checks + opt-in TORCHAO_MXFP8_VALIDATE_OFFSETS=1 value checks. - flat blocked E8M0 scale ABI; activation colwise scales are PER-GROUP blocked (K-groups layout) and sized by offsets[-1], which may be < R. - colwise qdata accepted in ANY major (dim1-native transposed memory, kernel un-transposed bytes, and mixes -- every combination probe-proven). - rows past offsets[-1]: caller-allocated outputs untouched; kernel-allocated outputs garbage and read-forbidden (verified with NaN-poisoned tails). Self-gate (GB200, FE 1.27.0/backend 92500, TE image build 401656373): 46/46 gates green -- full fwd+bwd chains at dbg/D!=F/G=1/16B shapes (GEMM outputs 87-159 dB vs dequantized-operand references, quant outputs 31.5 dB), A<R NaN-poison tail case, R==0, zero-token experts (dw slices written as zeros), per-op run-to-run bitwise determinism WITH use_dynamic_sched=True, fake metadata identical to eager for all 11 outputs, torch.compile(fullgraph) bitwise vs eager, and 8 malformed-input rejection cases.

29 tests, ~35 s on one GB200: full-chain numerics at four shape classes (debug with a zero-token expert, D!=F, G=1, 16B-class D=2048/F=1408/G=8) against two references per stage -- dequantized-operand refs with gates derived from the measured FP32 reduction-order band (-12 dB, capped 60 dB; ops measure 93-188 dB) and an independent no-quantization chain from the original BF16 tensors with gates at the measured MXFP8 band -6 dB (z 28.5, y 23.7 dB; op outputs land on the band). Also: the full 2x2 wgrad operand major-mode matrix, A<R strict tail with NaN/Inf-poisoned rows and NaN qdata/0xFF scale bytes in the colwise tail, per-op bitwise determinism, torch.compile fullgraph bitwise vs eager, FakeTensor metadata contracts, R==0, and adversarial negative controls (whole-matrix colwise scales, gate/up 32-block swap, single scale-byte flip, RCEIL-vs-FLOOR scale-byte identity) plus the malformed-input rejection matrix and the opt-in TORCHAO_MXFP8_VALIDATE_OFFSETS path.

…B/C) The cuDNN-frontend grouped-MLP ops (cudnn_grouped_mlp_{ops,validation} + the public wrapper) supersede the in-repo CuTe DSL kernels: same fused FC1+SwiGLU+dual-quant / FC2-dgrad+dSwiGLU / ragged-wgrad surface, with no dependency on this repo's CuTe DSL runtime shims. Delete the kernels (cutedsl_grouped_mlp.py), their ops/validation modules, the public mxfp8_grouped_mlp wrapper, tests, and bench; unwire both __init__ files. The quantizer/swizzle kernels in this package (quant.py, cutedsl/flydsl quantizers, cute_utils) are untouched -- the unfused baseline path and the cudnn composite's casts still use them. torchao.prototype.moe_training import smoke and the 29 cudnn grouped-MLP tests stay green. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

The four custom ops wrapping the cudnn-frontend CuTe DSL grouped-GEMM kernels take over the naming the removed in-repo kernel family vacated: - torchao::mxfp8_cudnn_grouped_mlp_fwd -> torchao::mxfp8_grouped_gemm_swiglu_fwd - torchao::mxfp8_cudnn_grouped_mm -> torchao::mxfp8_grouped_gemm - torchao::mxfp8_cudnn_grouped_mlp_bwd -> torchao::mxfp8_grouped_gemm_dswiglu_bwd - torchao::mxfp8_cudnn_grouped_mlp_wgrad -> torchao::mxfp8_grouped_gemm_wgrad Files follow: kernels/mxfp8/cudnn_grouped_mlp_{ops,validation}.py -> grouped_mlp_{ops,validation}.py; the public wrapper cudnn_grouped_mlp.py -> mxfp8_grouped_mlp.py (availability flag now _mxfp8_grouped_mlp_kernels_available); moe_training/__init__ re-exports the four ops plus is_supported. cudnn-frontend remains named in docstrings as the kernel provenance; identifiers no longer carry it. Also restyle the test module to match the other moe_training MXFP8 tests: module-level capability/availability skips (allow_module_level) instead of a custom availability marker, no per-import noqa, no __main__ block. All 29 tests semantically unchanged and green. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…nto agent/mxfp8-cudnn-grouped-mlp

The always-on validation tier is metadata-only, so its verdict is a pure function of the operands' metadata -- yet the training loop re-ran the full battery on every op call (hundreds per step with identical metadata, including tensors the previous op in the chain had just produced). Each validator now records a (shapes, strides, dtype, device, storage_offset) signature after PASSING and skips straight to the derived dims on repeats: 8.2 -> 3.4 us/call on the 16B fwd signature. First-call and torch.compile capture-time rejection behavior is unchanged (a rejected call never records its signature), the pointer-alignment gate stays covered because storage_offset is part of the signature, and the opt-in TORCHAO_MXFP8_VALIDATE_OFFSETS values check still runs on every call while enabled (values are not metadata). Signature recording caps at 4096 entries; beyond that new signatures simply validate every time. 29/29 ao tests and 13/13 torchtitan composite tests green. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Mirror the gated-act (SwiGLU) + MXFP8 PR's single-module structure: the metadata-validation helpers (grouped_mlp_validation.py) and the public wrapper module (moe_training/mxfp8_grouped_mlp.py) fold into kernels/mxfp8/grouped_mlp_ops.py, which is now self-contained -- availability probe, is_supported, condensed validation helpers, the four custom ops with their fakes, and the availability-gated public wrappers. Both package __init__s return to their upstream state (importing grouped_mlp_ops registers the ops; consumers import it directly). All validation logic and error messages are unchanged; docstrings condensed to contracts. Net: 3 modules / 1,363 lines -> 1 module / 1,196 lines. 29/29 tests green (rejection-message and opt-in offsets gates included); ruff 0.11.6 lint+format clean. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Follow the sibling file convention (cutedsl_gated_act_mxfp8.py, cutedsl_quantize_*.py) for the self-contained grouped-MLP module; the name also matches the slot the removed in-repo kernel family vacated. Content unchanged. 29/29 tests green. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Ops module (cutedsl_grouped_mlp.py): - Skip validation memoization when any dim/stride/storage_offset is a SymInt (unhashable under dynamic-shape compile); validation still runs. - Record signatures only on REAL-tensor passes so the first real call always runs the 16-byte data_ptr gate that fake passes skip. - Widen the fwd/bwd int32 element-index guards to R * max(D, 2F), mirroring the mm/wgrad validators' R * max(N, K). - Apply the 16-byte data_ptr gate (skipped for fakes) to validate_blocked_scales and validate_ragged_colwise_scales via a shared _check_pointer_alignment helper. - Make _cached_ones cross-stream safe: record a CUDA event after the first fill and have every cache hit wait on it from the consuming stream (the buffer is immutable once filled, so one event suffices). - Parametrize validate_feature_dims' dim names so mm/wgrad call sites report their actual N/K dims instead of D/F. - Tighten the availability gate to exactly capability (10, 0) (_is_sm_10x -> _is_sm100): the cudnn wrappers are *_sm100-specific. Tests (test_mxfp8_grouped_mlp.py): - Guard the cutedsl_grouped_mlp import with a module-level skip so a stale installed torchao skips instead of erroring at collection. - Extend test_r0_all_ops to the R==0 early returns of all four ops. - Restructure test_optin_offsets_validation to assert the default-build non-rejection on the validator directly, never launching a kernel with out-of-contract 128-row offsets; add a nondecreasing-offsets rejection. - Fix the _quant_weight_colwise docstring (it builds contiguous row-major bytes, not dim1-native strides), add a native=True path, and cover the production dim1-native weight major for ops 2 and 3. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

pytorch-bot · 2026-08-24T10:54:31Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4820

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure, 1 Unclassified Failure

As of commit 5a50196 with merge base be38123 ():

NEW FAILURE - The following job has failed:

Run PT2E GPU Tests / test-rocm (ROCM Nightly, linux.rocm.gpu.ecosystem.mi350.1, --pre torch torchvision --index-url ht... / linux-job (gh)
Could not load credentials from any providers

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

PR Label Check / Check PR Labels (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
##[error]This PR requires at least one label starting with 'module:'. Available modules can be found at: https://github.com/pytorch/ao/labels?q=module

This comment was automatically generated by Dr. CI and updates every 15 minutes.

vkuzo · 2026-08-24T11:37:35Z

this is interesting! Will take a further look, but I have two initial comments:

Can we change torchao/prototype/moe_training/kernels/mxfp8/cutedsl_grouped_mlp.py to something like cudnn_grouped_mlp.py, and make the op names state clearly that this is calling cudnn?
Also, for the tests, can we make sure that each op has a plain pytorch reference, and we are validating that the numerics match the reference?

wolfcomos and others added 13 commits August 16, 2026 21:12

Merge branch 'pytorch:main' into agent/mxfp8-cudnn-grouped-mlp

ddd8846

Merge remote-tracking branch 'origin/agent/mxfp8-cudnn-grouped-mlp' i…

4d4079d

…nto agent/mxfp8-cudnn-grouped-mlp

wolfcomos requested review from andrewor14, jerryzh168 and vkuzo as code owners August 24, 2026 10:54

meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 24, 2026

wolfcomos mentioned this pull request Aug 24, 2026

Wire the fully-fused grouped-GEMM MXFP8 plan behind fusion_plan wolfcomos/torchtitan#12

Draft

alanhuangyoo mentioned this pull request Aug 24, 2026

Gate MX/NVFP4 safetensors test configs on sm100 #4821

Open

Navigation Menu

Uh oh!

[mxfp8 training] Add fused grouped GEMM + SwiGLU + MXFP8 quantization ops for MoE routed experts - #4820

[mxfp8 training] Add fused grouped GEMM + SwiGLU + MXFP8 quantization ops for MoE routed experts#4820
wolfcomos wants to merge 13 commits into
pytorch:mainfrom
wolfcomos:agent/mxfp8-cudnn-grouped-mlp

wolfcomos commented Aug 24, 2026

Uh oh!

pytorch-bot Bot commented Aug 24, 2026 •

edited

Loading

Uh oh!

vkuzo commented Aug 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

Conversation

wolfcomos commented Aug 24, 2026

Summary

Contracts

Performance

API

Numerics and tests

Reproducing

Follow-ups

Uh oh!

pytorch-bot Bot commented Aug 24, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/4820

❌ 1 New Failure, 1 Unclassified Failure

Uh oh!

vkuzo commented Aug 24, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

pytorch-bot Bot commented Aug 24, 2026 •

edited

Loading