MXFP8 cudnn-fe Swiglu fused Grouped MLP fwd/bwd and wgrad group gemm - #5
Draft
wolfcomos wants to merge 13 commits into
Draft
MXFP8 cudnn-fe Swiglu fused Grouped MLP fwd/bwd and wgrad group gemm#5wolfcomos wants to merge 13 commits into
wolfcomos wants to merge 13 commits into
Conversation
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>
wolfcomos
changed the title
wolfcomos
changed the title
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>
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
1 participant
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.
Uh oh!
There was an error while loading. Please reload this page.