Repository navigation
Conversation
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
for more information, see https://pre-commit.ci
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Conflicts arose because the base branch was squash-merged as NVIDIA#3642. Resolved by 3-way merging against the pre-gated tip (bfdb240) instead of the stale fork point. CuTeDSL/__init__.py: keep the entrypoint loop and register get_mxfp8_requantization_function from NVIDIA#3560 in it. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
for more information, see https://pre-commit.ci
maybe_warn_cutedsl_not_chosen was renamed to report_cutedsl_not_chosen in NVIDIA#3746, so this file no longer compiled with NVTE_WITH_CUTEDSL. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
The colwise-only gated kernel ran 64 threads per CTA (one per column), half the warps of gated_mxfp8.cuh's colwise-only config at the same SMEM-limited CTAs per SM, and staged the gated result through SMEM before quantizing it. Use 128 threads with two per column, fuse the gated op into the quantization in registers, and combine the partial block amaxes through SMEM, as the CUDA kernel does. Bit-exact vs CUDA; colwise gated goes from ~0.93x to 1.13x of CUDA on GB200 (swiglu bwd bf16 0.81x -> 1.08x). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Rowwise-only gated quantization computed the gated op in a separate pass that read every input element by element down columns, wrote the result back to SMEM and synced, then read it again to quantize: 4.5-7x the CUDA kernel's SMEM loads and 5x its stores at the same occupancy. Compute the op in registers inside the rowwise quantization, with the general kernel's rowwise thread mapping and bank rotation, as gated_mxfp8.cuh does. Also size the colwise-only amax exchange buffer by the number of outputs, so FP32 forward fits the CUDA kernel's 6 CTAs per SM instead of 5. Bit-exact vs CUDA; on GB200 rowwise gated goes from 1.05x to 1.12x of CUDA (configs slower than 3%: 246 -> 2). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
Fixes pylint E0606 (possibly-used-before-assignment) in the FWD path. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Signed-off-by: Kaining Zhong <kainingz@nvidia.com>
…gister-resident MXFP8 kernel Add opt-in BF16 lookup tables (NVTE_ACTIVATION_TABLE=1) for GeLU, SiLU and their derivatives in the MXFP8 quantize. TE common builds each table on first use, once per device. The CuTeDSL register-resident kernel reads the tables from shared memory; the CUDA generic kernel and the other CuTeDSL kernels compute the values the tables hold, so both backends write the same bytes. The CuTeDSL register-resident kernel now also takes activation derivatives, on 256-column tiles, with fused dbias and amax, any row count, column tails, and FP16 and FP32 inputs. The CuTeDSL kernels add dbias in the CUDA kernels' row order, except FP32 dbias with rowwise-only output, where CUDA adds row pairs (r, r + 32) first. They give a block of only NaNs the CUDA kernels' scale. The shared-memory tile quantizer's ACTIVATION_TABLE parameter defaults to "none", so callers that read no table need not pass it. test_activation_table checks every BF16 input of each table against the computed activation. test_cast_mxfp8_act_exact checks the fused quantize bit for bit against the activation computed per element and then quantized, with the tables on and off. Signed-off-by: Teddy Do <tdophung@nvidia.com>
…ations Add MXFP8GatedQuantizeG2RKernel, a register-resident CuTeDSL kernel for every gated activation. A CTA covers 256 output columns and walks 32-row bands. Each thread loads its rows of the act and gate halves (and of the incoming gradient, in backward) into registers, computes the gated op there with gated_elementwise_op, rounds it to BF16 as the CUDA kernels do, and quantizes the band in both directions from registers. The backward quantizes its two output halves one after the other. The register-tile quantizer of the bidimensional kernel is factored out (quantize_mxfp8_rmem_tile_bidimensional) and shared by both kernels; the bidimensional kernel writes the same bytes as before. MXFP8GatedQuantizeEntry picks it for BF16 input with columnwise output, and the shared-memory gated kernel otherwise: FP16 or FP32 input, rowwise-only output, columnwise scales not 8-byte aligned, or more than 65535 row bands. Backward with GEMM-swizzled scales stays on the CUDA kernel. Co-authored-by: Kaining Zhong <kainingz@nvidia.com> Signed-off-by: Teddy Do <tdophung@nvidia.com>
tdophung
force-pushed
the
teddy/mxfp8-cutedsl-gated-g2r
branch
from
November 13, 2026 02:28
306e10c to
fe01795
Compare
8 of 13 tasks
…nch one band per gated CTA In the register-tile quantizer that the bidimensional and gated register kernels share, the threads that reduce the columnwise amaxes now also write the columnwise scales of their own columns, and the warps that hold no reducer skip that code: a loop of one trip or none, since ptxas predicates the body of an `if` in every warp. Before, those warps issued the whole reduction with all lanes off, and the first warp row wrote the scales after reading them back from shared memory. This removes those instructions and most of the gated backward's register spills. The gated register kernel now launches one CTA per 32-row band instead of walking several bands per CTA. Unlike the bidimensional kernel, it has no per-CTA setup (a table) to amortize over a walk. At 4096 tokens this doubles the grid, which is faster, mostly in the backward. Also correct the input-load comment: the kernel loads all its input rows in one step, so no fewer rows are live. Output bytes are unchanged. Signed-off-by: Teddy Do <tdophung@nvidia.com>
…ed SwiGLU register kernel SwiGLU's sigmoid, 1 / (1 + exp(-x)), divides with an IEEE reciprocal that ptxas guards with a per-element range check and a slow path. In the gated register kernel that is about a third of the forward's instructions. For 1 <= y < 2^126, rcp.rn.f32 is rcp.approx.ftz refined by one Newton step, so add rcp_f32_in_range, which issues that sequence without the check, and sigmoid(..., fast_rcp=True), which uses it. For x >= SIGMOID_FAST_RCP_MIN (-80), 1 + exp(-x) is in that range, with either exp. The gated register kernel checks each 8-value thread row once: when the NaN-propagating minimum of its act values is at least the bound, the row's SwiGLU, forward and backward, takes the fast reciprocal; otherwise (a value below the bound, -inf or NaN) it divides as before. Each path is a loop of one trip or none, so a warp whose rows all take one path skips the other. gated_elementwise_op takes the choice as FAST_RCP, False by default, so the shared-memory gated kernel, the other activations and the non-gated kernels are unchanged. The fast reciprocal equals the division bit for bit on every FP32 y in [1, 2^126), so output bytes are unchanged. A new test runs SwiGLU over every BF16 act value against the CUDA backend, on a shape with row and column tails too. Signed-off-by: Teddy Do <tdophung@nvidia.com>
… the CuTeDSL MXFP8 kernels MXFP8GatedQuantizeConfig no longer rejects the backward with GEMM-swizzled scales, and the register kernel takes it too: it splits the scales of its two output halves with _split_gated_scales, as the shared-memory kernel does. The split matches the CUDA kernel's layout when each half spans whole 128-column scale tiles (N % 128 == 0). The C++ hook keeps handing N % 128 != 0 to the CUDA kernel, now without reporting it, since that fallback is by design. Before, Python rejected every swizzled backward while C++ declined only N % 128 != 0, so with NVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN=ERROR a swizzled backward raised for every N. BF16 input with columnwise output takes the register kernel; FP16, FP32 and rowwise-only output take the shared-memory kernel. The backend tests now run the swizzled backward wherever N % 128 == 0, and a new test runs it for every input dtype and output direction on a shape with a row tail and a 128-column tail. That covers the shared-memory kernel's rowwise-only and colwise-only paths and the register kernel's row tail with swizzled scales. Signed-off-by: Teddy Do <tdophung@nvidia.com>
This branch has not been deployed
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
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.
Description
Stacked on #3840 (18b020d) and #3761. This branch is #3840's head (kainzhong's CuTeDSL gated MXFP8 kernels),
then #3761's first commit as it applies on #3840, then this PR's four commits. Review only this PR's commits
until #3840 and #3761 merge; it will be rebased as they do. #3761 itself sits on main. Its later commits overlap
this PR in two places: the shared quantizer's column-scale write (the same mechanism, written differently around
the gated path; whichever PR merges second resolves it) and the sigmoid helper (the same code, apart from #3840's
fast_mathargument).This PR builds on Kaining Zhong's CuTeDSL port of the gated MXFP8 kernel in #3840. It reuses #3840's
gated_elementwise_op,_split_gated_scalesandMXFP8GatedQuantizeConfig.The new kernel:
MXFP8GatedQuantizeG2RKernelis a register-resident CuTeDSL kernel for every gated activation:SwiGLU, GeGLU, ReGLU, QGeGLU, SReGLU, GLU, clamped SwiGLU and SiTU-GLU, forward and backward.
gated_elementwise_op, rounds it to BF16 as the CUDA kernel does,and quantizes the band in both directions from registers.
gated_mxfp8.cuhwrites.#3840's shared-memory kernel instead stages each tile through shared memory when it writes both directions. For
rowwise-only and columnwise-only output it fuses the gated op into the quantization in registers.
Shared quantizer: the register-tile quantizer of the bidimensional kernel (
quantize_mxfp8_rmem_tile_bidimensionaland its layout and bounds helpers) is factored out and shared by both kernels. The gated kernel loads its rows with
its own
_load_rows. The bidimensional kernel writes the same bytes as before.Kernel selection:
MXFP8GatedQuantizeEntrypicks the kernel per config and per call.These changes take effect only with the CuTeDSL backend on.
Commits after #3761's (review them in order; each is measured against its parent below):
amax also write its scale (one byte per store) in a warp-uniform loop, in the shared quantizer that the
non-gated kernel uses too; [Common] Add BF16 activation tables and fused dbias to the CuTeDSL register-resident MXFP8 kernel #3761 makes the same change there. The gated kernel launches one CTA per 32-row band.
Output bytes are unchanged. This removes the backward slowdown noted below.
rcp.approx.ftzplus one Newton step for1 / (1 + exp(-x)), which equals the rounded division there; any otherrow takes the precise division. Used by the register-resident and [Common] Add gated MXFP8 CuTeDSL kernel #3840's shared-memory gated kernels. The helper
in
activations.py/utils.pyis the one [Common] Add BF16 activation tables and fused dbias to the CuTeDSL register-resident MXFP8 kernel #3761 adds.FP32 and rowwise-only output run [Common] Add gated MXFP8 CuTeDSL kernel #3840's shared-memory kernel. The C++ dispatch declines
cols % 128 != 0silently, as Python does. A new test covers every dtype and direction at a row tail.
Type of change
Changes
CuTeDSL/cast/mxfp8/quantize_mxfp8_gated_g2r_bidim.py(new): the register-resident gated kernel.CuTeDSL/cast/mxfp8/quantize_mxfp8_g2r_bidim.py: the bidimensional register-tile quantizer, factored out(
quantize_mxfp8_rmem_tile_bidimensional), with aGATEDflag and NaN-to-0 helpers for the gated outputs.CuTeDSL/cast/mxfp8/quantize_mxfp8_common.py:MXFP8GatedQuantizeEntry, whichMXFP8GatedQuantizeConfig.compile()now builds in place of [Common] Add gated MXFP8 CuTeDSL kernel #3840's shared-memory kernel.tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py,tests/jax/test_mxfp8_cutedsl_backend.py) compare this kernel's output with the CUDA kernel's, bit for bit, onthe BF16 cases it takes.
Performance
One GB300 GPU, CuTeDSL backend on. This PR against its base (#3840's head 18b020d plus #3761's commit), on the same
GPU:
swiglu)MXFP8GatedQuantizeG2SKernel->MXFP8GatedQuantizeG2RKernelswiglu)dswiglu), compact scalesswiglu)swiglu)unchanged, with the same bytes (same GPU, base -> this PR, tables on, BF16 4096x18432): GeLU 62.01 -> 62.35 us
(0.99x; 1.00x by CUPTI), dbias + dGeLU 125.95 -> 125.95 us (1.00x).
Per commit, each against its parent, same GPU, event median (kernel time alone, CUPTI, agrees). BF16 in unless
noted; output bytes identical between parent and commit in every row.
Gated guard lanes, this PR's head against commit 1 (GB300, full mode): every lane is faster, 1.11x to 1.40x
(compact dSwiGLU 1.38x; FP16, FP32 and rowwise-only swizzled dSwiGLU 1.35x, 1.21x, 1.15x; SwiGLU tail 1.32x; GeGLU
1.11x; GeGLU tail 1.12x; dSwiGLU tail 1.40x). GB300 rows of commits 3 and 4 were measured on commit 2's earlier
form, which differed only in the compact non-gated scale store; Rubin rows are on the final commits.
Known issue, resolved by commit 2: moving onto #3840's
gated_elementwise_ophad made the compact backward 6.6%slower than this kernel's earlier code (148.3 against 139.1 us on #3840's earlier head). Commit 2 takes it from
149.50 to 117.76 us.
Measurement details:
nvcr.io/nvidia/pytorch:26.10-py3, build 477939455, index digestsha256:6d347d649d5dd108b4652f1b9cea6a3c859692ac408617a9108616da40e3da1f.NVTE_FRAMEWORK=pytorch NVTE_WITH_CUTEDSL=1 NVTE_CUDA_ARCHS=100(nvcc 13.4).NVTE_ENABLE_CUTEDSL_BACKEND=1 NVTE_ACTIVATION_TABLE=1in both arms; the tables donot apply to gated ops. The kernels each arm launched were checked with torch.profiler.
tex.swigluortex.dswiglucall with anMXFP8Quantizer(MXFP8 block scaling): BF16 input fromrandn, E4M3 output, GEMM-swizzled scales unless stated.launch is enqueued.
round.
NVTE_ENABLE_CUTEDSL_BACKEND=1 NVTE_ACTIVATION_TABLE=1 python bench_quantize.py --op swiglu --m 4096 --n 18432 --in-dtype bf16 --fp8 e4m3 --direction both --swizzle 1 --reps 500 --warmup 20 --flush read(
--m 16384;--op dswiglu --swizzle 0;--direction colwiseor--direction rowwisefor the last two; N is theoutput width).
The script is not in TE's
benchmarks/.Testing
All runs were on one GB300 (SM 10.3), and the C++ tests were built for sm_103a.
NVTE_WITH_CUTEDSL=0.NVTE_ENABLE_CUTEDSL_BACKEND=1andNVTE_REPORT_IF_CUTEDSL_BACKEND_NOT_CHOSEN=ERROR.tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py -k gated(bitwise against the CUDA backend)::test_gated_swizzled_bwd_tails(new; swizzled backward, 3 dtypes x 3 directions x 8 activations at 264x384)test_operator --gtest_filter='*CastMXFP8_GatedAct*'test_operator --gtest_filter='*CastMXFP8ActExact*'(the shared quantizer leaves the non-gated kernel's bytes unchanged)tests/pytorch/mxfp8/test_mxfp8_cutedsl_backend.py, non-gated cast and dtypes-k 'test_gated_sizes and SwiGLU and 32x32 and 256x1024'CUTEDSL_DEBUG_LOGGING.MXFP8GatedQuantizeG2RKerneland matched CUDA bit for bit: 336 cases,BF16 with columnwise output, forward and compact backward.
are
N % 128 != 0swizzled backward, which falls back to CUDA by design.and Rubin.
tests/jax/test_mxfp8_cutedsl_backend.py, the minimum-version builds, andSM 10.0.
Checklist: