Repository navigation
Conversation
…tructure Bring the NVFP4 tuned-1D CuTeDSL backend from NVIDIA#3223 onto main, whose CuTeDSL plumbing has changed since that PR branched: - Build and dispatch the backend only with NVTE_WITH_CUTEDSL, and enable it with NVTE_ENABLE_CUTEDSL_BACKEND like MXFP8. - Give NVFP4QuantConfig a to_id() and SM architecture so it uses TVMFFIConfigCache, and pass the device index to DLTensorWrapper, whose FP4 packing main already provides. - Pass the transposed output as (cols, rows): for a rank > 2 input TE shapes it {cols, d0, d1, ...}, which the default 2D flattening of DLTensorWrapper collapses to the wrong extents. - Fall back to CUDA, with a reason under NVTE_WARN_IF_CUTEDSL_BACKEND_NOT_CHOSEN, for configurations the compiled ABI cannot take: non-BF16 input, GEMM-swizzled or non-E4M3 scales, missing amax tensors, and flat dims not divisible by 32. - Keep the NVFP4-only helpers in CuTeDSL/cast/nvfp4/utils.py and philox_rng.py instead of the shared utils.py, using main's inline_ptx style, and register the native TVM-FFI function as MXFP8 does. - Start the packed abs-max from a zero held in a register: as a literal it became an immediate operand, which ptxas rejects for max.xorsign.abs.bf16x2. - Switch the test to main's backend switches and config keys. Signed-off-by: Jan Bielak <jbielak@nvidia.com> Signed-off-by: Teddy Do <tdophung@nvidia.com>
Lay out the NVFP4 backend the way the MXFP8 one is organized: quantize_nvfp4_common.py holds the config, a kernel base class with supports() / eligible_to_dispatch(), the NVFP4QuantizeEntry that picks a kernel at run time, and compilation and TVM-FFI registration. The TMA-staged tuned-1D kernel moves to quantize_nvfp4_g2s_general.py as NVFP4QuantizeGeneralG2SKernel, the fallback for every config, so that specialized kernels can be added beside it without touching it. No functional change. Signed-off-by: Teddy Do <tdophung@nvidia.com>
Add two specialized kernels to the CuTeDSL NVFP4 backend, the CuTeDSL counterparts of the CUDA register-resident 1D quantize kernel, and dispatch to them ahead of the general TMA-staged kernel for BF16 input without fast math or row-scaled NVFP4: - NVFP4QuantizeRowwiseG2RKernel (rowwise only) streams each 16-element block from one 32-byte load straight to its output, with the next wave of rows prefetched into L2. - NVFP4QuantizeBidimensionalG2RKernel (rowwise + columnwise) reads every row twice, once per direction, with the second read served by L2, and stages only the transposed output in shared memory. A cost model picks the number of 64-row slabs per CTA per launch. The per-block coefficient depends only on the block's scale byte, so its divide is computed once per CTA into a shared-memory table. Round-to- nearest output is byte-identical to the CUDA kernels. Stochastic rounding draws Philox once per output block, keyed by the block's coordinates and direction (philox_rng.py). Each draw is plain Philox4x32 of a known counter, as curanddx computes it, with the round count from NVTE_BUILD_NUM_PHILOX_ROUNDS; builds with fewer than 3 rounds keep the general kernel for stochastic rounding. The random bits for a given seed differ from the general kernel's, which follow its thread layout. The kernels launch with the driver's default shared-memory carveout: with min_blocks_per_mp set, CuTeDSL otherwise derives one that ignores the driver's per-CTA reservation and limits residency below the target. Test that the folded block draw equals plain Philox4x32, and that an input that is only 16-byte aligned is served by the general kernel. Signed-off-by: Teddy Do <tdophung@nvidia.com>
tdophung
changed the title
7 of 13 tasks
Follow test_mxfp8_cutedsl_backend.py: quantize with the CUDA and the CuTeDSL backend and compare bytes, guarding against a silent CUDA fallback by checking that the config's kernel is registered. Drop the call-counting proxies and the cases for configs the CuTeDSL backend does not serve, which only compared the CUDA kernel against itself. Add the g2r grid.y overflow cases that the g2s kernel must serve, and shorten the module docstring. 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
Adds register-resident (g2r) NVFP4 1D quantize kernels to the CuTeDSL NVFP4 backend: the CuTeDSL counterparts of the CUDA register kernel in the companion CUDA PR (
teddy/nvfp4-register-quantize-1d). The backend is organized like the MXFP8 one after #3642: a common module with the config, a kernel base class (supports()/eligible_to_dispatch()), theNVFP4QuantizeEntrydispatcher and compilation, plus one module per kernel.Depends on #3223. The first commit brings that PR's NVFP4 backend onto current main, whose CuTeDSL infrastructure has changed since it branched; it stays authored by its original author. This PR will be rebased onto #3223 once that lands, so the first commit is expected to shrink or disappear. Opened as a draft for visibility until then.
Type of change
Changes
NVTE_WITH_CUTEDSL,NVTE_ENABLE_CUTEDSL_BACKEND,TVMFFIConfigCache), with CUDA fallback (and a reason underNVTE_WARN_IF_CUTEDSL_BACKEND_NOT_CHOSEN) for configs the compiled ABI cannot take. Also fixes the transposed-output shape for rank > 2 inputs and a ptxas rejection of an immediate operand tomax.xorsign.abs.bf16x2.quantize_nvfp4_common.pyand per-kernel modules; the TMA-staged tuned-1D kernel becomesNVFP4QuantizeGeneralG2SKernel, the fallback for every config.NVFP4QuantizeRowwiseG2RKernel(rowwise only) andNVFP4QuantizeBidimensionalG2RKernel(rowwise + columnwise), dispatched for BF16 input without fast math or row-scaled NVFP4 when their alignment and grid checks pass; exactly one kernel launches per call. Round-to-nearest output is byte-identical to the CUDA kernels.philox_rng.py); each draw is plain Philox4x32 (as curanddx computes it) withNVTE_BUILD_NUM_PHILOX_ROUNDSrounds. Its random bits for a given seed differ from the general kernel's.test_nvfp4_cutedsl_backend.py, laid out like the MXFP8 one): CUDA vs CuTeDSL byte equality over shapes, directions, fast math, row scaling, zero blocks and the noop flag; inputs the g2r kernels must hand to g2s (16-byte-aligned pointer, grid.y overflow); stochastic rounding checked statistically and for determinism; the per-block draw against plain Philox4x32.Results (GB200, ncu kernel time, BF16 randn, E4M3 scales, cold L2)
Speedup over TE main (
quantize_transpose_tuned_1D), shapes 1024x1024 to 32768x16384:Per-shape numbers and the diff: https://claude.ai/artifact/3eqb1p2shCLjyyD68SuQ3f?sk=R5I3hkJLPOMfmm88DJYRUA
Testing
tests/pytorch/nvfp4/test_nvfp4_cutedsl_backend.pywithNVTE_ENABLE_CUTEDSL_BACKEND=1: 53 passed.test_operatorwith the backend on,--gtest_filter='OperatorTest/FusedCastTransposeNVFP4*:OperatorTestRowScaled/FusedCastTransposeNVFP4*:OperatorTestUE5M3/FusedCastTransposeNVFP4*': 112 passed.Checklist: