fp8_autocast bug fix when switching from non-fp8 execution - #2
Merged
Merged
Conversation
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
13 tasks
8 of 13 tasks
Merged
2 of 13 tasks
4 of 13 tasks
timmoon10
added a commit
that referenced
this pull request
Oct 4, 2026
* Prototype NVFP4 with UE5M3 scales Co-authored-by: Teddy Do <tdophung@nvidia.com> Co-authored-by: Varun Thumbe <vthumbe@nvidia.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Enable e5m3 fused GEMM kernels from cuDNN (#2) * [PyTorch] Enable e5m3 fused GEMM kernels from cuDNN Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * have to pad to 256 to use cuDNN Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * fix: need to pass scale_dtype Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * route wgrad to cuDNN's wgrad API Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * Support grouped linear with NVFP4-UE5M3 NVFP4-UE5M3 grouped GEMM falls back to dense GEMMs. Generalize usage of wgrad kernel and use when tensors sizes are not 256-aligned. Fix inconsistent m,n,k GEMM notation. Remove ue5m3 hacks in op fuser tests. Add ue5m3 to grouped MLP tests. Signed-off-by: Tim Moon <tmoon@nvidia.com> * Fix typos Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> --------- Signed-off-by: Kaining Zhong <kainingz@nvidia.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Tim Moon <tmoon@nvidia.com> Co-authored-by: Codex <noreply@openai.com> * Use custom recipe for NVFP4-UE5M3 tests Signed-off-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add grouped MLP kernel for GGEMM+SwiGLU+RHT+quant Signed-off-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Debug integration with GGEMM+GLU+RHT+quant Signed-off-by: Tim Moon <tmoon@nvidia.com> * Remove scale max helper functions from NVFP4 cast utils Signed-off-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix compile error Signed-off-by: Tim Moon <tmoon@nvidia.com> * Fix NVFP4 scale dtype ABI Signed-off-by: Tim Moon <tmoon@nvidia.com> * remove redundant output alloc Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * no need to pad N now Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * fix linting errors Signed-off-by: Kaining Zhong <kainingz@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Disable cuDNN GGEMM+GLU+RHT+quant kernel Signed-off-by: Tim Moon <tmoon@nvidia.com> * Guard NVFP4 alpha scaling by scaling mode Signed-off-by: Tim Moon <tmoon@nvidia.com> * Restore UE5M3 NVFP4 cast support Signed-off-by: Tim Moon <tmoon@nvidia.com> * Localize NVFP4 4over6 scale policy Signed-off-by: Tim Moon <tmoon@nvidia.com> * Tweak arg order in C API functions Signed-off-by: Tim Moon <tmoon@nvidia.com> * Rename cuDNN GGEMM helper functions for general_gemm Signed-off-by: Tim Moon <tmoon@nvidia.com> * Fix compilation error in C++ test Signed-off-by: Tim Moon <tmoon@nvidia.com> * Treat nvfp4_e4m3_max=0 as unset value Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> * Debug torch.compile test failure Signed-off-by: Tim Moon <tmoon@nvidia.com> * Debug test failures Signed-off-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Revert accidental CUTLASS commit change Signed-off-by: Tim Moon <tmoon@nvidia.com> * Enable cuDNN GGEMM+GLU+RHT+quant kernel Signed-off-by: Tim Moon <tmoon@nvidia.com> * Remove incorrect scale_inv_dtype arg to NVFP4Tensor constructor Signed-off-by: Tim Moon <tmoon@nvidia.com> * Fix bug when selecting cuDNN GGEMM+GLU+RHT+quant kernel Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Restore GGEMM+GLU+RHT+amax kernel with RHT sign mask Signed-off-by: Tim Moon <tmoon@nvidia.com> * Debug minor test failures Signed-off-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix incorrect scale dtypes in grouped tensor builder method Signed-off-by: Tim Moon <tmoon@nvidia.com> * Avoid redundant amax ptr check in row-scaled NVFP4 quantize Review suggestion from @ptrendx Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Clean up C++ unit tests Signed-off-by: Tim Moon <tmoon@nvidia.com> * Address some review comments from @ptrendx Signed-off-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Respect no-tensor-scaling in grouped tensor helper function Signed-off-by: Tim Moon <tmoon@nvidia.com> * Debug C++ test compilation error Signed-off-by: Tim Moon <tmoon@nvidia.com> * Make sure GEMM alpha/beta scales are on GPU, even without amaxes Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> * Create new Mcore DDP integration functions rather than breaking backward compatibility Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> * Remove unrelated tests from NVFP4 dequantize C++ unit tests Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> * Add reference impl and test for NVFP4-UE5M3 Co-authored-by: Codex <noreply@openai.com> Signed-off-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make NVFP4-UE5M3 grouped GEMM CUDA-graph safe (#4) * Make NVFP4-UE5M3 grouped GEMM CUDA-graph safe The UE5M3 grouped GEMM builds its one-element int32 offsets tensor inline at two call sites: torch.tensor([tokens], dtype=torch.int32, device=a_tensor.device) # wgrad torch.tensor([N], dtype=torch.int32, device=device) # fprop Each of these materializes an unpinned CPU tensor and copies it host-to-device. CUDA rejects that during graph capture: RuntimeError: Cannot copy between CPU and CUDA tensors during CUDA graph capture unless the CPU tensor is pinned. Please use tensor.pin_memory() or allocate the tensor with pin_memory=True. Any model capturing these GEMMs in a CUDA graph therefore fails at capture time. This was hit end-to-end on DeepSeek-V3 671B, where every convergence config captures full forward/backward graphs; it aborts during warmup capture and no training step completes. Route both sites through a cached helper, mirroring get_cached_ones_tensor in this same file -- which both functions already use for their `ones` tensor a few lines earlier, and whose docstring notes it keeps "stable data pointers across CUDA graph replays". Caching by (value, device) moves the host-to-device copy to warmup and keeps the pointer stable across replays. The offsets are constant for a given shape, so this is numerically neutral. Signed-off-by: Michal Futrega <mfutrega@nvidia.com> * Review suggestions Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --------- Signed-off-by: Michal Futrega <mfutrega@nvidia.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Fix JAX NVFP4 GEMM device scalar handling Signed-off-by: tdophung <tdophung@nvidia.com> * Replace NVFP4 amax heuristic with explicit alpha/beta residency flag b3415c1 gated the NVFP4 per-tensor scale branch in cublas_gemm() on have_nvfp4_amax, an OR over the two operands' amax pointers. That branch is also what moves alpha/beta from host into device workspace, and CUBLASLT_POINTER_MODE_DEVICE is set unconditionally for every NVFP4 GEMM. So when both operands had a null amax the branch was skipped, leaving alpha/beta as host stack pointers that cuBLAS then dereferenced as device addresses -- garbage output, not a mis-scaled result. That is exactly the shape of the CI failure: only the both_unit_global_scale parametrization of test_gemm_with_missing_nvfp4_amax went red, while the x_/w_ cases passed because the OR still held. Note that nvte_nvfp4_compute_per_tensor_scale with a null amax is an exact identity on alpha (it substitutes scale_max * fp4_max), so skipping it never saved meaningful work -- its real job on that path is materializing the scalars on device. Making the call is always safe; skipping it is not. Replace the heuristic with kNVTEMatmulConfigAlphaBetaOnDevice, which states pointer residency directly instead of inferring it. Default false, so PyTorch (nvte_cublas_gemm/_scaled pass host pointers, and generic_gemm reaches nvte_cublas_gemm_v2 without setting the flag) is restored to its pre-b3415c19 behavior. JAX sets it when is_nvfp4_scaling(), matching the device buffers it binds for alpha/beta, and folds the per-tensor scale into alpha itself. Also document in the JAX helper that its block-scale bound is hardcoded to E4M3's 448 while the C++ get_nvfp4_scale_max() returns 114688 for UE5M3, which will need reconciling once UE5M3 reaches the JAX path. Verified on GB200: tests/pytorch/nvfp4 9310 passed / 0 failed / 6852 skipped (CI at b3415c1 was 9309 passed / 1 failed / 6852 skipped), tests/pytorch/mxfp8 371 passed, and JAX nvfp4+gemm 382 passed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> Signed-off-by: tdophung <tdophung@nvidia.com> * Document bugs with `alpha_beta_on_device` GEMM config Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * [PyTorch] Do not swizzle operand scales in place in the NVFP4-UE5M3 GEMM (#5) * Do not swizzle operand scales in place in the NVFP4-UE5M3 GEMM The cuDNN UE5M3 GEMM wrapper swizzled A/B scales in place and marked the tensors as swizzled. For weights that mark persisted on the parameter. The master-weight cast (_cast_master_weights_to_nvfp4_2d, used by Megatron's fp4 param gather after every optimizer step) then rewrote the scales in the unswizzled layout without clearing the mark, so every later GEMM and dequantize read unswizzled scales as swizzled. Gaussian weights hide it (near-uniform scales, ~0.2 relative error); real DeepSeek-V3 weights are destroyed (2.6 relative error, loss 12.98 vs 7.86 for E4M3 under identical routing). Swizzle clones of the scale tensors instead and leave the operands untouched, matching what the cuBLAS path does in C++. * Test that the NVFP4-UE5M3 GEMM leaves operand scales untouched * Avoid suppressing errors in swizzle function. Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --------- Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Return the swizzled scales the UE5M3 GEMM helper actually produced (#6) tex.swizzle_scales_for_gemm_ assigns new tensors to the scale attributes instead of writing into the buffers it is given, so returning the clones made before the call hands the cuDNN GEMM unswizzled scales. On DeepSeek-V3 671B this turned the forward into garbage: routing collapsed onto expert 0 and the EP-rank-0 GPUs ran out of memory in the first expert forward. Read the attributes back after the swizzle, and make the regression test check the GEMM result against the dequantized reference. * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Debug compilation errors Signed-off-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Debug compile error in NVRTC NVFP4 4over6 kernel Signed-off-by: Tim Moon <tmoon@nvidia.com> * Debug NVFP4 UE5M3 in grouped MLP tests Use list of NVFP4 variant names rather than repeating at callsites. Exclude SReLU variants with UE5M3. Signed-off-by: Tim Moon <tmoon@nvidia.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix typo Review suggestion from @greptile-apps. Signed-off-by: Tim Moon <tmoon@nvidia.com> * Debug NVFP4 4over6 tests Signed-off-by: Tim Moon <tmoon@nvidia.com> --------- Signed-off-by: Tim Moon <tmoon@nvidia.com> Signed-off-by: Kaining Zhong <kainingz@nvidia.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Michal Futrega <mfutrega@nvidia.com> Signed-off-by: tdophung <tdophung@nvidia.com> Co-authored-by: Teddy Do <tdophung@nvidia.com> Co-authored-by: Varun Thumbe <vthumbe@nvidia.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kaining Zhong <44538064+kainzhong@users.noreply.github.com> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Kaining Zhong <kainingz@nvidia.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Michal Futrega <michal.futrega@gmail.com> Co-authored-by: Claude Opus 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
2 participants
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.
Signed-off-by: Kirthi Shankar Sivamani ksivamani@nvidia.com