From fff2245cdcab9feade6cda49c5581a714c4fe9a7 Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Mon, 20 Apr 2026 11:45:47 -0700 Subject: [PATCH 001/180] Changed version to 2.16.0.dev0 Signed-off-by: Przemek Tredak --- build_tools/VERSION.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/build_tools/VERSION.txt b/build_tools/VERSION.txt index 34ab1df06..36334f690 100644 --- a/build_tools/VERSION.txt +++ b/build_tools/VERSION.txt @@ -1 +1 @@ -2.15.0.dev0 +2.16.0.dev0 From 264da2b99fa5e027a19159bded6f5b107d976281 Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Tue, 21 Apr 2026 02:24:11 +0200 Subject: [PATCH 002/180] [Common] Reduced padding kernel compilation time (#2827) * Reduced padding kernel compilation time Signed-off-by: Oleg Goncharov * Completely removed unroll for better performance Signed-off-by: Oleg Goncharov --------- Signed-off-by: Oleg Goncharov --- transformer_engine/common/util/padding.cu | 2 -- 1 file changed, 2 deletions(-) diff --git a/transformer_engine/common/util/padding.cu b/transformer_engine/common/util/padding.cu index 835923828..67f98a130 100644 --- a/transformer_engine/common/util/padding.cu +++ b/transformer_engine/common/util/padding.cu @@ -87,7 +87,6 @@ __global__ void __launch_bounds__(threads_per_block) multi_padding_kernel(MultiP // Note: Each thread loads n_iterations subtiles, casts to output // type, and transposes in registers. Type local_zero = static_cast(0.f); -#pragma unroll for (int iter = 0; iter < n_iterations; ++iter) { const int i1 = tidy + iter * bdimy; const int j1 = tidx; @@ -171,7 +170,6 @@ __global__ void __launch_bounds__(threads_per_block) multi_unpadding_kernel(Mult // Note: Each thread loads n_iterations subtiles, casts to output // type, and transposes in registers. Type local_zero = static_cast(0.f); -#pragma unroll for (int iter = 0; iter < n_iterations; ++iter) { const int i1 = tidy + iter * bdimy; const int j1 = tidx; From 2d92aa6aae029f3caef74c92d3991d7c1ca0db10 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Tue, 21 Apr 2026 15:44:02 -0400 Subject: [PATCH 003/180] [PyTorch] Fix cuteDSL kernel incorrect numerics when K is 64 aligned (#2905) Zero out padded region when swizzling via group quantize Signed-off-by: Kirthi Shankar Sivamani --- .../cast/mxfp8/group_quantize_mxfp8.cuh | 20 +++++++++++++------ .../pytorch/ops/fused/backward_grouped_mlp.py | 6 ++++-- 2 files changed, 18 insertions(+), 8 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index ce6917aa4..ce827d24e 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -109,11 +109,15 @@ __device__ __forceinline__ void process_colwise_stage( const size_t global_scales_offset_Y = scales_offset_Y_colwise + stage; const size_t global_scales_offset_X = scales_offset_X_colwise; + const bool colwise_scale_is_within_bounds = global_scales_offset_X < cols; + size_t scale_idx = 0; if constexpr (WITH_GEMM_SWIZZLED_SCALES) { const size_t tensor_base_row = tensor_base_for_scales / cols; const size_t tensor_scales_offset_Y_base = tensor_base_row / SCALE_DIM_Y; - const size_t tensor_scales_offset_colwise_base = tensor_base_for_scales / SCALE_DIM_Y; + const size_t cols_padded = DIVUP(cols, static_cast(scale_tensor_alignment_X_colwise)) * + static_cast(scale_tensor_alignment_X_colwise); + const size_t tensor_scales_offset_colwise_base = tensor_base_row * cols_padded / SCALE_DIM_Y; const size_t local_scales_offset_Y = global_scales_offset_Y - tensor_scales_offset_Y_base; scale_idx = tensor_scales_offset_colwise_base + transformer_engine::dispatch::mxfp8::swizzle::gemm_swizzled_scale_idx( @@ -164,7 +168,9 @@ __device__ __forceinline__ void process_colwise_stage( const e8m0_t biased_exponent = ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); - scales_colwise[scale_idx] = biased_exponent; + // OOB padded region needs to be zeroed out. + scales_colwise[scale_idx] = + colwise_scale_is_within_bounds ? biased_exponent : static_cast(0); const bf16 block_scale_inverse = ptx::exp2f_rcp(biased_exponent); const ptx::bf16x2 block_scale_inverse_bf16_x2 = {block_scale_inverse, block_scale_inverse}; @@ -234,7 +240,9 @@ __device__ __forceinline__ void process_colwise_stage( const e8m0_t biased_exponent = ptx::float_to_e8m0(thread_amax * Quantized_Limits::max_norm_rcp); - scales_colwise[scale_idx] = biased_exponent; + // OOB padded region needs to be zeroed out. + scales_colwise[scale_idx] = + colwise_scale_is_within_bounds ? biased_exponent : static_cast(0); const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); #pragma unroll @@ -393,9 +401,9 @@ __device__ __forceinline__ void process_rowwise_stage( } else { scale_idx = stage_scales_offset_Y * scale_stride_rowwise + stage_scales_offset_X; } - if (rowwise_scale_is_within_bounds) { - scales_rowwise[scale_idx] = biased_exponent; - } + // OOB padded region needs to be zeroed out. + scales_rowwise[scale_idx] = + rowwise_scale_is_within_bounds ? biased_exponent : static_cast(0); const bf16 block_scale_inverse_bf16 = ptx::exp2f_rcp(biased_exponent); const ptx::bf16x2 block_scale_inverse_bf16_x2 = {block_scale_inverse_bf16, diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 3eb57c356..fc69b522d 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -63,10 +63,12 @@ def _cudnn_compute_wgrad( # b_tensor = X = (total_tokens, in_features) column-major b_tensor = grouped_x.columnwise_data.view(dtype=fp8_dtype).view(total_tokens, in_features) - sfa_tensor = grouped_dy.columnwise_scale_inv.view(out_features, -1).view( + sfa_leading_dim = ((out_features + 127) // 128) * 128 + sfb_leading_dim = ((in_features + 127) // 128) * 128 + sfa_tensor = grouped_dy.columnwise_scale_inv.view(sfa_leading_dim, -1).view( dtype=torch.float8_e8m0fnu ) - sfb_tensor = grouped_x.columnwise_scale_inv.view(in_features, -1).view( + sfb_tensor = grouped_x.columnwise_scale_inv.view(sfb_leading_dim, -1).view( dtype=torch.float8_e8m0fnu ) From 0e8ff35553f106400f94a2d6a69fbd26436b006d Mon Sep 17 00:00:00 2001 From: Santosh Bhavani Date: Tue, 21 Apr 2026 12:44:33 -0700 Subject: [PATCH 004/180] fix(readme): update broken links and modernize project description (#2879) * fix broken links in README Signed-off-by: Santosh Bhavani * update README to modernize description and standardize terminology Signed-off-by: Santosh Bhavani * Update README.rst Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Santosh Bhavani * Removed the duplicate line Signed-off-by: Przemek Tredak --------- Signed-off-by: Santosh Bhavani Signed-off-by: Przemek Tredak Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Przemek Tredak --- README.rst | 36 +++++++++++++++++------------------- 1 file changed, 17 insertions(+), 19 deletions(-) diff --git a/README.rst b/README.rst index e537b7a1f..884d8170e 100644 --- a/README.rst +++ b/README.rst @@ -38,21 +38,19 @@ precision-like API that can be used seamlessly with your framework-specific code framework agnostic C++ API that can be integrated with other deep learning libraries to enable FP8 support for Transformers. -As the number of parameters in Transformer models continues to grow, training and inference for -architectures such as BERT, GPT and T5 become very memory and compute-intensive. Most deep learning -frameworks train with FP32 by default. This is not essential, however, to achieve full accuracy for -many deep learning models. Using mixed-precision training, which combines single-precision (FP32) -with lower precision (e.g. FP16) format when training a model, results in significant speedups with -minimal differences in accuracy as compared to FP32 training. With Hopper GPU -architecture FP8 precision was introduced, which offers improved performance over FP16 with no -degradation in accuracy. Although all major deep learning frameworks support FP16, FP8 support is -not available natively in frameworks today. - -TE addresses the problem of FP8 support by providing APIs that integrate with popular Large Language -Model (LLM) libraries. It provides a Python API consisting of modules to easily build a Transformer -layer as well as a framework-agnostic library in C++ including structs and kernels needed for FP8 -support. Modules provided by TE internally maintain scaling factors and other values needed for FP8 -training, greatly simplifying mixed precision training for users. +As Transformer models scale to hundreds of billions of parameters across large language models, +MoE architectures, and multimodal models, training and inference become increasingly +memory and compute-intensive. Mixed-precision training, which combines single-precision (FP32) with +lower precision formats, delivers significant speedups with minimal impact on accuracy. FP8, introduced +with the Hopper GPU architecture, offers further performance gains over FP16 with no degradation in +accuracy, and newer formats like MXFP8 and NVFP4 on Blackwell push efficiency even further. + +TE integrates with popular LLM frameworks and provides optimizations that make low-precision training +work seamlessly with advanced features like MoE, tensor/sequence/context parallelism, and fused +operations. It provides a Python API consisting of modules to easily build a Transformer layer as +well as a framework-agnostic library in C++ including structs and kernels needed for FP8 support. +Modules provided by TE internally maintain scaling factors and other values needed for FP8 training, +greatly simplifying mixed precision training for users. Highlights ========== @@ -140,7 +138,7 @@ Flax for _ in range(10): loss, (param_grads, other_grads) = fwd_bwd_fn(params, other_variables, inp) -For a more comprehensive tutorial, check out our `Getting Started Guide `_. +For a more comprehensive tutorial, check out our `Getting Started Guide `_. .. overview-end-marker-do-not-remove @@ -383,7 +381,7 @@ FP8 and MXFP8 have been tested extensively across different model architectures +------------+------------------+---------------------------------------------------------------------------------------------------------+ | Model | Framework | Source | +============+==================+=========================================================================================================+ -| MPT-1.3B | Mosaic Composer | https://www.mosaicml.com/blog/coreweave-nvidia-h100-part-1 | +| MPT-1.3B | Mosaic Composer | https://www.databricks.com/blog/coreweave-nvidia-h100-part-1 | +------------+------------------+---------------------------------------------------------------------------------------------------------+ | LLama2-7B | Alibaba Pai | https://mp.weixin.qq.com/s/NQT0uKXLbXyh5031zBdeBQ | +------------+------------------+---------------------------------------------------------------------------------------------------------+ @@ -471,8 +469,8 @@ Previous News :alt: H200 * [11/2023] `Inflection-2: The Next Step Up `_ -* [11/2023] `Unleashing The Power Of Transformers With NVIDIA Transformer Engine `_ +* [11/2023] `Unleashing The Power Of Transformers With NVIDIA Transformer Engine `_ * [11/2023] `Accelerating PyTorch Training Workloads with FP8 `_ * [09/2023] `Transformer Engine added to AWS DL Container for PyTorch Training `_ * [06/2023] `Breaking MLPerf Training Records with NVIDIA H100 GPUs `_ -* [04/2023] `Benchmarking Large Language Models on NVIDIA H100 GPUs with CoreWeave (Part 1) `_ +* [04/2023] `Benchmarking Large Language Models on NVIDIA H100 GPUs with CoreWeave (Part 1) `_ From ee5dcec2258db9697d87c86c773a04d56b3023d7 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Tue, 21 Apr 2026 14:16:13 -0700 Subject: [PATCH 005/180] Add MXFP8 attention (#2719) * initial implementation for mxfp8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * semi-working FP8; broken F16 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up last commit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * comment out F16 pass Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * pull in grouped_quantize for MXFP8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * grouped tensor - pytorch Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * quantize mxfp8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix shapes/strides Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix unfused; clean up Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * split d to d_qk/d_v; attempt at bwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix last merge Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * attempt at SWA/MLA Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove prints Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove leftover prints Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "update FE" This reverts commit d9ff5662aa4b4b6267c77baf614aada6602fa133. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix MLA O strides; add bottom_right_diagonal Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * attempt at bwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix get_quantizers; attempt at bwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix fprop; add o_format Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * attempt at bwd with o_format/d_out_format/dqkv_layout Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix dtype/o_format/etc in bwd calls Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix generateMatrixStridesWithFormats and _v1; fix padding for mxfp8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix upon last commit for paddedsizes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add mxfp8 env var Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * disable FA for mxfp8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add mha test Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * attempt at bwd; force determinism; fix shapes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove prints Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE from pre-merge branch to post-merge develop Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * allow MXFP8 linear + f16 attn Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * test cp a2a Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove prints temporarily Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * test cp p2p Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fixes for mla Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * open up a2a for mla Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * test ag Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweaks for last commit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * enable mla ag Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix merge Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix merge Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert to main grouped tensor impl Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor tweaks to return to main Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove prints Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix combine_and_quantize for f16 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor tweaks Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix ds descale_o Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "fix ds descale_o" This reverts commit cd0bd82e239ff01210338b4e34cb8784109d22ec. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fixes for p2p and ag Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * tweak cp test skips Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update FE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix bwd KV tensors Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak recipe control and backend selection Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak quantizer logic Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fixes after last two commits Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * improve generate strides Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fixes for previous commit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix bwd for current/delayed Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak test configs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix dO/dO_f16 strides Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix tests: SWA logic/test configs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix ag Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add fp8 sink attn Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix a2a comm for F16 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove nan/inf print in test Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix fa a2a Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix fa a2a+p2p f16 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE to include new fixes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix thd for bwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * refactor a2a for fu/fa Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE to fix d64 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * refactor ag Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * refactor p2p/a2a+p2p; mostly regarding shapes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add shadow f16 fwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE to fix SWA/BRCM Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * switch to GH FE temporarily Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * switch back to GL FE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE to latest commit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update group tensor usage after merge main Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * env vars for qdq(q,k), o_f16 tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * allow other recipes than mxfp8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix grouped tensor for MLA Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * change cp test configs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add shadow f16 bwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix a2a+p2p for sbhd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix last commit and causal flag for fa Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * enable fp8 sink and disable fp8_mha Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor cleanup for cp/non-cp Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update FE for FP8 sink Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix TE for FP8 sink Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * temporary: random sink/print sink Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "temporary: random sink/print sink" This reverts commit 706095f802e04cbdd5d88ee53849cc5ec938203f. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * replace d_out_format with do_format Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix compare_and_assert for None cases Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove logic for b and simplify logic for dqkv types Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor fix for ndim_q/kv Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add explanation of fp8_output/grad in MHA Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tidy up FP8 checks for bhsd/learnable Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove leading underscores in nvte_convert_qkv_format Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * simplify logic in generateMatrixStridesWithLayout Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up strides/ifelse-recipe logic Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak checks in utils.py Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak UnfusedDPA Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * enable testing for ag+swa and disable fp8_mha Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak FusedAttn, fp8/f16 tensor naming/docstring Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * replace d_out_format with do_format Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up a2a Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up ag Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up p2p/a2a+p2p Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak test configs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * qdq dO in bwd shadow f16 path Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * tweak qdq dO logic Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove prints in shadow paths Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE to allow non-determinism Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fuse qkv transposes; first pass Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remap parallelism to grid(bh, splits, 3) block(s/splits x d); use nvec = 128 bits Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * allocate contiguous block for qkv Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix grouped tensor row/col scale_inv offsets Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * use fused permute kernels Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * quantize row/col as needed in fwd/bwd, non-cp/cp Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "quantize row/col as needed in fwd/bwd, non-cp/cp" This reverts commit ca5376956e8b8f662c7fa88661695b3e9eda4f8f. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Reapply "quantize row/col as needed in fwd/bwd, non-cp/cp" This reverts commit f19e852be3463210f2b3be5839ae8931e5ad92d0. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix v_col format when row is quantized Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add back necessary bwd quants for shadow paths/cp a2a Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove ZInv for all layouts except T3HD Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix cp p2p with zinv Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * temporarily switch to GH FE main Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * switch back to GL FE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix ag after merge main Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add condition for qdq(do) to not affect other tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix custom_mha_fp8 test Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix amax dqkv Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix fp8_recipe in DPA utils Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove use of amax for mxfp8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add o_format/do_format/dqkv_layout to cache indicators for fp8 and f16 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * enable sink attn + FP8 in CP Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE to GH v1.22.0 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix for inconsistent kwarg name in permute to grouped tensor Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add TMA permute Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "add TMA permute" This reverts commit 2532a50e829144bee290fc94acb8f3f154a62ea9. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * TMA load for bhsd transposes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix some lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * temp: quant+perm+swizzle, rope, perm_fused Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove mla_rope for now; clean up quant+permute+pad_swizzle; create multi_tensor_swizzle Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix last commit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * implement narrow-m for col swizzle; reorder to pad+perm+swizzle Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fused pad into perm; remove at::zeros as zeros done in perm kernels Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove shadow code Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor fix for permute shapes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * check smem size before entering narrow-k/m kernels Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * expand permute to multi_tensor_ Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * refactor qkv/do quant; create a fast_path call Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * cleanup grouped tensor fix Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove _with_amax for create_unquantized_tensor Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix last commit Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * reimplement inplace_multi_tensor_swizzle Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix last commit; set swizzled flag in python Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove permute_to_grouped_tensor_bwd; clean up fwd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add doxygen for multi_tensor_swizzle Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * clean up nvte_convert_qkv_format Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fixes based on code review Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * group layouts/formats in APIs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * rename nvte_convert_qkv_format to nvte_convert_qkv_shape Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove MXFP8 create_unquantized_tensor Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * rename permute_to_grouped_tensor to transpose_to_bhsd Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add multi_tensor_swizzle_xx_unchecked and split the calls/paths Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * straighten up indexing for multi_tensor_pad Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * batch up kernel calls per-16-tensors for pad and permute Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * remove nvec128; rename nvec64 back to nvec Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add Macros/arch specifics for compilation Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix lint Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * attempt 1: MLA RoPE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "attempt 1: MLA RoPE" This reverts commit 79229248718d26a0ae7029206adc26c687bb42a7. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix kv_cache tests for Fused, is_page=True Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * attempt 2: MLA RoPE Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * use DIVUP/_TO_MULTIPLE for pad_s_d_for_mxfp8 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove CUDNN_VERSION 8900 macros Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add narrow-k/m swizzle tests Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * compile flash_attn.cu with special archs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * Revert "attempt 2: MLA RoPE" This reverts commit 3b854b29a3677de2005fecff821d801ccd9bf5d4. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * make contiguous instead of check is_contiguous Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove unused s_q/s_kv Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove unused issue_tma_store_strided Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * add version gate for mxfp8 for CPP users Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * replace nvte_get_qkv_shape with AttentionShape Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * populate nvte_ changes to Jax Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * update FE to 1.22.1 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix minor merge issue Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert to FE 1.21 since it's what mxfp8 needs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * udpate jax attention shapes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Revert "revert to FE 1.21 since it's what mxfp8 needs" This reverts commit f09961a03bd7f5a316474b5d77b8292c7a49c1a6. Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * pick FE 1.22 to support mxfp8 and avoid rng issue in 1.22.1 Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * fix CP AG test on Hopper Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- 3rdparty/cudnn-frontend | 2 +- tests/cpp/operator/CMakeLists.txt | 1 + tests/cpp/operator/test_multi_swizzle.cu | 415 +++++ tests/cpp/operator/test_swizzle.cu | 8 + .../attention/run_attention_with_cp.py | 45 +- tests/pytorch/attention/test_attention.py | 101 +- .../attention/test_attention_with_cp.py | 240 +-- tests/pytorch/utils.py | 4 + transformer_engine/common/CMakeLists.txt | 2 +- transformer_engine/common/common.h | 2 +- .../common/fused_attn/flash_attn.cu | 753 +++++++- .../common/fused_attn/fused_attn.cpp | 179 +- .../fused_attn_f16_arbitrary_seqlen.cu | 105 +- .../fused_attn_f16_arbitrary_seqlen.h | 18 +- .../fused_attn_f16_max512_seqlen.cu | 2 - .../fused_attn/fused_attn_f16_max512_seqlen.h | 2 - .../common/fused_attn/fused_attn_fp8.cu | 969 +++++++--- .../common/fused_attn/fused_attn_fp8.h | 46 +- transformer_engine/common/fused_attn/utils.cu | 21 + transformer_engine/common/fused_attn/utils.h | 209 ++- .../include/transformer_engine/fused_attn.h | 108 +- .../include/transformer_engine/swizzle.h | 19 +- transformer_engine/common/swizzle/swizzle.cu | 646 +++++-- .../common/transformer_engine.cpp | 8 +- .../common/util/pybind_helper.h | 7 +- .../jax/csrc/extensions/attention.cpp | 128 +- .../dot_product_attention/backends.py | 351 ++-- .../dot_product_attention/context_parallel.py | 1598 ++++++++++++----- .../dot_product_attention.py | 99 +- .../attention/dot_product_attention/utils.py | 510 +++++- .../pytorch/attention/multi_head_attention.py | 26 +- .../pytorch/cpp_extensions/fused_attn.py | 80 +- transformer_engine/pytorch/csrc/extensions.h | 30 +- .../pytorch/csrc/extensions/attention.cpp | 302 +++- .../pytorch/csrc/extensions/pybind.cpp | 17 + .../pytorch/csrc/extensions/swizzle.cpp | 133 +- transformer_engine/pytorch/csrc/util.h | 3 + .../tensor/storage/grouped_tensor_storage.py | 28 +- 38 files changed, 5558 insertions(+), 1659 deletions(-) create mode 100644 tests/cpp/operator/test_multi_swizzle.cu diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend index 7b9b711c2..97f6cb3b8 160000 --- a/3rdparty/cudnn-frontend +++ b/3rdparty/cudnn-frontend @@ -1 +1 @@ -Subproject commit 7b9b711c22b6823e87150213ecd8449260db8610 +Subproject commit 97f6cb3b88cacff507cca1280db5650a457d92b3 diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index f83c4ae06..a5ea74171 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -32,6 +32,7 @@ add_executable(test_operator test_multi_unpadding.cu test_causal_softmax.cu test_swizzle.cu + test_multi_swizzle.cu test_swap_first_dims.cu test_grouped_gemm.cu ../test_common.cu) diff --git a/tests/cpp/operator/test_multi_swizzle.cu b/tests/cpp/operator/test_multi_swizzle.cu new file mode 100644 index 000000000..4984b7783 --- /dev/null +++ b/tests/cpp/operator/test_multi_swizzle.cu @@ -0,0 +1,415 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include + +#include "../test_common.h" +#include "transformer_engine/transformer_engine.h" + +using namespace transformer_engine; + +template +void compute_ref_swizzle(const uint8_t *h_input, uint8_t *h_output, + const size_t M, const size_t K) { + constexpr int NEW_SF_TILE_DIM_M = SF_TILE_DIM_M / 4; + constexpr int NEW_SF_TILE_DIM_K = SF_TILE_DIM_K * 4; + constexpr int SF_TILE_SIZE = SF_TILE_DIM_M * SF_TILE_DIM_K; + + for (size_t m = 0; m < M; m++) { + for (size_t k = 0; k < K; k++) { + int tile_id_m = m / SF_TILE_DIM_M; + int tile_id_k = k / SF_TILE_DIM_K; + int m_in_tile = m % SF_TILE_DIM_M; + int k_in_tile = k % SF_TILE_DIM_K; + + int row_in_new_tile = m_in_tile % NEW_SF_TILE_DIM_M; + int col_in_new_tile = m_in_tile / NEW_SF_TILE_DIM_M * SF_TILE_DIM_K + k_in_tile; + + int tile_output_ptr = tile_id_m * SF_TILE_DIM_M * K + tile_id_k * SF_TILE_SIZE; + int out_index = tile_output_ptr + row_in_new_tile * NEW_SF_TILE_DIM_K + col_in_new_tile; + if constexpr (row_scaling) + h_output[out_index] = h_input[k + m * K]; + else + h_output[out_index] = h_input[k * M + m]; + } + } +} + +static void zero_scale_inv_padding(uint8_t *buf, + size_t padded_rows, size_t padded_cols, + size_t orig_rows, size_t orig_cols) { + for (size_t r = 0; r < padded_rows; ++r) { + for (size_t c = 0; c < padded_cols; ++c) { + if (r >= orig_rows || c >= orig_cols) { + buf[r * padded_cols + c] = 0; + } + } + } +} + +// =================================================================== +// Multi-tensor swizzle test +// =================================================================== + +void performTestMultiTensorSwizzle(const int num_tensors, const size_t M, const size_t K, + bool rowwise) { + using namespace test; + constexpr size_t BLOCK_SIZE = 32; + const std::vector shape{M, K}; + + std::vector> input_tensors; + std::vector> output_tensors; + std::vector input_handles; + std::vector output_handles; + + for (int i = 0; i < num_tensors; ++i) { + auto input = std::make_unique("input_" + std::to_string(i), shape, + DType::kFloat8E4M3, rowwise, !rowwise, + NVTE_MXFP8_1D_SCALING); + auto output = std::make_unique("output_" + std::to_string(i), shape, + DType::kFloat8E4M3, rowwise, !rowwise, + NVTE_MXFP8_1D_SCALING); + fillUniform(input.get()); + output->set_with_gemm_swizzled_scales(true); + + input->to_cpu(); + if (rowwise) { + const NVTEShape rs = input->rowwise_scale_inv_shape(); + zero_scale_inv_padding(input->rowwise_cpu_scale_inv_ptr(), + rs.data[0], rs.data[1], + M, (K + BLOCK_SIZE - 1) / BLOCK_SIZE); + } else { + const NVTEShape cs = input->columnwise_scale_inv_shape(); + zero_scale_inv_padding(input->columnwise_cpu_scale_inv_ptr(), + cs.data[0], cs.data[1], + (M + BLOCK_SIZE - 1) / BLOCK_SIZE, K); + } + input->from_cpu(); + + input_handles.push_back(input->data()); + output_handles.push_back(output->data()); + input_tensors.emplace_back(std::move(input)); + output_tensors.emplace_back(std::move(output)); + } + + nvte_multi_tensor_swizzle_scaling_factors(input_handles.data(), output_handles.data(), + num_tensors, 0); + + cudaDeviceSynchronize(); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + for (int i = 0; i < num_tensors; ++i) { + output_tensors[i]->to_cpu(); + if (rowwise) { + const NVTEShape rs = input_tensors[i]->rowwise_scale_inv_shape(); + const size_t numel = rs.data[0] * rs.data[1]; + std::unique_ptr ref = std::make_unique(numel); + compute_ref_swizzle<128, 4, true>( + input_tensors[i]->rowwise_cpu_scale_inv_ptr(), + ref.get(), rs.data[0], rs.data[1]); + compareResults("multi_tensor_swizzle_row_" + std::to_string(i), + output_tensors[i]->rowwise_cpu_scale_inv_ptr(), + ref.get(), numel); + } else { + const NVTEShape cs = input_tensors[i]->columnwise_scale_inv_shape(); + const size_t numel = cs.data[0] * cs.data[1]; + std::unique_ptr ref = std::make_unique(numel); + compute_ref_swizzle<128, 4, false>( + input_tensors[i]->columnwise_cpu_scale_inv_ptr(), + ref.get(), cs.data[1], cs.data[0]); + compareResults("multi_tensor_swizzle_col_" + std::to_string(i), + output_tensors[i]->columnwise_cpu_scale_inv_ptr(), + ref.get(), numel); + } + } +} + +// =================================================================== +// Multi-tensor unswizzle test (uses single-tensor swizzle to prepare) +// =================================================================== + +void performTestMultiTensorUnswizzle(const int num_tensors, const size_t M, const size_t K, + bool rowwise) { + using namespace test; + constexpr size_t BLOCK_SIZE = 32; + const std::vector shape{M, K}; + + std::vector> orig_tensors, swizzled_tensors, output_tensors; + std::vector swizzled_handles, output_handles; + + for (int i = 0; i < num_tensors; ++i) { + auto orig = std::make_unique("orig_" + std::to_string(i), shape, + DType::kFloat8E4M3, rowwise, !rowwise, + NVTE_MXFP8_1D_SCALING); + auto swizzled = std::make_unique("swizzled_" + std::to_string(i), shape, + DType::kFloat8E4M3, rowwise, !rowwise, + NVTE_MXFP8_1D_SCALING); + auto output = std::make_unique("output_" + std::to_string(i), shape, + DType::kFloat8E4M3, rowwise, !rowwise, + NVTE_MXFP8_1D_SCALING); + fillUniform(orig.get()); + swizzled->set_with_gemm_swizzled_scales(true); + + orig->to_cpu(); + if (rowwise) { + const NVTEShape rs = orig->rowwise_scale_inv_shape(); + zero_scale_inv_padding(orig->rowwise_cpu_scale_inv_ptr(), + rs.data[0], rs.data[1], + M, (K + BLOCK_SIZE - 1) / BLOCK_SIZE); + } else { + const NVTEShape cs = orig->columnwise_scale_inv_shape(); + zero_scale_inv_padding(orig->columnwise_cpu_scale_inv_ptr(), + cs.data[0], cs.data[1], + (M + BLOCK_SIZE - 1) / BLOCK_SIZE, K); + } + orig->from_cpu(); + + nvte_swizzle_scaling_factors(orig->data(), swizzled->data(), 0); + + swizzled_handles.push_back(swizzled->data()); + output_handles.push_back(output->data()); + orig_tensors.emplace_back(std::move(orig)); + swizzled_tensors.emplace_back(std::move(swizzled)); + output_tensors.emplace_back(std::move(output)); + } + + nvte_multi_tensor_unswizzle_scaling_factors(swizzled_handles.data(), output_handles.data(), + num_tensors, 0); + + cudaDeviceSynchronize(); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + for (int i = 0; i < num_tensors; ++i) { + orig_tensors[i]->to_cpu(); + output_tensors[i]->to_cpu(); + if (rowwise) { + const NVTEShape rs = orig_tensors[i]->rowwise_scale_inv_shape(); + const size_t numel = rs.data[0] * rs.data[1]; + compareResults("multi_unswizzle_row_" + std::to_string(i), + output_tensors[i]->rowwise_cpu_scale_inv_ptr(), + orig_tensors[i]->rowwise_cpu_scale_inv_ptr(), + numel); + } else { + const NVTEShape cs = orig_tensors[i]->columnwise_scale_inv_shape(); + const size_t numel = cs.data[0] * cs.data[1]; + compareResults("multi_unswizzle_col_" + std::to_string(i), + output_tensors[i]->columnwise_cpu_scale_inv_ptr(), + orig_tensors[i]->columnwise_cpu_scale_inv_ptr(), + numel); + } + } +} + +// =================================================================== +// Multi-tensor swizzle -> unswizzle roundtrip test +// =================================================================== + +void performTestMultiTensorRoundtrip(const int num_tensors, const size_t M, const size_t K, + bool rowwise) { + using namespace test; + constexpr size_t BLOCK_SIZE = 32; + const std::vector shape{M, K}; + + std::vector> orig_tensors, mid_tensors, final_tensors; + std::vector orig_handles, mid_handles, final_handles; + + for (int i = 0; i < num_tensors; ++i) { + auto orig = std::make_unique("orig_" + std::to_string(i), shape, + DType::kFloat8E4M3, rowwise, !rowwise, + NVTE_MXFP8_1D_SCALING); + auto mid = std::make_unique("mid_" + std::to_string(i), shape, + DType::kFloat8E4M3, rowwise, !rowwise, + NVTE_MXFP8_1D_SCALING); + auto fin = std::make_unique("fin_" + std::to_string(i), shape, + DType::kFloat8E4M3, rowwise, !rowwise, + NVTE_MXFP8_1D_SCALING); + fillUniform(orig.get()); + mid->set_with_gemm_swizzled_scales(true); + + orig->to_cpu(); + if (rowwise) { + const NVTEShape rs = orig->rowwise_scale_inv_shape(); + zero_scale_inv_padding(orig->rowwise_cpu_scale_inv_ptr(), + rs.data[0], rs.data[1], + M, (K + BLOCK_SIZE - 1) / BLOCK_SIZE); + } else { + const NVTEShape cs = orig->columnwise_scale_inv_shape(); + zero_scale_inv_padding(orig->columnwise_cpu_scale_inv_ptr(), + cs.data[0], cs.data[1], + (M + BLOCK_SIZE - 1) / BLOCK_SIZE, K); + } + orig->from_cpu(); + + orig_handles.push_back(orig->data()); + mid_handles.push_back(mid->data()); + final_handles.push_back(fin->data()); + orig_tensors.emplace_back(std::move(orig)); + mid_tensors.emplace_back(std::move(mid)); + final_tensors.emplace_back(std::move(fin)); + } + + nvte_multi_tensor_swizzle_scaling_factors(orig_handles.data(), mid_handles.data(), + num_tensors, 0); + nvte_multi_tensor_unswizzle_scaling_factors(mid_handles.data(), final_handles.data(), + num_tensors, 0); + + cudaDeviceSynchronize(); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + for (int i = 0; i < num_tensors; ++i) { + orig_tensors[i]->to_cpu(); + final_tensors[i]->to_cpu(); + if (rowwise) { + const NVTEShape rs = orig_tensors[i]->rowwise_scale_inv_shape(); + const size_t numel = rs.data[0] * rs.data[1]; + compareResults("multi_roundtrip_row_" + std::to_string(i), + final_tensors[i]->rowwise_cpu_scale_inv_ptr(), + orig_tensors[i]->rowwise_cpu_scale_inv_ptr(), + numel); + } else { + const NVTEShape cs = orig_tensors[i]->columnwise_scale_inv_shape(); + const size_t numel = cs.data[0] * cs.data[1]; + compareResults("multi_roundtrip_col_" + std::to_string(i), + final_tensors[i]->columnwise_cpu_scale_inv_ptr(), + orig_tensors[i]->columnwise_cpu_scale_inv_ptr(), + numel); + } + } +} + +// =================================================================== +// Test suites +// =================================================================== + +class MultiTensorSwizzleTestSuite + : public ::testing::TestWithParam> {}; + +TEST_P(MultiTensorSwizzleTestSuite, TestMultiTensorSwizzle) { + const auto num_tensors = std::get<0>(GetParam()); + const auto M = std::get<1>(GetParam()); + const auto K = std::get<2>(GetParam()); + const auto rowwise = std::get<3>(GetParam()); + performTestMultiTensorSwizzle(num_tensors, M, K, rowwise); +} + +class MultiTensorUnswizzleTestSuite + : public ::testing::TestWithParam> {}; + +TEST_P(MultiTensorUnswizzleTestSuite, TestMultiTensorUnswizzle) { + const auto num_tensors = std::get<0>(GetParam()); + const auto M = std::get<1>(GetParam()); + const auto K = std::get<2>(GetParam()); + const auto rowwise = std::get<3>(GetParam()); + performTestMultiTensorUnswizzle(num_tensors, M, K, rowwise); +} + +class MultiTensorRoundtripTestSuite + : public ::testing::TestWithParam> {}; + +TEST_P(MultiTensorRoundtripTestSuite, TestMultiTensorRoundtrip) { + const auto num_tensors = std::get<0>(GetParam()); + const auto M = std::get<1>(GetParam()); + const auto K = std::get<2>(GetParam()); + const auto rowwise = std::get<3>(GetParam()); + performTestMultiTensorRoundtrip(num_tensors, M, K, rowwise); +} + +namespace { + +// Shapes that exercise the narrow_k kernel (rowwise) / narrow_m kernel (colwise): +// Narrow-K fires when ALL tensors have scale num_tiles_k < TB_DIM (32), +// i.e. padded ceil(K/32) < 128. +// Narrow-M fires analogously for colwise when padded K < 4096 +// (since colwise m = K padded to 128, num_tiles_m = m / 128 < 32). +// +// Shapes that bypass narrow and use the regular multi_tensor kernel: +// K >= 4096 makes num_tiles_k >= 32 (rowwise) and num_tiles_m >= 32 (colwise). + +std::vector> multi_tensor_test_cases = { + // --- Narrow path cases (K small → narrow_k for row, narrow_m for col) --- + // M and K both aligned to 128 + {3, 256, 256, true}, + {3, 256, 256, false}, + {4, 128, 128, true}, + {4, 128, 128, false}, + // M not divisible by 128 (but must be divisible by 32 for colwise — + // the kernel computes original_K = M / BLOCK_SIZE using floor division) + {3, 192, 256, true}, + {3, 192, 256, false}, + {2, 64, 256, true}, + {2, 64, 256, false}, + // Larger narrow K (num_tiles_k = 8, shared mem = 128 KB) + {2, 128, 1024, true}, + {2, 128, 1024, false}, + // K not divisible by 128 + {3, 256, 160, true}, + {3, 256, 160, false}, + // Neither M nor K divisible by 128 + {3, 192, 160, true}, + {3, 192, 160, false}, + // Minimum sizes (M=32 is the MXFP8 block size minimum for colwise) + {2, 32, 32, true}, + {2, 32, 32, false}, + {4, 32, 64, true}, + {4, 32, 64, false}, + + // --- Non-narrow path cases (K >= 4096 → regular multi_tensor kernel) --- + {3, 256, 4096, true}, + {3, 256, 4096, false}, + {2, 128, 8192, true}, + {2, 128, 8192, false}, +}; + +} // namespace + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + MultiTensorSwizzleTestSuite, + ::testing::ValuesIn(multi_tensor_test_cases), + [](const testing::TestParamInfo& info) { + return "n" + std::to_string(std::get<0>(info.param)) + + "_M" + std::to_string(std::get<1>(info.param)) + + "_K" + std::to_string(std::get<2>(info.param)) + + (std::get<3>(info.param) ? "_row" : "_col"); + }); + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + MultiTensorUnswizzleTestSuite, + ::testing::ValuesIn(multi_tensor_test_cases), + [](const testing::TestParamInfo& info) { + return "n" + std::to_string(std::get<0>(info.param)) + + "_M" + std::to_string(std::get<1>(info.param)) + + "_K" + std::to_string(std::get<2>(info.param)) + + (std::get<3>(info.param) ? "_row" : "_col"); + }); + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + MultiTensorRoundtripTestSuite, + ::testing::ValuesIn(multi_tensor_test_cases), + [](const testing::TestParamInfo& info) { + return "n" + std::to_string(std::get<0>(info.param)) + + "_M" + std::to_string(std::get<1>(info.param)) + + "_K" + std::to_string(std::get<2>(info.param)) + + (std::get<3>(info.param) ? "_row" : "_col"); + }); diff --git a/tests/cpp/operator/test_swizzle.cu b/tests/cpp/operator/test_swizzle.cu index 806a2482a..1ea82f19c 100644 --- a/tests/cpp/operator/test_swizzle.cu +++ b/tests/cpp/operator/test_swizzle.cu @@ -613,6 +613,14 @@ std::vector> num_tiles = { {65, 257}, {65, 258}, {65, 259}, + // Additional narrow-path coverage: narrow_k (row) when num_tiles_K < 32, + // narrow_m (col) when num_tiles_M < 32. + {1, 4}, // narrow_k with 4 K-tiles + {1, 8}, // narrow_k with 8 K-tiles + {4, 1}, // narrow_m with 4 M-tiles + {8, 1}, // narrow_m with 8 M-tiles + {31, 1}, // narrow_m at boundary (31 < TB_DIM=32) + {1, 31}, // narrow_k at boundary (31 < TB_DIM=32) }; // Raw {M, K} data shapes for unswizzle tests. Includes aligned cases (scale dims diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 0f36a8816..8dfea644a 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -19,8 +19,14 @@ DotProductAttention, Float8Quantizer, Float8CurrentScalingQuantizer, + MXFP8Quantizer, +) +from transformer_engine.common.recipe import ( + DelayedScaling, + Float8CurrentScaling, + MXFP8BlockScaling, + Format, ) -from transformer_engine.common.recipe import DelayedScaling, Float8CurrentScaling from utils import ModelConfig, compare_and_assert dtypes = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp8": torch.bfloat16} @@ -180,6 +186,7 @@ def run_dpa_with_cp( scaling_mode="delayed", f16_O="False", is_training="True", + deterministic="False", log_level=logging.WARNING, ): """Test DotProductAttention module with context parallelism""" @@ -188,11 +195,15 @@ def run_dpa_with_cp( is_training = is_training == "True" # set up environment variables and config + if deterministic == "True": + os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "0" + else: + os.environ["NVTE_ALLOW_NONDETERMINISTIC_ALGO"] = "1" fp8_bwd = fp8_bwd == "True" and dtype == "fp8" os.environ["NVTE_FP8_DPA_BWD"] = "1" if fp8_bwd else "0" fp8_dpa = fp8_dpa == "True" and dtype == "fp8" - fp8_mha = fp8_mha == "True" and dtype == "fp8" - f16_O = dtype == "fp8" and scaling_mode == "current" and f16_O == "True" + fp8_mha = fp8_mha == "True" and dtype == "fp8" and scaling_mode != "mxfp8" + f16_O = dtype == "fp8" and scaling_mode in ["current", "mxfp8"] and f16_O == "True" os.environ["NVTE_DPA_FP8CS_O_in_F16"] = "1" if f16_O else "0" os.environ["NVTE_FLASH_ATTN"] = "0" os.environ["NVTE_FUSED_ATTN"] = "0" @@ -247,6 +258,8 @@ def run_dpa_with_cp( fp8_recipe = DelayedScaling(fp8_dpa=fp8_dpa, fp8_mha=fp8_mha) if scaling_mode == "current": fp8_recipe = Float8CurrentScaling(fp8_dpa=fp8_dpa, fp8_mha=fp8_mha) + if scaling_mode == "mxfp8": + fp8_recipe = MXFP8BlockScaling(fp8_format=Format.E4M3, fp8_dpa=fp8_dpa, fp8_mha=fp8_mha) # instantiate attention module core_attn = DotProductAttention( @@ -302,10 +315,25 @@ def run_dpa_with_cp( fp8_dtype=tex.DType.kFloat8E5M2, device="cuda", ) + if scaling_mode == "mxfp8": + qkv_quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + ) + qkv_quantizer.optimize_for_gemm = True + qkv_quantizer.internal = False + dout_quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E5M2, + rowwise=True, + columnwise=True, + ) + dout_quantizer.optimize_for_gemm = True + dout_quantizer.internal = False qkv_layout = "_".join([qkv_format] * 3) q, k, v, dout = [x.clone().detach() for x in [q_orig, k_orig, v_orig, dout_orig]] if fp8_mha: - q, k, v = combine_and_quantize(qkv_layout, q, k, v, qkv_quantizer) + q, k, v, qkv_layout, _ = combine_and_quantize(qkv_layout, q, k, v, qkv_quantizer) for x in [q, k, v]: x.requires_grad = True @@ -413,7 +441,7 @@ def run_dpa_with_cp( dout_quantizer.scale.fill_(1.0) dout_quantizer.amax.fill_(0.0) if fp8_mha: - q_, k_, v_ = combine_and_quantize(qkv_layout, q_, k_, v_, qkv_quantizer) + q_, k_, v_, qkv_layout, _ = combine_and_quantize(qkv_layout, q_, k_, v_, qkv_quantizer) if is_training: q_, k_, v_ = [x.requires_grad_() for x in [q_, k_, v_]] if bias_ is not None: @@ -494,6 +522,7 @@ def run_dpa_with_cp( # get outputs tensors = [out, dq, dk, dv, dbias, out_, dq_, dk_, dv_, dbias_] + names = ["out", "dq", "dk", "dv", "dbias", "out_cp", "dq_cp", "dk_cp", "dv_cp", "dbias_cp"] if fp8_mha: tensors_to_deq = [out, out_] if not fp8_bwd else tensors for i, tensor in enumerate(tensors_to_deq): @@ -502,11 +531,11 @@ def run_dpa_with_cp( tensors_to_deq[i] = tensor.dequantize() if not fp8_bwd: tensors[0], tensors[5] = tensors_to_deq - for tensor in tensors: + for i, tensor in enumerate(tensors): # dbias/dbias_ could be None, so skip check for it if tensor is not None: - assert torch.all(~torch.isnan(tensor)) - assert torch.all(~torch.isinf(tensor)) + assert torch.all(~torch.isnan(tensor)), f"{names[i]} contains NaN" + assert torch.all(~torch.isinf(tensor)), f"{names[i]} contains Inf" out, dq, dk, dv, dbias, out_, dq_, dk_, dv_, dbias_ = tensors ############ compare results between CP and no-CP ############ diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 38d8626b4..c9ea79144 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -1936,20 +1936,45 @@ def get_model(dtype, config): return outputs +attn_mask_type = "causal" model_configs_fp8_vs_f16 = { # test: ModelConfig(b, sq, hq, dqk) - "fp8_9": ModelConfig(2, 2048, 16, 128), - "fp8_10": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12), - "fp8_11": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4), - "fp8_12": ModelConfig(2, 2048, 16, 128, attn_mask_type="causal"), - "fp8_13": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="causal"), - "fp8_14": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal"), - "fp8_15": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding"), - "fp8_16": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="padding"), - "fp8_17": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding"), - "fp8_18": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding_causal"), - "fp8_19": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="padding_causal"), - "fp8_20": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding_causal"), + "fp8_9": ModelConfig( + 2, + 4096, + 128, + 192, + head_dim_v=128, + ), + "fp8_10": ModelConfig( + 1, + 4096, + 128, + 192, + head_dim_v=128, + attn_mask_type="causal", + ), + "fp8_11": ModelConfig( + 2, + 4096, + 128, + 192, + head_dim_v=128, + attn_mask_type="causal_bottom_right", + ), + "fp8_12": ModelConfig(2, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="causal"), + "fp8_13": ModelConfig(2, 8192, 32, 128, attn_mask_type="causal", window_size=(128, 0)), + "fp8_14": ModelConfig(2, 8192, 64, 64, num_gqa_groups=8, attn_mask_type="causal"), + "fp8_15": ModelConfig(2, 8192, 64, 64, attn_mask_type="causal", window_size=(128, 0)), + "fp8_16": ModelConfig( + 2, 8192, 64, 64, num_gqa_groups=8, attn_mask_type="causal", softmax_type="learnable" + ), + "fp8_17": ModelConfig( + 2, 8192, 64, 64, attn_mask_type="causal", window_size=(128, 0), softmax_type="learnable" + ), + "fp8_18": ModelConfig(1, 8192, 32, 128, num_gqa_groups=4, attn_mask_type="padding"), + "fp8_19": ModelConfig(2, 2048, 16, 128, attn_mask_type="padding_causal"), + "fp8_20": ModelConfig(2, 2048, 24, 128, num_gqa_groups=12, attn_mask_type="padding_causal"), } param_types_fp8_vs_f16 = [torch.float16, torch.bfloat16] @@ -1966,7 +1991,7 @@ def get_model(dtype, config): @pytest.mark.parametrize("fp8_dpa_bwd", [True, False]) @pytest.mark.parametrize("RoPE", [True, False]) @pytest.mark.parametrize("is_training", [True, False]) -@pytest.mark.parametrize("scaling_mode", ["delayed", "current"]) +@pytest.mark.parametrize("scaling_mode", ["delayed", "current", "mxfp8"]) def test_mha_fp8_vs_f16( dtype, model, @@ -1997,6 +2022,12 @@ def test_mha_fp8_vs_f16( fp8_dpa=True, fp8_mha=True, ) + elif scaling_mode == "mxfp8": + fp8_recipe = recipe.MXFP8BlockScaling( + fp8_format=recipe.Format.E4M3, + fp8_dpa=True, + fp8_mha=False, + ) fp8_meta = {} fp8_meta["recipe"] = fp8_recipe available_backends, _, _ = get_available_attention_backends( @@ -2216,7 +2247,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: @pytest.mark.parametrize("qkv_layout", qkv_layout_fp8_vs_f16) @pytest.mark.parametrize("fp8_dpa_bwd", [True, False]) @pytest.mark.parametrize("is_training", [True, False]) -@pytest.mark.parametrize("scaling_mode", ["delayed", "current"]) +@pytest.mark.parametrize("scaling_mode", ["delayed", "current", "mxfp8"]) def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scaling_mode): """Test DotProductAttention module in FP8""" config = model_configs_fp8_vs_f16[model] @@ -2248,6 +2279,12 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal fp8_format=recipe.Format.HYBRID, fp8_dpa=True, ) + elif scaling_mode == "mxfp8": + fp8_recipe = recipe.MXFP8BlockScaling( + fp8_format=recipe.Format.E4M3, + fp8_dpa=True, + fp8_mha=False, + ) fp8_meta = {} fp8_meta["recipe"] = fp8_recipe available_backends, _, _ = get_available_attention_backends( @@ -2319,7 +2356,7 @@ def test_dpa_fp8_vs_f16(dtype, model, qkv_layout, fp8_dpa_bwd, is_training, scal atol = 5e-1 rtol = 5e-2 rmse_tol = 0.11 - bwd_names = ["dq", "dk", "dv"] + bwd_names = ["dq", "dk", "dv", "d_softmax_offset"] if flash_attn_supported and fused_attn_supported_f16: logging.debug("========== {:^25s} ==========".format("flash fp8 vs fused f16:")) logging.debug("========== {:^25s} ==========".format("forward output")) @@ -2408,7 +2445,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: with quantized_model_init(enabled=fp8_dpa): dpa = DotProductAttention( config.num_heads, - config.head_dim_qk, + (config.head_dim_qk, config.head_dim_v), num_gqa_groups=config.num_gqa_groups, attention_dropout=config.dropout_p, sequence_parallel=False, @@ -2418,6 +2455,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: layer_number=1, attention_type="self", qkv_format=qkv_format, + softmax_type=config.softmax_type, ).to(dtype=dtype, device="cuda") if not is_training: dpa = dpa.eval() @@ -2453,7 +2491,8 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: "skv": config.max_seqlen_kv, "h": config.num_heads, "hg": config.num_gqa_groups, - "d": config.head_dim_qk, + "dqk": config.head_dim_qk, + "dv": config.head_dim_v, "t": cu_seqlens_q[-1], "tg": cu_seqlens_kv[-1], "3": 3, @@ -2469,6 +2508,10 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: layout = layout.replace("s", "skv") layout = layout.replace("h", "hg") layout = layout.replace("t", "tg") + if i == 2: + layout = layout.replace("d", "dv") + else: + layout = layout.replace("d", "dqk") tensor_shape = [dim_to_num[j] for j in layout.split("_")] if config.dropout_p == 0.0: tensor = torch.randn(tensor_shape, dtype=dtype, device="cuda") @@ -2493,6 +2536,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: qkv_format_kv = "_".join(qkv_format) qkv_format_kv = qkv_format_kv.replace("s", "sq") + qkv_format_kv = qkv_format_kv.replace("d", "dv") out_grad_shape = [dim_to_num[i] for i in qkv_format_kv.split("_")] out_grad_shape_new = [*out_grad_shape[:-2], out_grad_shape[-2] * out_grad_shape[-1]] out_grad = torch.randn(out_grad_shape_new, dtype=dtype, device="cuda") @@ -2503,6 +2547,7 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: inp[1], inp[2], qkv_format=qkv_format, + window_size=config.window_size, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, max_seqlen_q=config.max_seqlen_q, @@ -2510,14 +2555,16 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: attn_mask_type=config.attn_mask_type, checkpoint_core_attention=False, core_attention_bias_type=config.attn_bias_type, - fp8_output=fp8_dpa, ) if is_training: out.backward(out_grad) + d_softmax_offset = None + if is_training and config.softmax_type != "vanilla": + d_softmax_offset = dpa.softmax_offset.grad if is_training: - return out, (inp[0].grad, inp[1].grad, inp[2].grad) - return out, (None, None, None) + return out, (inp[0].grad, inp[1].grad, inp[2].grad, d_softmax_offset) + return out, (None, None, None, d_softmax_offset) model_configs_fp8 = { @@ -2769,6 +2816,8 @@ def forward( quantization_params=qkv_quantizer, use_split_accumulator=_2X_ACC_FPROP, ) + qkv_layout = "bs3hd" if cudnn_frontend_version == 1 else "t3hd" + o_format = "bshd" if cudnn_frontend_version == 1 else "thd" qkv = qkv.view(-1, 3, h, d) qkv_fp16 = qkv.dequantize().view(b, max_s, 3, h, d).contiguous() torch.save(qkv_fp16, "qkv.pt") @@ -2797,7 +2846,8 @@ def forward( attn_scale=None, dropout=p_dropout, fast_zero_fill=fast_zero_fill, - qkv_layout="bs3hd" if cudnn_frontend_version == 1 else "t3hd", + qkv_layout=qkv_layout, + o_format=o_format, attn_bias_type="no_bias", attn_mask_type=mask_type if cudnn_frontend_version == 1 else "padding", rng_gen=None, @@ -2820,6 +2870,8 @@ def forward( ctx.num_heads = num_heads ctx.mask_type = mask_type ctx.dtype = inp.dtype + ctx.qkv_layout = qkv_layout + ctx.o_format = o_format ctx.dQKV_quantizer = dQKV_quantizer ctx.dO_quantizer = dO_quantizer @@ -2837,7 +2889,6 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], (q, k, v, inp_fp8, qkv_weight_fp8, out) = restore_from_func_ctx(ctx) proj_dgrad = ctx.dO_quantizer(grad_output) - fp8_dtype_backward = get_fp8_te_dtype(ctx.fp8_meta["recipe"], fprop_tensor=False) dq, dk, dv, *rest = fused_attn_bwd( ctx.max_s, @@ -2850,7 +2901,6 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], out, proj_dgrad.view_as(out), ctx.qkv_dtype, - fp8_dtype_backward, ctx.aux_ctx_tensors, FusedAttnBackend["FP8"], None, @@ -2861,7 +2911,10 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], attn_scale=None, dropout=ctx.p_dropout, fast_zero_fill=ctx.fast_zero_fill, - qkv_layout="bs3hd" if cudnn_frontend_version == 1 else "t3hd", + qkv_layout=ctx.qkv_layout, + o_format=ctx.o_format, + do_format=ctx.o_format, + dqkv_layout=ctx.qkv_layout, attn_bias_type="no_bias", attn_mask_type=ctx.mask_type if cudnn_frontend_version == 1 else "padding", ) diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 5aaf67061..23d1bfdd8 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -17,6 +17,8 @@ from transformer_engine.common.recipe import ( DelayedScaling, Float8CurrentScaling, + MXFP8BlockScaling, + Format, ) from transformer_engine.pytorch.attention.dot_product_attention.utils import FlashAttentionUtils @@ -26,6 +28,12 @@ pytest_logging_level = logging.getLevelName(logging.root.level) +# Get determinism +_deterministic = ( + not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) + or torch.are_deterministic_algorithms_enabled() +) + # Initialize RNG state seed = 1234 torch.manual_seed(seed) @@ -39,13 +47,11 @@ "cp_1_1": ModelConfig(2, 4096, 12, 128), # MHA "cp_1_2": ModelConfig(2, 4096, 12, 128, attn_mask_type="causal", window_size=(512, 0)), # MHA "cp_1_3": ModelConfig(2, 4096, 12, 128, window_size=(512, 512)), # MHA - "cp_2_0": ModelConfig(2, 4096, 12, 128, num_gqa_groups=2, attn_mask_type="causal"), # GQA + "cp_2_0": ModelConfig(2, 4096, 32, 128, num_gqa_groups=4, attn_mask_type="causal"), # GQA "cp_2_1": ModelConfig(2, 4096, 12, 128, num_gqa_groups=2), # GQA - "cp_2_2": ModelConfig( - 2, 4096, 12, 128, num_gqa_groups=2, attn_mask_type="causal", window_size=(512, 0) - ), # GQA + "cp_2_2": ModelConfig(2, 4096, 32, 128, attn_mask_type="causal", window_size=(128, 0)), # GQA "cp_2_3": ModelConfig(2, 4096, 12, 128, num_gqa_groups=2, window_size=(512, 512)), # GQA - "cp_3_0": ModelConfig(2, 4096, 12, 192, attn_mask_type="causal", head_dim_v=128), # MLA + "cp_3_0": ModelConfig(2, 4096, 128, 192, attn_mask_type="causal", head_dim_v=128), # MLA "cp_3_1": ModelConfig(2, 4096, 12, 192, head_dim_v=128), # MLA "cp_3_2": ModelConfig( 2, 4096, 12, 192, attn_mask_type="causal", window_size=(512, 0), head_dim_v=128 @@ -73,7 +79,7 @@ def get_bash_arguments(num_gpus_per_node, **kwargs): qkv_formats = ["bshd", "sbhd", "thd"] cp_comm_types = ["p2p", "all_gather", "a2a", "a2a+p2p"] if test_essential: - configs = ["cp_1_0", "cp_1_2", "cp_2_1", "cp_3_2", "cp_3_3"] + configs = ["cp_2_0", "cp_2_2", "cp_3_0", "cp_3_3"] model_configs_flash_attn = {k: model_configs_flash_attn[k] for k in configs} dtypes = ["bf16"] qkv_formats = ["sbhd", "thd"] @@ -94,25 +100,34 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): config.context_parallel = True config.cp_comm_type = cp_comm_type - if "p2p" in cp_comm_type and config.window_size != (-1, 0) and config.window_size != (-1, -1): - pytest.skip("CP implementation with KV P2P does not support sliding window yet!") - if cp_comm_type == "all_gather" and config.attn_bias_type != "no_bias": - pytest.skip("CP implementation with KV all-gather does not support bias yet!") - if qkv_format == "thd": - if cp_comm_type == "all_gather": - pytest.skip("CP implementation with KV all-gather does not support THD format yet!") - if cp_comm_type == "a2a+p2p": - pytest.skip( - "CP implementation with QKVO A2A+P2P (Hierarchical A2A) does not support THD format" - " yet!" - ) - if "a2a" in cp_comm_type and config.attn_bias_type != "no_bias": - pytest.skip("CP implementation with QKVO A2A does not support bias yet!") - if "a2a" in cp_comm_type and (config.num_heads % 2 != 0 or config.num_gqa_groups % 2 != 0): + if config.attn_bias_type != "no_bias" and qkv_format == "thd": + pytest.skip("No support for bias with THD format!") + if config.attn_bias_type != "no_bias" and cp_comm_type in ["all_gather", "a2a", "a2a+p2p"]: + pytest.skip("No support for bias with cp_comm_type={all_gather, a2a, a2a+p2p}!") + + if qkv_format == "thd" and cp_comm_type in ["all_gather", "a2a+p2p"]: + pytest.skip("No support for THD format with cp_comm_type={all_gather, a2a+p2p}!") + + if ( + config.window_size != (-1, 0) + and config.window_size != (-1, -1) + and cp_comm_type + in [ + "p2p", + "a2a+p2p", + ] + ): + pytest.skip("No support for SWA with cp_comm_type={p2p, a2a+p2p}!") + + if cp_comm_type in ["a2a", "a2a+p2p"] and ( + config.num_heads % 2 != 0 or config.num_gqa_groups % 2 != 0 + ): pytest.skip( - f"CP implementation with QKVO A2A requires num_heads ({config.num_heads}) and" - f" num_gqa_groups ({config.num_gqa_groups}) to be divisible by cp_size (2)!" + f"cp_comm_type=a2a requires num_heads ({config.num_heads}) and" + f" num_gqa_groups ({config.num_gqa_groups}) divisible by 2!" ) + + # FlashAttention / CP implementation specific: MLA only with KV P2P if "p2p" not in cp_comm_type and config.head_dim_qk != config.head_dim_v: pytest.skip("MLA CP currently only support KV P2P!") dtypes = {"fp16": torch.float16, "bf16": torch.bfloat16} @@ -150,8 +165,22 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): 2, 4096, 12, 128, attn_bias_type="post_scale_bias", bias_shape="bhss" ), # MHA "cp_1_5": ModelConfig(2, 4096, 12, 128, attn_mask_type="causal", window_size=(512, 512)), # MHA - "cp_2_0": ModelConfig(2, 4096, 12, 128, num_gqa_groups=2, attn_mask_type="causal"), # GQA - "cp_2_1": ModelConfig(2, 4096, 12, 128, num_gqa_groups=2), # GQA + "cp_2_0": ModelConfig( + 2, + 4096, + 32, + 128, + num_gqa_groups=4, + attn_mask_type="causal", + ), # GQA + "cp_2_1": ModelConfig( + 2, + 4096, + 32, + 128, + attn_mask_type="causal", + window_size=(128, 0), + ), # GQA "cp_2_2": ModelConfig( 2, 4096, @@ -189,7 +218,7 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): 2, 4096, 12, 128, num_gqa_groups=2, attn_mask_type="causal", window_size=(512, 512) ), # GQA "cp_3_0": ModelConfig(2, 4096, 12, 128, attn_mask_type="causal", head_dim_v=64), # MLA - "cp_3_1": ModelConfig(2, 4096, 12, 128, head_dim_v=64), # MLA + "cp_3_1": ModelConfig(2, 4096, 128, 192, head_dim_v=128, attn_mask_type="causal"), # MLA "cp_3_2": ModelConfig( 2, 4096, 12, 128, attn_mask_type="causal", attn_bias_type="post_scale_bias", head_dim_v=64 ), # MLA @@ -206,6 +235,9 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): "cp_4_2": ModelConfig( 2, 4096, 64, 64, num_gqa_groups=8, attn_mask_type="causal", softmax_type="learnable" ), # GQA + "cp_4_3": ModelConfig( + 2, 4096, 64, 64, attn_mask_type="causal", window_size=(128, 0), softmax_type="learnable" + ), # GQA } @@ -215,16 +247,15 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): if test_essential: configs = [ "cp_1_0", - "cp_1_1", - "cp_1_4", - "cp_1_5", "cp_2_0", + "cp_2_1", "cp_2_2", - "cp_2_3", "cp_2_4", + "cp_3_1", "cp_3_2", "cp_3_4", "cp_4_2", + "cp_4_3", ] model_configs_fused_attn = {k: model_configs_fused_attn[k] for k in configs} dtypes = ["bf16", "fp8"] @@ -240,96 +271,81 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): @pytest.mark.parametrize("fp8_bwd", [True, False]) @pytest.mark.parametrize("fp8_mha", [True, False]) @pytest.mark.parametrize("fp8_dpa", [True, False]) -@pytest.mark.parametrize("scaling_mode", [None, "delayed", "current"]) +@pytest.mark.parametrize("scaling_mode", [None, "delayed", "current", "mxfp8"]) @pytest.mark.parametrize("f16_O", [True, False]) def test_cp_with_fused_attention( dtype, model, qkv_format, cp_comm_type, fp8_bwd, fp8_mha, fp8_dpa, scaling_mode, f16_O ): + config = model_configs_fused_attn[model] + config.context_parallel = True + config.cp_comm_type = cp_comm_type + num_gpus = 4 if cp_comm_type == "a2a+p2p" else 2 if num_gpus > torch.cuda.device_count(): - pytest.skip(f"Test requires {num_gpus} GPUs, but found {torch.cuda.device_count()}") + pytest.skip(f"Test requires {num_gpus} GPUs, but found {torch.cuda.device_count()} GPUs.") + + if get_device_compute_capability() < (9, 0) and qkv_format == "thd": + pytest.skip("Only sm90+ architectures support THD format!") + if get_device_compute_capability() < (9, 0) and dtype == "fp8": + pytest.skip("Only sm90+ architectures support FP8 attention!") - if qkv_format == "thd" and get_device_compute_capability() < (9, 0): - pytest.skip("THD format is only supported on sm90+!") - if cp_comm_type == "all_gather" and get_cudnn_version() < (9, 3, 0): - pytest.skip("CP implementation with KV all-gather is only supported with cuDNN >= 9.3.0!") - if dtype == "fp8" and get_device_compute_capability() < (9, 0): - pytest.skip("FP8 attention is only supported on sm90+!") + if dtype == "fp8" and not (fp8_mha or fp8_dpa): + pytest.skip("dtype=fp8 requires fp8_dpa=True or fp8_mha=True!") if dtype == "fp8" and not fp8_dpa and fp8_mha: pytest.skip("Duplicate tests to fp8_dpa=True and fp8_mha=True!") if dtype != "fp8" and fp8_bwd: - pytest.skip("Only fp8 works with fp8_bwd=True!") - - config = model_configs_fused_attn[model] - config.context_parallel = True - config.cp_comm_type = cp_comm_type + pytest.skip("fp8_bwd=True requires dtype=fp8!") + if dtype != "fp8" and (fp8_mha or fp8_dpa): + pytest.skip("dtype!=fp8 requires fp8_dpa=False and fp8_mha=False!") - if qkv_format == "thd" and config.attn_bias_type == "post_scale_bias": - pytest.skip("THD format does not support post_scale_bias yet!") - if qkv_format == "thd": - if cp_comm_type == "all_gather": - pytest.skip("CP implementation with KV all-gather does not support THD format yet!") - if cp_comm_type == "a2a+p2p": - pytest.skip( - "CP implementation with QKVO A2A+P2P (Hierarchical A2A) does not support THD format" - " yet!" - ) - if dtype == "fp8" and cp_comm_type == "all_gather": - pytest.skip( - "CP implementation with KV all-gather does not support FP8 + context parallelism yet!" - ) if dtype == "fp8" and qkv_format == "thd": - pytest.skip("FP8 attention cannot work with THD format yet!") + pytest.skip("No support for FP8 attention with THD format!") if dtype == "fp8" and config.attn_bias_type != "no_bias": - pytest.skip("FP8 attention cannot work with bias yet!") - if dtype == "fp8" and config.window_size != (-1, 0) and config.window_size != (-1, -1): - pytest.skip("FP8 attention cannot work with sliding window yet!") - if "p2p" in cp_comm_type and config.window_size != (-1, 0) and config.window_size != (-1, -1): - pytest.skip("CP implementation with KV P2P does not support sliding window yet!") - if cp_comm_type == "all_gather" and config.attn_bias_type != "no_bias": - pytest.skip("CP implementation with KV all-gather does not support bias yet!") - if "a2a" in cp_comm_type and config.attn_bias_type != "no_bias": - pytest.skip("CP implementation with QKVO A2A does not support bias yet!") - if "a2a" in cp_comm_type and (config.num_heads % 2 != 0 or config.num_gqa_groups % 2 != 0): - pytest.skip( - f"CP implementation with QKVO A2A requires num_heads ({config.num_heads}) and" - f" num_gqa_groups ({config.num_gqa_groups}) to be divisible by cp_size (2)!" - ) - if dtype != "fp8" and (fp8_mha or fp8_dpa): - pytest.skip("Only fp8 works with fp8_dpa=True or fp8_mha=True!") - if dtype == "fp8" and not (fp8_mha or fp8_dpa): - pytest.skip("fp8 only works with fp8_dpa=True or fp8_mha=True!") - if dtype != "fp8" and scaling_mode is not None: - pytest.skip("Only fp8 works with scaling_mode != None!") - if dtype == "fp8" and scaling_mode is None: - pytest.skip("fp8 only works with scaling_mode != None!") - if ( - dtype == "fp8" - and scaling_mode == "current" - and cp_comm_type not in ["p2p", "a2a+p2p", "a2a"] + pytest.skip("No support for FP8 attention with bias!") + + if config.attn_bias_type != "no_bias" and qkv_format == "thd": + pytest.skip("No support for bias with THD format!") + if config.attn_bias_type != "no_bias" and cp_comm_type in ["all_gather", "a2a", "a2a+p2p"]: + pytest.skip("No support for bias with cp_comm_type={all_gather, a2a, a2a+p2p}!") + + if qkv_format == "thd" and cp_comm_type in ["all_gather", "a2a+p2p"]: + pytest.skip("No support for THD format with cp_comm_type={all_gather, a2a+p2p}!") + + if (config.window_size[0] != -1 or config.window_size[1] not in [-1, 0]) and cp_comm_type in [ + "p2p", + "a2a+p2p", + ]: + pytest.skip("No support for SWA with cp_comm_type={p2p, a2a+p2p}!") + + if cp_comm_type in ["a2a", "a2a+p2p"] and ( + config.num_heads % 2 != 0 or config.num_gqa_groups % 2 != 0 ): - pytest.skip("fp8 only works with P2P, A2A and A2A+P2P for scaling_mode = current!") - if f16_O and (dtype != "fp8" or scaling_mode != "current"): - pytest.skip("f16_O only needs to be tested for dtype = fp8 and scaling_mode = current!") - if "p2p" not in cp_comm_type and config.head_dim_qk != config.head_dim_v: - pytest.skip("MLA CP currently only support KV P2P!") - if dtype == "fp8" and config.head_dim_qk != config.head_dim_v: - pytest.skip("MLA CP currently does not support FP8 attention!") - if dtype == "fp8" and config.softmax_type != "vanilla": - pytest.skip("CP implementation does not support non-vanilla softmax types in FP8!") - if config.softmax_type != "vanilla" and cp_comm_type != "a2a": pytest.skip( - "CP implementation only supports cp_comm_type=a2a for non-vanilla softmax types!" + f"cp_comm_type=a2a requires num_heads ({config.num_heads}) and" + f" num_gqa_groups ({config.num_gqa_groups}) divisible by 2!" ) + + if config.softmax_type != "vanilla" and cp_comm_type != "a2a": + pytest.skip(f"No support for non-vanilla softmax with cp_comm_type={cp_comm_type}!") if ( - get_cudnn_version() < (9, 18, 0) - and config.softmax_type != "vanilla" + config.softmax_type != "vanilla" and qkv_format == "thd" + and get_cudnn_version() < (9, 18, 0) ): - pytest.skip( - "Unless cudnn version >= 9.18.0, CP implementation does not support qkv_format=thd for" - " non-vanilla softmax types!" - ) + pytest.skip("No support for non-vanilla softmax with THD format and cuDNN < 9.18.0!") + + if dtype == "fp8" and scaling_mode is None: + pytest.skip("dtype=fp8 requires scaling_mode != None!") + if dtype != "fp8" and scaling_mode is not None: + pytest.skip("dtype!=fp8 requires scaling_mode = None!") + if dtype != "fp8" and not f16_O: + pytest.skip("dtype!=fp8 requires f16_O=True!") + if scaling_mode == "delayed" and f16_O: + pytest.skip("scaling_mode=delayed requires f16_O=False!") + if scaling_mode == "mxfp8" and not f16_O: + pytest.skip("scaling_mode=mxfp8 requires f16_O=True!") + if scaling_mode == "mxfp8" and fp8_mha: + pytest.skip("No support for scaling_mode=mxfp8 with fp8_mha=True!") dtypes = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp8": torch.bfloat16} @@ -353,6 +369,12 @@ def test_cp_with_fused_attention( Float8CurrentScaling(fp8_dpa=True), DelayedScaling(fp8_dpa=True), ] + if fp8 and scaling_mode == "mxfp8": + fp8_meta["recipe"] = MXFP8BlockScaling(fp8_format=Format.E4M3, fp8_dpa=True) + fp8_meta["local_recipes"] = [ + MXFP8BlockScaling(fp8_format=Format.E4M3, fp8_dpa=True), + ] + # For 111s, dbias calculation is not supported as of cuDNN 9.18, hence, test fwd only for 111s. is_training = False if config.bias_shape == "111s" else True available_backends, _, fused_attn_backends = get_available_attention_backends( @@ -362,8 +384,23 @@ def test_cp_with_fused_attention( fp8=fp8, fp8_meta=fp8_meta, is_training=is_training, + deterministic=_deterministic, ) _, fused_attn_supported, _ = available_backends + if fused_attn_supported and config.attn_mask_type in ["causal", "padding_causal"]: + config_copy = copy.deepcopy(config) + config_copy.context_parallel = False + config_copy.attn_mask_type = config.attn_mask_type + "_bottom_right" + available_backends, _, fused_attn_backends = get_available_attention_backends( + config_copy, + qkv_dtype=dtypes[dtype] if dtype != "fp8" else torch.float8_e4m3fn, + qkv_layout="_".join([qkv_format] * 3), + fp8=fp8, + fp8_meta=fp8_meta, + is_training=is_training, + deterministic=_deterministic, + ) + _, fused_attn_supported, _ = available_backends if not fused_attn_supported: pytest.skip("No attention backend available.") @@ -381,6 +418,7 @@ def test_cp_with_fused_attention( scaling_mode=scaling_mode, f16_O=f16_O, is_training=is_training, + deterministic=_deterministic, log_level=pytest_logging_level, ), ) diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index fd9a6416e..8f8852edc 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -198,6 +198,10 @@ def reset_rng_states() -> None: def compare_and_assert(a, b, name_a, name_b, atol, rtol, rmse_tol, is_fp8): + if a is None and b is None: + logging.debug(f"{name_a} vs {name_b}: both are None") + return + if not is_fp8: torch.testing.assert_close(a, b, atol=atol, rtol=rtol) return diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index a21c1ee7e..53f9773a7 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -179,7 +179,6 @@ list(APPEND transformer_engine_cuda_sources transpose/quantize_transpose_vector_blockwise.cu transpose/swap_first_dims.cu dropout/dropout.cu - fused_attn/flash_attn.cu fused_attn/context_parallel.cu fused_attn/kv_cache.cu fused_attn/fused_attn_f16_max512_seqlen.cu @@ -210,6 +209,7 @@ list(APPEND transformer_engine_cuda_sources comm_gemm_overlap/userbuffers/userbuffers.cu) list(APPEND transformer_engine_cuda_arch_specific_sources + fused_attn/flash_attn.cu activation/gelu.cu activation/glu.cu activation/relu.cu diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 6e207370d..68aa0f4c5 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -1003,7 +1003,7 @@ size_t typeToSize(const DType type); size_t typeToNumBits(const DType type); void CheckNoopTensor(const Tensor &t, const std::string &name); -void CheckInputTensor(const Tensor &t, const std::string &name); +void CheckInputTensor(const Tensor &t, const std::string &name, bool check_scale_inv_shapes = true); void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empty = false); /*! \brief Update a tensor's FP8 scale-inverse diff --git a/transformer_engine/common/fused_attn/flash_attn.cu b/transformer_engine/common/fused_attn/flash_attn.cu index 6c66746e6..5037be828 100644 --- a/transformer_engine/common/fused_attn/flash_attn.cu +++ b/transformer_engine/common/fused_attn/flash_attn.cu @@ -4,12 +4,30 @@ * See LICENSE for license information. ************************************************************************/ +#include +#include + #include "../common.h" +#include "../util/cuda_driver.h" +#include "../util/cuda_runtime.h" +#include "../util/ptx.cuh" +#include "../utils.cuh" #include "transformer_engine/fused_attn.h" namespace transformer_engine { + +// ============================================================================ +// prepare_flash_attn: SBH3D <-> BSHD_BSHD_BSHD for the FlashAttention backend +// ============================================================================ + namespace flash_attention { +/// Packed vector of N elements of T; alignment matches a single wide load/store of N * sizeof(T) bytes. +template +struct alignas(sizeof(T) * N) Vec { + T data[N]; +}; + constexpr int warp_size = 32; constexpr int type_size = 2; // FP16 or BF16 constexpr int nvec = sizeof(uint64_t) / type_size; @@ -35,8 +53,8 @@ __launch_bounds__(block_size) __global__ T *my_output = qkv + offset_output; for (int i = 0; i < Z; ++i) { - uint64_t *out = reinterpret_cast(my_output + i * load_size); - *out = *reinterpret_cast(my_input + i * load_size * 3); + Vec *const out = reinterpret_cast *>(my_output + i * load_size); + *out = *reinterpret_cast *>(my_input + i * load_size * 3); } } @@ -61,8 +79,8 @@ __launch_bounds__(block_size) __global__ T *my_output = qkv + offset_output; for (int i = 0; i < Z; ++i) { - uint64_t *out = reinterpret_cast(my_output + i * load_size * 3); - *out = *reinterpret_cast(my_input + i * load_size); + Vec *const out = reinterpret_cast *>(my_output + i * load_size * 3); + *out = *reinterpret_cast *>(my_input + i * load_size); } } @@ -134,6 +152,696 @@ void prepare_flash_attn_bwd(Tensor q, Tensor k, Tensor v, Tensor qkv, cudaStream } } // namespace flash_attention + +// ============================================================================ +// multi_tensor_transpose_to_bhsd: BSHD/SBHD -> BHSD +// ============================================================================ + +namespace multi_tensor_transpose_to_bhsd { + +using flash_attention::Vec; + +constexpr int kMaxPermuteTensors = 16; + +struct PermuteSlot { + const void *input; + void *output; + size_t S, H, D_in, D_out; +}; + +struct PermuteParams { + PermuteSlot slots[kMaxPermuteTensors]; +}; + +struct TmaMapParams { + CUtensorMap maps[kMaxPermuteTensors]; +}; + +// ---------- path 3: fallback_not_vec_aligned ---------- + +__device__ __forceinline__ void copy_row_bytes(const char *__restrict__ src, char *__restrict__ dst, + size_t D_bytes) { + size_t off = 0; + for (; off + 16 <= D_bytes; off += 16) { + uint4 tmp; + memcpy(&tmp, src + off, 16); + memcpy(dst + off, &tmp, 16); + } + for (; off + 8 <= D_bytes; off += 8) { + uint2 tmp; + memcpy(&tmp, src + off, 8); + memcpy(dst + off, &tmp, 8); + } + for (; off + 4 <= D_bytes; off += 4) { + unsigned int tmp; + memcpy(&tmp, src + off, 4); + memcpy(dst + off, &tmp, 4); + } + for (; off + 2 <= D_bytes; off += 2) { + uint16_t tmp; + memcpy(&tmp, src + off, 2); + memcpy(dst + off, &tmp, 2); + } + for (; off < D_bytes; ++off) dst[off] = src[off]; +} + +__device__ __forceinline__ void copy_and_pad_row_bytes(const char *__restrict__ src, + char *__restrict__ dst, size_t D_bytes, + size_t D_out_bytes) { + copy_row_bytes(src, dst, D_bytes); + for (size_t off = D_bytes; off < D_out_bytes; ++off) dst[off] = 0; +} + +constexpr int TRANSPOSE_TILE = 32; +constexpr int TRANSPOSE_BLOCK = 256; +constexpr int TRANSPOSE_WARPS = TRANSPOSE_BLOCK / 32; // 8 + +template +__launch_bounds__(TRANSPOSE_BLOCK) __global__ + void transpose_to_bhsd_fallback_not_vec_aligned_kernel(PermuteParams params, size_t b, + unsigned int s_tiles) { + const auto &slot = params.slots[blockIdx.z]; + const T *__restrict__ in = reinterpret_cast(slot.input); + T *__restrict__ out = reinterpret_cast(slot.output); + const size_t S = slot.S; + const size_t H = slot.H; + const size_t D = slot.D_in; + const size_t D_out = slot.D_out; + const size_t D_bytes = D * sizeof(T); + const size_t D_out_bytes = D_out * sizeof(T); + const size_t D_smem_pad = (D_bytes + 3u) & ~size_t(3); + + const size_t tile_s = static_cast(blockIdx.x) % static_cast(s_tiles); + const size_t b_i = static_cast(blockIdx.x) / static_cast(s_tiles); + if (b_i >= b) return; + const size_t tile_h = static_cast(blockIdx.y); + + const size_t s_base = tile_s * TRANSPOSE_TILE; + const size_t h_base = tile_h * TRANSPOSE_TILE; + + extern __shared__ char smem[]; + const size_t smem_row = static_cast(TRANSPOSE_TILE) * D_smem_pad + 4; + + // ---- Phase 1: global → smem (sweep consecutive H → coalesced reads) ---- + for (unsigned int warp_off = threadIdx.x >> 5; warp_off < TRANSPOSE_TILE; + warp_off += TRANSPOSE_WARPS) { + const size_t local_s = warp_off; + const size_t local_h = threadIdx.x & 31u; + const size_t s_i = s_base + local_s; + const size_t h_i = h_base + local_h; + if (s_i < S && h_i < H) { + const char *__restrict__ src; + if constexpr (kIsBshd) + src = reinterpret_cast(in + b_i * S * H * D + s_i * H * D + h_i * D); + else + src = reinterpret_cast(in + s_i * b * H * D + b_i * H * D + h_i * D); + copy_row_bytes(src, smem + local_s * smem_row + local_h * D_smem_pad, D_bytes); + } + } + + __syncthreads(); + + // ---- Phase 2: smem → global (sweep consecutive S → coalesced writes, with padding) ---- + for (unsigned int warp_off = threadIdx.x >> 5; warp_off < TRANSPOSE_TILE; + warp_off += TRANSPOSE_WARPS) { + const size_t local_h = warp_off; + const size_t local_s = threadIdx.x & 31u; + const size_t s_i = s_base + local_s; + const size_t h_i = h_base + local_h; + if (s_i < S && h_i < H) { + copy_and_pad_row_bytes( + smem + local_s * smem_row + local_h * D_smem_pad, + reinterpret_cast(out + b_i * H * S * D_out + h_i * S * D_out + s_i * D_out), + D_bytes, D_out_bytes); + } + } +} + +// ---------- path 2: fallback_vec_aligned ---------- + +constexpr int fallback_permute_threads = 1024; + +template +__device__ __forceinline__ void permute_vec_loop(const T *__restrict__ in, T *__restrict__ out, + size_t b, size_t S, size_t H, size_t D, + size_t D_out, size_t b_i, size_t h_i, + size_t s_begin, size_t S_chunk) { + const size_t out_base = b_i * H * S * D_out + h_i * S * D_out; + const size_t d_vec = D / static_cast(N); + const size_t total_work = S_chunk * d_vec; + for (size_t w = static_cast(threadIdx.x); w < total_work; + w += static_cast(blockDim.x)) { + const size_t s_local = w / d_vec; + const size_t s_i = s_begin + s_local; + const size_t d_off = (w % d_vec) * static_cast(N); + const T *__restrict__ in_ptr; + if constexpr (kIsBshd) { + in_ptr = in + b_i * (S * H * D) + s_i * (H * D) + h_i * D + d_off; + } else { + in_ptr = in + s_i * (b * H * D) + b_i * (H * D) + h_i * D + d_off; + } + T *__restrict__ out_ptr = out + out_base + s_i * D_out + d_off; + *reinterpret_cast *>(out_ptr) = *reinterpret_cast *>(in_ptr); + } + if (D_out > D) { + const size_t pad_elems = D_out - D; + const size_t total_pad = S_chunk * pad_elems; + for (size_t w = static_cast(threadIdx.x); w < total_pad; + w += static_cast(blockDim.x)) { + const size_t s_local = w / pad_elems; + const size_t s_i = s_begin + s_local; + const size_t d_off = D + (w % pad_elems); + out[out_base + s_i * D_out + d_off] = static_cast(0); + } + } +} + +template +__launch_bounds__(fallback_permute_threads) __global__ + void transpose_to_bhsd_fallback_vec_aligned_kernel(PermuteParams params, size_t b, + unsigned int permute_s_splits, + size_t h_grid) { + const auto &slot = params.slots[blockIdx.z]; + const T *__restrict__ in = reinterpret_cast(slot.input); + T *__restrict__ out = reinterpret_cast(slot.output); + const size_t S = slot.S; + const size_t H = slot.H; + const size_t D = slot.D_in; + const size_t D_out = slot.D_out; + + const size_t b_i = static_cast(blockIdx.x) / h_grid; + const size_t h_i = static_cast(blockIdx.x) % h_grid; + if (b_i >= b) return; + if (h_i >= H) return; + + const unsigned int s_part = blockIdx.y; + const size_t s_begin = (S * static_cast(s_part)) / static_cast(permute_s_splits); + const size_t s_end = + (S * static_cast(s_part + 1)) / static_cast(permute_s_splits); + if (s_begin >= s_end) return; + const size_t S_chunk = s_end - s_begin; + + const size_t D_bytes = D * sizeof(T); + + if (D_bytes % 16 == 0) { + constexpr size_t N = 16 / sizeof(T); + permute_vec_loop(in, out, b, S, H, D, D_out, b_i, h_i, s_begin, S_chunk); + return; + } + if (D_bytes % 8 == 0) { + constexpr size_t N = 8 / sizeof(T); + permute_vec_loop(in, out, b, S, H, D, D_out, b_i, h_i, s_begin, S_chunk); + return; + } + if constexpr (sizeof(T) <= 4) { + if (D_bytes % 4 == 0) { + constexpr size_t N = 4 / sizeof(T); + permute_vec_loop(in, out, b, S, H, D, D_out, b_i, h_i, s_begin, S_chunk); + return; + } + } +} + +// ---------- path 1: TMA ---------- + +constexpr int tma_permute_threads = 128; +constexpr int tma_permute_s_tile_default = 32; + +__device__ __forceinline__ void cp_async_bulk_tensor_4d_global_to_shared( + void *dst_shmem, const CUtensorMap *tensor_map, uint32_t c0, uint32_t c1, uint32_t c2, + uint32_t c3, uint64_t *mbar) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t dst = __cvta_generic_to_shared(dst_shmem); + uint32_t bar = __cvta_generic_to_shared(mbar); + asm volatile( + "cp.async.bulk.tensor.4d.shared::cluster.global.tile" + ".mbarrier::complete_tx::bytes [%0], [%1, {%2, %3, %4, %5}], [%6];" ::"r"(dst), + "l"(tensor_map), "r"(c0), "r"(c1), "r"(c2), "r"(c3), "r"(bar) + : "memory"); +#else + NVTE_DEVICE_ERROR("cp_async_bulk_tensor_4d_global_to_shared requires SM 10.0+."); +#endif +} + +__device__ __forceinline__ void cp_async_bulk_tensor_4d_shared_to_global( + const CUtensorMap *tensor_map, uint32_t c0, uint32_t c1, uint32_t c2, uint32_t c3, + void *src_shmem) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + uint32_t src = __cvta_generic_to_shared(src_shmem); + asm volatile( + "cp.async.bulk.tensor.4d.global.shared::cta.bulk_group" + " [%0, {%1, %2, %3, %4}], [%5];" ::"l"(tensor_map), + "r"(c0), "r"(c1), "r"(c2), "r"(c3), "r"(src) + : "memory"); +#else + NVTE_DEVICE_ERROR("cp_async_bulk_tensor_4d_shared_to_global requires SM 10.0+."); +#endif +} + +static void create_4D_tensor_map(CUtensorMap &tensorMap, void *dataPtr, DType dtype, uint64_t dim0, + uint64_t dim1, uint64_t dim2, uint64_t dim3, uint32_t box0, + uint32_t box1, uint32_t box2, uint32_t box3) { + cuda_driver::ensure_context_exists(); + static PFN_cuTensorMapEncodeTiled_v12000 cuDriverTensorMapEncodeTiled = []() { + void *ptr = cuda_driver::get_symbol("cuTensorMapEncodeTiled"); + return reinterpret_cast(ptr); + }(); + + CUtensorMapDataType tma_dtype; + size_t elem_bytes; + switch (dtype) { + case DType::kFloat16: + tma_dtype = CU_TENSOR_MAP_DATA_TYPE_FLOAT16; + elem_bytes = 2; + break; + case DType::kBFloat16: + tma_dtype = CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; + elem_bytes = 2; + break; + case DType::kFloat8E4M3: + case DType::kFloat8E5M2: + case DType::kFloat8E8M0: + case DType::kByte: + tma_dtype = CU_TENSOR_MAP_DATA_TYPE_UINT8; + elem_bytes = 1; + break; + default: + NVTE_ERROR("create_4D_tensor_map: unsupported dtype ", to_string(static_cast(dtype))); + } + + constexpr uint32_t rank = 4; + uint64_t size[rank] = {dim0, dim1, dim2, dim3}; + uint64_t stride[rank - 1] = { + dim0 * elem_bytes, + dim0 * dim1 * elem_bytes, + dim0 * dim1 * dim2 * elem_bytes, + }; + uint32_t boxSize[rank] = {box0, box1, box2, box3}; + uint32_t elemStride[rank] = {1, 1, 1, 1}; + + const auto oob_fill = (tma_dtype == CU_TENSOR_MAP_DATA_TYPE_UINT8) + ? CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE + : CU_TENSOR_MAP_FLOAT_OOB_FILL_NAN_REQUEST_ZERO_FMA; + + NVTE_CHECK_CUDA_DRIVER(cuDriverTensorMapEncodeTiled( + &tensorMap, tma_dtype, rank, dataPtr, size, stride, boxSize, elemStride, + CU_TENSOR_MAP_INTERLEAVE_NONE, CU_TENSOR_MAP_SWIZZLE_NONE, CU_TENSOR_MAP_L2_PROMOTION_NONE, + oob_fill)); +} + +template +__device__ __forceinline__ void issue_tma_load_strided(T *smem_buf, const CUtensorMap *tma, + size_t h_i, size_t s_tile, size_t b_i, + uint64_t *mbar, size_t tile_bytes) { + ptx::mbarrier_arrive_expect_tx(mbar, static_cast(tile_bytes)); + if constexpr (kIsBshd) { + cp_async_bulk_tensor_4d_global_to_shared(smem_buf, tma, 0, static_cast(h_i), + static_cast(s_tile), + static_cast(b_i), mbar); + } else { + cp_async_bulk_tensor_4d_global_to_shared(smem_buf, tma, 0, static_cast(h_i), + static_cast(b_i), + static_cast(s_tile), mbar); + } +} + +__device__ __forceinline__ void st_global_cs_uint4(uint4 *ptr, uint4 val) { + asm volatile("st.global.cs.v4.b32 [%0], {%1, %2, %3, %4};" ::"l"(ptr), "r"(val.x), "r"(val.y), + "r"(val.z), "r"(val.w) + : "memory"); +} +// TMA loads from strided input to smem + non-temporal stores to contiguous output in gmem + +template +__launch_bounds__(tma_permute_threads) __global__ + void transpose_to_bhsd_kernel(const __grid_constant__ TmaMapParams tma_maps, + PermuteParams params, size_t b, size_t h_grid, + unsigned int permute_s_splits, size_t s_tile_size) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + const auto &slot = params.slots[blockIdx.z]; + const CUtensorMap *tma_in = &tma_maps.maps[blockIdx.z]; + T *__restrict__ tensor_out = reinterpret_cast(slot.output); + const size_t Sdim = slot.S; + const size_t Hdim = slot.H; + const size_t Ddim = slot.D_in; + const size_t Ddim_out = slot.D_out; + + const size_t b_i = static_cast(blockIdx.x) / h_grid; + const size_t h_i = static_cast(blockIdx.x) % h_grid; + + if (b_i >= b) return; + if (h_i >= Hdim) return; + + const unsigned int s_part = blockIdx.y; + const size_t s_begin = + (Sdim * static_cast(s_part)) / static_cast(permute_s_splits); + const size_t s_end = + (Sdim * static_cast(s_part + 1)) / static_cast(permute_s_splits); + if (s_begin >= s_end) return; + + const size_t out_base = b_i * Hdim * Sdim * Ddim_out + h_i * Sdim * Ddim_out; + + extern __shared__ __align__(128) char smem_raw[]; + T *smem = reinterpret_cast(smem_raw); + + __shared__ __align__(8) uint64_t mbar; + const bool is_leader = (threadIdx.x == 0); + + if (is_leader) { + ptx::mbarrier_init(&mbar, static_cast(blockDim.x)); + ptx::fence_proxy_async_shared_cta(); + } + __syncthreads(); + + const size_t S_TILE = s_tile_size; + const uint32_t tile_bytes = static_cast(S_TILE * Ddim * sizeof(T)); + int parity = 0; + + for (size_t s_tile = s_begin; s_tile < s_end; s_tile += S_TILE) { + const size_t tile_rows = min(S_TILE, s_end - s_tile); + + if (is_leader) { + issue_tma_load_strided(smem, tma_in, h_i, s_tile, b_i, &mbar, tile_bytes); + } else { + ptx::mbarrier_arrive(&mbar); + } + + ptx::mbarrier_wait_parity(&mbar, parity); + parity ^= 1; + + T *__restrict__ out_ptr = tensor_out + out_base + s_tile * Ddim_out; + constexpr size_t vec_elems = sizeof(uint4) / sizeof(T); + + if (Ddim_out == Ddim) { + const size_t total_elems = tile_rows * Ddim; + for (size_t i = threadIdx.x * vec_elems; i < total_elems; + i += static_cast(blockDim.x) * vec_elems) { + uint4 v = *reinterpret_cast(smem + i); + st_global_cs_uint4(reinterpret_cast(out_ptr + i), v); + } + } else { + const size_t total_out_elems = tile_rows * Ddim_out; + for (size_t i = threadIdx.x * vec_elems; i < total_out_elems; + i += static_cast(blockDim.x) * vec_elems) { + const size_t row = i / Ddim_out; + const size_t col = i % Ddim_out; + uint4 v; + if (col + vec_elems <= Ddim) { + v = *reinterpret_cast(smem + row * Ddim + col); + } else { + memset(&v, 0, sizeof(v)); + const size_t smem_off = row * Ddim + col; + size_t copy_elems = (col < Ddim) ? (Ddim - col) : 0; + if (copy_elems > 0) memcpy(&v, smem + smem_off, copy_elems * sizeof(T)); + } + st_global_cs_uint4(reinterpret_cast(out_ptr + i), v); + } + } + + __syncthreads(); + } + + if (is_leader) { + ptx::mbarrier_invalid(&mbar); + } +#endif +} + +// 4D TMA descriptor: +// [B, S, H, D]: TMA dims [D, H, S, B], box [D, 1, S_TILE, 1] +// [S, B, H, D]: TMA dims [D, H, B, S], box [D, 1, 1, S_TILE] + +static void create_strided_tensor_map(CUtensorMap &map, void *ptr, DType dtype, size_t b, size_t s, + size_t h, size_t d, size_t s_tile, bool is_bshd) { + if (is_bshd) { + create_4D_tensor_map(map, ptr, dtype, static_cast(d), static_cast(h), + static_cast(s), static_cast(b), + static_cast(d), 1, static_cast(s_tile), 1); + } else { + create_4D_tensor_map(map, ptr, dtype, static_cast(d), static_cast(h), + static_cast(b), static_cast(s), + static_cast(d), 1, 1, static_cast(s_tile)); + } +} + +void multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_tensors, + NVTE_QKV_Format original_format, cudaStream_t stream) { + using namespace transformer_engine; + if (num_tensors == 0) return; + NVTE_CHECK(num_tensors <= static_cast(kMaxPermuteTensors), "num_tensors must be in [1, ", + kMaxPermuteTensors, "], got ", num_tensors, "."); + + const bool is_bshd = (original_format == NVTE_QKV_Format::NVTE_BSHD); + const DType dtype = inputs[0].dtype(); + const size_t elem_size = typeToSize(dtype); + const size_t b = outputs[0].shape()[0]; + + PermuteParams params{}; + size_t s_max = 0, h_max = 0, s_min = SIZE_MAX; + size_t d_in_max = 0, d_out_max = 0; + bool any_not_vec_aligned = false; + bool all_tma_ok = true; + + for (size_t i = 0; i < num_tensors; ++i) { + const size_t H = outputs[i].shape()[1]; + const size_t S = outputs[i].shape()[2]; + const size_t D_in = inputs[i].shape()[inputs[i].shape().size() - 1]; + const size_t D_out = outputs[i].shape()[3]; + params.slots[i] = {inputs[i].data.dptr, outputs[i].data.dptr, S, H, D_in, D_out}; + s_max = std::max(s_max, S); + h_max = std::max(h_max, H); + s_min = std::min(s_min, S); + d_in_max = std::max(d_in_max, D_in); + d_out_max = std::max(d_out_max, D_out); + if ((D_in * elem_size) % 4 != 0) any_not_vec_aligned = true; + const size_t inner = D_in * elem_size; + if (inner < 32 || inner % 16 != 0) all_tma_ok = false; + } + + if (all_tma_ok) { + const int sm = cuda::sm_arch(cuda::current_device()); + if (sm < 100) { + all_tma_ok = false; + } else { + switch (dtype) { + case DType::kFloat16: + case DType::kBFloat16: + case DType::kFloat8E4M3: + case DType::kFloat8E5M2: + case DType::kFloat8E8M0: + case DType::kByte: + break; + default: + all_tma_ok = false; + } + } + } + + // Dispatch order: + // 1. TMA path: SM 10.0+, D_in*elem >= 32 && 16-aligned, supported dtype, + // and s_tile*D_in*elem is uint4-aligned. + // 2. Fallback path (vec-aligned): vectorized loads/stores when D_in*elem % 4 == 0. + // 3. Fallback path (not-vec-aligned): shared-memory transpose when D_in*elem % 4 != 0. + if (all_tma_ok) { + const size_t s_tile = std::min(static_cast(tma_permute_s_tile_default), s_min); + bool tma_aligned = true; + for (size_t i = 0; i < num_tensors && tma_aligned; ++i) { + if ((s_tile * params.slots[i].D_in * elem_size) % sizeof(uint4) != 0) tma_aligned = false; + } + + if (tma_aligned) { + TmaMapParams tma_maps{}; + for (size_t i = 0; i < num_tensors; ++i) { + const auto &slot = params.slots[i]; + create_strided_tensor_map(tma_maps.maps[i], const_cast(slot.input), dtype, b, + slot.S, slot.H, slot.D_in, s_tile, is_bshd); + } + + const unsigned int permute_s_splits = std::max(1u, static_cast(s_min / s_tile)); + dim3 grid(static_cast(b * h_max), permute_s_splits, + static_cast(num_tensors)); + const size_t smem_bytes = s_tile * d_in_max * elem_size; + + if (is_bshd) { + TRANSFORMER_ENGINE_TYPE_SWITCH_ALL( + dtype, dtype_t, auto kernel = transpose_to_bhsd_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes)); + kernel<<>>(tma_maps, params, b, h_max, + permute_s_splits, s_tile);); + } else { + TRANSFORMER_ENGINE_TYPE_SWITCH_ALL( + dtype, dtype_t, auto kernel = transpose_to_bhsd_kernel; + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes)); + kernel<<>>(tma_maps, params, b, h_max, + permute_s_splits, s_tile);); + } + NVTE_CHECK_CUDA(cudaGetLastError()); + return; + } + } + + if (!any_not_vec_aligned) { + const unsigned int permute_s_splits = std::max( + 1u, static_cast(s_min / static_cast(fallback_permute_threads))); + dim3 grid(static_cast(b * h_max), permute_s_splits, + static_cast(num_tensors)); + + if (is_bshd) { + TRANSFORMER_ENGINE_TYPE_SWITCH_ALL( + dtype, dtype_t, + transpose_to_bhsd_fallback_vec_aligned_kernel + <<>>(params, b, permute_s_splits, h_max);); + } else { + TRANSFORMER_ENGINE_TYPE_SWITCH_ALL( + dtype, dtype_t, + transpose_to_bhsd_fallback_vec_aligned_kernel + <<>>(params, b, permute_s_splits, h_max);); + } + } else { + const unsigned int st = + static_cast((s_max + TRANSPOSE_TILE - 1) / TRANSPOSE_TILE); + const unsigned int ht = + static_cast((h_max + TRANSPOSE_TILE - 1) / TRANSPOSE_TILE); + dim3 grid(static_cast(b) * st, ht, static_cast(num_tensors)); + const size_t D_pad = (d_in_max * elem_size + 3u) & ~size_t(3); + const size_t smem_bytes = + static_cast(TRANSPOSE_TILE) * (static_cast(TRANSPOSE_TILE) * D_pad + 4); + + if (is_bshd) { + TRANSFORMER_ENGINE_TYPE_SWITCH_ALL( + dtype, dtype_t, + transpose_to_bhsd_fallback_not_vec_aligned_kernel + <<>>(params, b, st);); + } else { + TRANSFORMER_ENGINE_TYPE_SWITCH_ALL( + dtype, dtype_t, + transpose_to_bhsd_fallback_not_vec_aligned_kernel + <<>>(params, b, st);); + } + } + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +} // namespace multi_tensor_transpose_to_bhsd + +// =================================================================================== +// multi_tensor_pad_last_dim: pad the last dim of multiple tensors to certain alignment +// =================================================================================== + +namespace multi_tensor_pad_last_dim { + +constexpr int pad_threads_per_block = 256; +constexpr int kMaxPadTensors = 16; + +struct PadLastDimArgs { + const uint8_t *input; + uint32_t *output; + size_t n_uint32; + uint32_t in_row_bytes; + uint32_t out_row_uint32; +}; + +struct MultiPadParams { + PadLastDimArgs tensors[kMaxPadTensors]; +}; + +__launch_bounds__(pad_threads_per_block) __global__ + void multi_tensor_pad_last_dim_kernel(MultiPadParams params) { + const auto &a = params.tensors[blockIdx.y]; + + for (size_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; idx < a.n_uint32; + idx += static_cast(gridDim.x) * blockDim.x) { + const uint32_t col_byte = (idx % a.out_row_uint32) * 4; + const size_t row = idx / a.out_row_uint32; + const uint8_t *__restrict__ src = a.input + row * static_cast(a.in_row_bytes); + + uint32_t val; + if (col_byte + 4 <= a.in_row_bytes) { + memcpy(&val, src + col_byte, 4); + } else if (col_byte >= a.in_row_bytes) { + val = 0; + } else { + val = 0; + memcpy(&val, src + col_byte, a.in_row_bytes - col_byte); + } + a.output[idx] = val; + } +} + +void launch_pad_batch(MultiPadParams ¶ms, int kernel_count, size_t max_n_uint32, + cudaStream_t stream) { + if (kernel_count == 0) return; + constexpr int threads = pad_threads_per_block; + const int blocks_x = static_cast( + std::min(DIVUP(max_n_uint32, static_cast(threads)), static_cast(65535))); + dim3 grid(blocks_x, kernel_count); + multi_tensor_pad_last_dim_kernel<<>>(params); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +void multi_tensor_pad_last_dim(Tensor *inputs, Tensor *outputs, size_t num_tensors, + cudaStream_t stream) { + using namespace transformer_engine; + + if (num_tensors == 0) return; + + MultiPadParams params{}; + size_t max_n_uint32 = 0; + int kernel_count = 0; + + for (size_t i = 0; i < num_tensors; ++i) { + auto &inp = inputs[i]; + auto &out = outputs[i]; + + NVTE_CHECK(inp.data.shape.size() == 2, "Expected 2D input tensor at index ", i, "."); + NVTE_CHECK(out.data.shape.size() == 2, "Expected 2D output tensor at index ", i, "."); + NVTE_CHECK(inp.data.dtype == out.data.dtype, "Dtype mismatch at index ", i, "."); + + const size_t rows = inp.data.shape[0]; + const size_t in_cols = inp.data.shape[1]; + const size_t out_cols = out.data.shape[1]; + + NVTE_CHECK(out.data.shape[0] == rows, "Row count mismatch at index ", i, "."); + NVTE_CHECK(out_cols >= in_cols, "out_cols < in_cols at index ", i, "."); + + if (rows == 0) continue; + + if (in_cols == out_cols) { + const size_t total_bytes = rows * in_cols * typeToSize(inp.data.dtype); + NVTE_CHECK_CUDA(cudaMemcpyAsync(out.data.dptr, inp.data.dptr, total_bytes, + cudaMemcpyDeviceToDevice, stream)); + continue; + } + + if (kernel_count == kMaxPadTensors) { + launch_pad_batch(params, kernel_count, max_n_uint32, stream); + params = MultiPadParams{}; + kernel_count = 0; + max_n_uint32 = 0; + } + + const size_t elem_size = typeToSize(inp.data.dtype); + const auto in_row_bytes = static_cast(in_cols * elem_size); + const auto out_row_bytes = static_cast(out_cols * elem_size); + NVTE_CHECK(out_row_bytes % 4 == 0, "Padded row size in bytes (", out_row_bytes, + ") must be a multiple of 4."); + + const uint32_t out_row_uint32 = out_row_bytes / 4; + const size_t n_uint32 = rows * out_row_uint32; + + params.tensors[kernel_count] = {reinterpret_cast(inp.data.dptr), + reinterpret_cast(out.data.dptr), n_uint32, + in_row_bytes, out_row_uint32}; + max_n_uint32 = std::max(max_n_uint32, n_uint32); + ++kernel_count; + } + + launch_pad_batch(params, kernel_count, max_n_uint32, stream); +} + +} // namespace multi_tensor_pad_last_dim } // namespace transformer_engine void nvte_prepare_flash_attn_fwd(NVTETensor qkvi, NVTETensor qkv, cudaStream_t stream) { @@ -153,3 +861,40 @@ void nvte_prepare_flash_attn_bwd(NVTETensor q, NVTETensor k, NVTETensor v, NVTET *convertNVTETensorCheck(v), *convertNVTETensorCheck(qkv), stream); } + +void nvte_multi_tensor_transpose_to_bhsd(NVTETensor *inputs, NVTETensor *outputs, + size_t num_tensors, NVTE_QKV_Format original_format, + cudaStream_t stream) { + NVTE_API_CALL(nvte_multi_tensor_transpose_to_bhsd); + NVTE_CHECK(original_format == NVTE_QKV_Format::NVTE_BSHD || + original_format == NVTE_QKV_Format::NVTE_SBHD, + "nvte_multi_tensor_transpose_to_bhsd: only BSHD/SBHD -> BHSD is currently " + "supported."); + using namespace transformer_engine; + + std::vector in_vec(num_tensors), out_vec(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + in_vec[i] = *convertNVTETensorCheck(inputs[i]); + out_vec[i] = *convertNVTETensorCheck(outputs[i]); + } + constexpr size_t kBatch = multi_tensor_transpose_to_bhsd::kMaxPermuteTensors; + for (size_t offset = 0; offset < num_tensors; offset += kBatch) { + const size_t batch = std::min(num_tensors - offset, kBatch); + multi_tensor_transpose_to_bhsd::multi_tensor_transpose_to_bhsd( + in_vec.data() + offset, out_vec.data() + offset, batch, original_format, stream); + } +} + +void nvte_multi_tensor_pad_last_dim(NVTETensor *inputs, NVTETensor *outputs, size_t num_tensors, + cudaStream_t stream) { + NVTE_API_CALL(nvte_multi_tensor_pad_last_dim); + using namespace transformer_engine; + + std::vector in_vec(num_tensors), out_vec(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + in_vec[i] = *convertNVTETensorCheck(inputs[i]); + out_vec[i] = *convertNVTETensorCheck(outputs[i]); + } + multi_tensor_pad_last_dim::multi_tensor_pad_last_dim(in_vec.data(), out_vec.data(), num_tensors, + stream); +} diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 3d6e3a0aa..141767b80 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -131,6 +131,8 @@ NVTE_QKV_Layout_Group nvte_get_qkv_layout_group(NVTE_QKV_Layout qkv_layout) { case NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_SBHD_SBHD: case NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD: return NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD; + case NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD: + return NVTE_QKV_Layout_Group::NVTE_SD_SD_SD; default: NVTE_ERROR("Unsupported qkv_layout ", transformer_engine::to_string(qkv_layout), " in nvte_get_qkv_layout_group."); @@ -172,6 +174,8 @@ NVTE_QKV_Format nvte_get_qkv_format(NVTE_QKV_Layout qkv_layout) { case NVTE_QKV_Layout::NVTE_THD_SBHD_SBHD: case NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD: return NVTE_QKV_Format::NVTE_THD_2SBHD; + case NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD: + return NVTE_QKV_Format::NVTE_BHSD; default: NVTE_ERROR("Unsupported qkv_layout ", transformer_engine::to_string(qkv_layout), " in nvte_get_qkv_format."); @@ -192,6 +196,8 @@ NVTE_QKV_Format nvte_get_q_format(NVTE_QKV_Layout qkv_layout) { case NVTE_QKV_Format::NVTE_THD_2BSHD: case NVTE_QKV_Format::NVTE_THD_2SBHD: return NVTE_QKV_Format::NVTE_THD; + case NVTE_QKV_Format::NVTE_BHSD: + return NVTE_QKV_Format::NVTE_BHSD; default: NVTE_ERROR("Unsupported qkv_format ", transformer_engine::to_string(qkv_format), " in nvte_get_q_format."); @@ -212,6 +218,8 @@ NVTE_QKV_Format nvte_get_kv_format(NVTE_QKV_Layout qkv_layout) { return NVTE_QKV_Format::NVTE_BSHD; case NVTE_QKV_Format::NVTE_THD: return NVTE_QKV_Format::NVTE_THD; + case NVTE_QKV_Format::NVTE_BHSD: + return NVTE_QKV_Format::NVTE_BHSD; default: NVTE_ERROR("Unsupported qkv_format ", transformer_engine::to_string(qkv_format), " in nvte_get_kv_format."); @@ -269,9 +277,22 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK))) && - (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD) && - !requires_64bit_ragged_offset && (softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) && + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)) || + // 9.21: d_qk=192, d_v=128 + (cudnn_runtime_version >= 92100 && sm_arch_ >= 100 && head_dim_qk <= 192 && + head_dim_v <= 128 && head_dim_qk % 16 == 0 && head_dim_v % 16 == 0 && + (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK))) && + // pre-9.21: {bshd, sbhd}, {vanilla} + // 9.21+: {bshd, sbhd, bhsd}, {vanilla, off-by-one, learnable} + ((cudnn_runtime_version < 92100 && + (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD) && + softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) || + (cudnn_runtime_version >= 92100 && + (qkv_format == NVTE_QKV_Format::NVTE_BSHD || qkv_format == NVTE_QKV_Format::NVTE_SBHD || + qkv_format == NVTE_QKV_Format::NVTE_BHSD))) && + !requires_64bit_ragged_offset && // 9.10.0: known bugs with SDPA FP8 (cudnn_runtime_version != 91000) && !return_max_logit) { if (cudnn_runtime_version >= 8900) { @@ -410,12 +431,15 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS)) && // qkv format (qkv_format == NVTE_QKV_Format::NVTE_SBHD || qkv_format == NVTE_QKV_Format::NVTE_BSHD || + qkv_format == NVTE_QKV_Format::NVTE_BHSD || (qkv_format == NVTE_QKV_Format::NVTE_THD && sm_arch_ >= 90 && ((cudnn_runtime_version >= 90100 && num_attn_heads == num_gqa_groups) || cudnn_runtime_version >= 90600)) || ((q_format == NVTE_QKV_Format::NVTE_SBHD || q_format == NVTE_QKV_Format::NVTE_BSHD || + q_format == NVTE_QKV_Format::NVTE_BHSD || (q_format == NVTE_QKV_Format::NVTE_THD && sm_arch_ >= 90) || kv_format == NVTE_QKV_Format::NVTE_SBHD || kv_format == NVTE_QKV_Format::NVTE_BSHD || + kv_format == NVTE_QKV_Format::NVTE_BHSD || (kv_format == NVTE_QKV_Format::NVTE_THD && sm_arch_ >= 90)) && cudnn_runtime_version >= 90700)) && // sliding window @@ -565,7 +589,8 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso const NVTETensor page_table_v, const NVTETensor rng_state, size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, bool return_max_logit, bool cuda_graph, float attn_scale, float dropout, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream) { @@ -587,23 +612,24 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso Tensor *output_O = convertNVTETensorCheck(O); Tensor *wkspace = convertNVTETensor(workspace); - auto ndim = input_Q->data.shape.size(); - auto ndim_kv = input_K->data.shape.size(); - size_t b = input_cu_seqlens_q->data.shape[0] - 1; - size_t h_q = input_Q->data.shape[ndim - 2]; - size_t h_kv = input_K->data.shape[ndim_kv - 2]; - size_t d_qk = input_Q->data.shape[ndim - 1]; - size_t d_v = input_V->data.shape[ndim_kv - 1]; - size_t t_q = 0; - size_t t_kv = 0; NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + auto *q_dims = input_Q->data.shape.data(); + auto *k_dims = input_K->data.shape.data(); + auto *v_dims = input_V->scaling_mode != NVTE_MXFP8_1D_SCALING + ? input_V->data.shape.data() + : input_V->columnwise_data.shape.data(); + AttentionShape q_shape(q_format, q_dims); + AttentionShape k_shape(kv_format, k_dims); + AttentionShape v_shape(kv_format, v_dims); + size_t b = q_shape.b(), h_q = q_shape.h(), d_qk = q_shape.d(), t_q = q_shape.t(); + size_t h_kv = k_shape.h(), t_kv = k_shape.t(), d_v = v_shape.d(); if (q_format == NVTE_QKV_Format::NVTE_THD) { - t_q = input_Q->data.shape[0]; - } - if (kv_format == NVTE_QKV_Format::NVTE_THD) { - t_kv = input_K->data.shape[0]; + b = input_cu_seqlens_q->data.shape[0] - 1; + } else if (kv_format == NVTE_QKV_Format::NVTE_THD) { + b = input_cu_seqlens_kv->data.shape[0] - 1; } + int64_t num_pages_k = 0; int64_t num_pages_v = 0; int64_t page_size_k = 0; @@ -642,38 +668,26 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso return_max_logit, cuda_graph, false); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { -#if (CUDNN_VERSION >= 8901) fused_attn_max_512_fwd(b, h_q, max_seqlen_q, max_seqlen_kv, d_qk, is_training, attn_scale, dropout, qkv_layout, bias_type, attn_mask_type, input_Q, input_K, input_V, input_Bias, output_O, Aux_CTX_Tensors, input_cu_seqlens_q, input_cu_seqlens_kv, input_rng_state, wkspace, stream, handle); -#else - NVTE_ERROR("cuDNN 8.9.1 is required for BF16/FP16 fused attention with max_seqlen<=512. \n"); -#endif } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { -#if (CUDNN_VERSION >= 8900) fused_attn_arbitrary_seqlen_fwd( b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, num_pages_k, num_pages_v, page_size_k, page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, is_training, - return_max_logit, attn_scale, dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, input_Q, input_K, input_V, - input_Bias, input_SoftmaxOffset, output_O, Aux_CTX_Tensors, input_cu_seqlens_q, + return_max_logit, attn_scale, dropout, qkv_layout, o_format, bias_type, attn_mask_type, + softmax_type, window_size_left, window_size_right, bottom_right_diagonal, input_Q, input_K, + input_V, input_Bias, input_SoftmaxOffset, output_O, Aux_CTX_Tensors, input_cu_seqlens_q, input_cu_seqlens_kv, input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_page_table_k, input_page_table_v, input_rng_state, wkspace, stream, handle); -#else - NVTE_ERROR( - "cuDNN 8.9.0 is required for BF16/FP16 fused attention with arbitrary sequence length. " - "\n"); -#endif } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { -#if (CUDNN_VERSION >= 8900) - fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, is_training, attn_scale, - dropout, qkv_layout, bias_type, attn_mask_type, input_Q, input_K, input_V, + fused_attn_fp8_fwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, is_training, + attn_scale, dropout, qkv_layout, o_format, qkv_scale_inv_format, bias_type, + attn_mask_type, softmax_type, window_size_left, window_size_right, + bottom_right_diagonal, input_Q, input_K, input_V, input_SoftmaxOffset, input_output_S, output_O, Aux_CTX_Tensors, input_cu_seqlens_q, input_cu_seqlens_kv, input_rng_state, wkspace, stream, handle); -#else - NVTE_ERROR("cuDNN 8.9.0 is required for FP8 fused attention. \n"); -#endif } else { NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); } @@ -687,11 +701,13 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float dropout, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, bool cuda_graph, - NVTETensor workspace, cudaStream_t stream) { + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, + NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, + NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, + bool cuda_graph, NVTETensor workspace, cudaStream_t stream) { NVTE_API_CALL(nvte_flash_attn_bwd); using namespace transformer_engine; const Tensor *input_cu_seqlens_q = convertNVTETensorCheck(cu_seqlens_q); @@ -712,22 +728,20 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso Tensor *output_dSoftmaxOffset = convertNVTETensorCheck(dSoftmaxOffset); Tensor *wkspace = convertNVTETensor(workspace); - auto ndim = input_Q->data.shape.size(); - auto ndim_kv = input_K->data.shape.size(); - size_t b = input_cu_seqlens_q->data.shape[0] - 1; - size_t h_q = input_Q->data.shape[ndim - 2]; - size_t h_kv = input_K->data.shape[ndim_kv - 2]; - size_t d_qk = input_Q->data.shape[ndim - 1]; - size_t d_v = input_V->data.shape[ndim_kv - 1]; - size_t t_q = 0; - size_t t_kv = 0; NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + auto *q_dims = input_Q->data.shape.data(); + auto *k_dims = input_K->data.shape.data(); + auto *v_dims = input_V->data.shape.data(); + AttentionShape q_shape(q_format, q_dims); + AttentionShape k_shape(kv_format, k_dims); + AttentionShape v_shape(kv_format, v_dims); + size_t b = q_shape.b(), h_q = q_shape.h(), d_qk = q_shape.d(), t_q = q_shape.t(); + size_t h_kv = k_shape.h(), t_kv = k_shape.t(), d_v = v_shape.d(); if (q_format == NVTE_QKV_Format::NVTE_THD) { - t_q = input_Q->data.shape[0]; - } - if (kv_format == NVTE_QKV_Format::NVTE_THD) { - t_kv = input_K->data.shape[0]; + b = input_cu_seqlens_q->data.shape[0] - 1; + } else if (kv_format == NVTE_QKV_Format::NVTE_THD) { + b = input_cu_seqlens_kv->data.shape[0] - 1; } auto handle = cudnnExecutionPlanManager::Instance().GetHandle(); @@ -740,17 +754,12 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso cuda_graph, deterministic); if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { -#if (CUDNN_VERSION >= 8901) Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[0]); fused_attn_max_512_bwd(b, h_q, max_seqlen_q, max_seqlen_kv, d_qk, attn_scale, dropout, qkv_layout, bias_type, attn_mask_type, input_Q, input_K, input_V, input_dO, output_S, output_dQ, output_dK, output_dV, output_dBias, input_cu_seqlens_q, input_cu_seqlens_kv, wkspace, stream, handle); -#else - NVTE_ERROR("cuDNN 8.9.1 is required for BF16/FP16 fused attention with max_seqlen<=512. \n"); -#endif } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { -#if (CUDNN_VERSION >= 8900) size_t i = 0; Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); Tensor *input_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); @@ -763,30 +772,36 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso } fused_attn_arbitrary_seqlen_bwd( b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, attn_scale, dropout, - qkv_layout, bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, - bottom_right_diagonal, deterministic, input_Q, input_K, input_V, input_O, input_dO, - input_Bias, input_SoftmaxOffset, output_S, output_dQ, output_dK, output_dV, output_dBias, - output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, input_cu_seqlens_q_padded, - input_cu_seqlens_kv_padded, input_rng_state, wkspace, stream, handle); -#else - const char *err_msg = - "cuDNN 8.9.0 is required for BF16/FP16 fused attention " - "with arbitrary sequence length. \n"; - NVTE_ERROR(err_msg); -#endif + qkv_layout, o_format, do_format, dqkv_layout, bias_type, attn_mask_type, softmax_type, + window_size_left, window_size_right, bottom_right_diagonal, deterministic, input_Q, input_K, + input_V, input_O, input_dO, input_Bias, input_SoftmaxOffset, output_S, output_dQ, output_dK, + output_dV, output_dBias, output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, + input_cu_seqlens_q_padded, input_cu_seqlens_kv_padded, input_rng_state, wkspace, stream, + handle); } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { -#if (CUDNN_VERSION >= 8900) - const Tensor *input_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[0]); - const Tensor *input_ZInv = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[1]); - const Tensor *input_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[2]); - fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, attn_scale, dropout, - qkv_layout, bias_type, attn_mask_type, deterministic, input_Q, input_K, - input_V, input_O, input_dO, input_M, input_ZInv, input_S, input_output_dP, - output_dQ, output_dK, output_dV, input_cu_seqlens_q, input_cu_seqlens_kv, - input_rng_state, wkspace, stream, handle); -#else - NVTE_ERROR("cuDNN 8.9.0 is required for FP8 fused attention. \n"); -#endif + size_t i = 0; + const Tensor *input_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + const Tensor *input_ZInv = nullptr; + if (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD) { + input_ZInv = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + } + const Tensor *input_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + const Tensor *input_SoftmaxOffset = nullptr; + if (softmax_type != NVTE_VANILLA_SOFTMAX) { + input_SoftmaxOffset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + } + const Tensor *input_dO_f16 = nullptr; + if (input_dO->scaling_mode == NVTE_MXFP8_1D_SCALING) { + input_dO_f16 = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + } + fused_attn_fp8_bwd(b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, attn_scale, dropout, + qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, + do_scale_inv_format, bias_type, attn_mask_type, softmax_type, + window_size_left, window_size_right, bottom_right_diagonal, deterministic, + input_Q, input_K, input_V, input_O, input_dO, input_dO_f16, input_M, + input_ZInv, input_S, input_SoftmaxOffset, input_output_dP, output_dQ, + output_dK, output_dV, output_dSoftmaxOffset, input_cu_seqlens_q, + input_cu_seqlens_kv, input_rng_state, wkspace, stream, handle); } else { NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); } diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu index eed674074..6df7ad35c 100644 --- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu +++ b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu @@ -19,7 +19,6 @@ #include "fused_attn_f16_arbitrary_seqlen.h" #include "utils.h" -#if (CUDNN_VERSION >= 8900) #define Q_ID 1 #define K_ID 2 #define V_ID 3 @@ -54,11 +53,11 @@ void fused_attn_arbitrary_seqlen_fwd_impl( int64_t page_size_k, int64_t page_size_v, int64_t max_pages_per_seq_k, int64_t max_pages_per_seq_v, int64_t bias_b, int64_t bias_h, int64_t bias_sq, int64_t bias_skv, bool is_training, bool return_max_logit, float scaling_factor, float dropout_probability, - NVTE_QKV_Layout layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, void *devPtrQ, void *devPtrK, void *devPtrV, void *devPtrBias, - void *devPtrSoftmaxOffset, void *devPtrS1, void *devPtrS2, void *devPtrO, - void *devPtrDropoutSeed, void *devPtrDropoutOffset, void *devPtrCuSeqlensQ, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_Bias_Type bias_type, + NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, void *devPtrQ, void *devPtrK, + void *devPtrV, void *devPtrBias, void *devPtrSoftmaxOffset, void *devPtrS1, void *devPtrS2, + void *devPtrO, void *devPtrDropoutSeed, void *devPtrDropoutOffset, void *devPtrCuSeqlensQ, void *devPtrCuSeqlensKV, void *devPtrPageTableK, void *devPtrPageTableV, void *devPtrSeqOffsetsQ, void *devPtrSeqOffsetsKV, cudnn_frontend::DataType_t tensorType, void *workspace, size_t *workspace_size, cudaStream_t stream, cudnnHandle_t handle) { @@ -80,8 +79,8 @@ void fused_attn_arbitrary_seqlen_fwd_impl( } bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); bool is_dropout = (is_training && dropout_probability != 0.0f); - NVTE_QKV_Format q_format = nvte_get_q_format(layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(layout); + NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); bool is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); const auto cudnn_runtime_version = cudnnGetVersion(); @@ -89,7 +88,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( const int sm_arch_ = cuda::sm_arch(device_id); bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(layout); + NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); if (is_paged_kv) { NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); @@ -135,7 +134,12 @@ void fused_attn_arbitrary_seqlen_fwd_impl( scaling_factor, is_training, dropout_probability, - layout, + qkv_layout, + o_format, + NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Layout_NOT_SET, + NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, @@ -202,17 +206,17 @@ void fused_attn_arbitrary_seqlen_fwd_impl( std::vector q_stride(4); std::vector k_stride(4); std::vector v_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d_qk, q_stride.data(), layout, + generateMatrixStrides(b, h, s_q, s_kv, d_qk, q_stride.data(), qkv_layout, NVTE_QKV_Matrix::NVTE_Q_Matrix); if (is_paged_kv) { generateMatrixStrides(num_pages_k, hg, page_size_k, page_size_v, d_qk, k_stride.data(), - layout, NVTE_QKV_Matrix::NVTE_K_Matrix); + qkv_layout, NVTE_QKV_Matrix::NVTE_K_Matrix); generateMatrixStrides(num_pages_v, hg, page_size_k, page_size_v, d_v, v_stride.data(), - layout, NVTE_QKV_Matrix::NVTE_V_Matrix); + qkv_layout, NVTE_QKV_Matrix::NVTE_V_Matrix); } else { - generateMatrixStrides(b, hg, s_q, s_kv, d_qk, k_stride.data(), layout, + generateMatrixStrides(b, hg, s_q, s_kv, d_qk, k_stride.data(), qkv_layout, NVTE_QKV_Matrix::NVTE_K_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), layout, + generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), qkv_layout, NVTE_QKV_Matrix::NVTE_V_Matrix); } @@ -368,7 +372,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( auto [O, Stats] = mha_graph->sdpa(Q, K, V, std::move(sdpa_options)); std::vector o_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), layout, + generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), qkv_layout, NVTE_QKV_Matrix::NVTE_O_Matrix); O->set_output(true).set_dim({b, h, s_q, d_v}).set_stride(o_stride); if (is_ragged_q) { @@ -513,7 +517,7 @@ void fused_attn_arbitrary_seqlen_fwd_impl( (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * num_bytes_per_ragged_offset; } - const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(layout); + const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); cu_seqlens_padded_to_offsets<<>>( layout_group, actual_b, b, h, hg, d_qk, d_v, static_cast(devPtrSeqOffsetsQ), static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, @@ -551,7 +555,8 @@ void fused_attn_arbitrary_seqlen_bwd_impl( int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, int64_t max_b, int64_t max_t_q, int64_t max_t_kv, int64_t bias_b, int64_t bias_h, int64_t bias_sq, int64_t bias_skv, float scaling_factor, float dropout_probability, - NVTE_QKV_Layout layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, + NVTE_QKV_Layout dqkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, void *devPtrQ, void *devPtrKTranspose, void *devPtrVTranspose, void *devPtrO, void *devPtrSoftmaxStats, void *devPtrBias, @@ -578,8 +583,8 @@ void fused_attn_arbitrary_seqlen_bwd_impl( } bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); bool is_dropout = (dropout_probability != 0.0f); - NVTE_QKV_Format q_format = nvte_get_q_format(layout); - NVTE_QKV_Format kv_format = nvte_get_kv_format(layout); + NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); bool is_ragged_q = (q_format == NVTE_QKV_Format::NVTE_THD); bool is_ragged_kv = (kv_format == NVTE_QKV_Format::NVTE_THD); const auto cudnn_runtime_version = cudnnGetVersion(); @@ -587,7 +592,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( const int sm_arch_ = cuda::sm_arch(device_id); bool use_ragged_stats = is_ragged_q && cudnn_runtime_version >= 90600 && sm_arch_ != 120; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(layout); + NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); bool is_paged_kv = (layout_group == NVTE_QKV_Layout_Group::NVTE_Paged_KV_HD_HD_HD); if (is_paged_kv) { NVTE_CHECK(is_padding, "Paged attention requires padding mask!"); @@ -632,7 +637,12 @@ void fused_attn_arbitrary_seqlen_bwd_impl( scaling_factor, true, dropout_probability, - layout, + qkv_layout, + o_format, + do_format, + dqkv_layout, + NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, @@ -703,13 +713,13 @@ void fused_attn_arbitrary_seqlen_bwd_impl( std::vector k_stride(4); std::vector v_stride(4); std::vector o_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d_qk, q_stride.data(), layout, + generateMatrixStrides(b, h, s_q, s_kv, d_qk, q_stride.data(), qkv_layout, NVTE_QKV_Matrix::NVTE_Q_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d_qk, k_stride.data(), layout, + generateMatrixStrides(b, hg, s_q, s_kv, d_qk, k_stride.data(), qkv_layout, NVTE_QKV_Matrix::NVTE_K_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), layout, + generateMatrixStrides(b, hg, s_q, s_kv, d_v, v_stride.data(), qkv_layout, NVTE_QKV_Matrix::NVTE_V_Matrix); - generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), layout, + generateMatrixStrides(b, h, s_q, s_kv, d_v, o_stride.data(), qkv_layout, NVTE_QKV_Matrix::NVTE_O_Matrix); q = mha_graph->tensor(fe::graph::Tensor_attributes() @@ -1024,7 +1034,7 @@ void fused_attn_arbitrary_seqlen_bwd_impl( (static_cast(is_ragged_q) + static_cast(is_ragged_kv)) * 2 * num_bytes_per_ragged_offset; } - const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(layout); + const NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); cu_seqlens_padded_to_offsets<<>>( layout_group, actual_b, b, h, hg, d_qk, d_v, static_cast(devPtrSeqOffsetsQ), static_cast(devPtrSeqOffsetsKV), ragged_offset_type, devOffsetsQ, devOffsetsK, @@ -1067,13 +1077,14 @@ void fused_attn_arbitrary_seqlen_fwd( size_t num_tokens_kv, size_t num_pages_k, size_t num_pages_v, size_t page_size_k, size_t page_size_v, size_t max_pages_per_seq_k, size_t max_pages_per_seq_v, bool is_training, bool return_max_logit, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, - const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, const Tensor *input_Bias, - const Tensor *input_SoftmaxOffset, Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, - const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, - const Tensor *cu_seqlens_kv_padded, const Tensor *page_table_k, const Tensor *page_table_v, - const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle) { + NVTE_QKV_Format o_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, + const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, Tensor *output_O, + NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, + const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, + const Tensor *page_table_k, const Tensor *page_table_v, const Tensor *rng_state, + Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; const auto QKV_type = input_Q->data.dtype; @@ -1202,12 +1213,12 @@ void fused_attn_arbitrary_seqlen_fwd( batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, max_batch_size, max_tokens_q, max_tokens_kv, num_pages_k, num_pages_v, page_size_k, page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, bias_b, bias_h, bias_sq, bias_skv, - is_training, return_max_logit, attn_scale, p_dropout, qkv_layout, bias_type, mask_type, - softmax_type, window_size_left, window_size_right, bottom_right_diagonal, devPtrQ, devPtrK, - devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS1, devPtrS2, devPtrO, devPtrDropoutSeed, - devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrPageTableK, devPtrPageTableV, - devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, get_cudnn_fe_dtype(QKV_type), workspace->data.dptr, - &workspace_size, stream, handle); + is_training, return_max_logit, attn_scale, p_dropout, qkv_layout, o_format, bias_type, + mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, devPtrQ, + devPtrK, devPtrV, devPtrBias, devPtrSoftmaxOffset, devPtrS1, devPtrS2, devPtrO, + devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrPageTableK, + devPtrPageTableV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, get_cudnn_fe_dtype(QKV_type), + workspace->data.dptr, &workspace_size, stream, handle); if (workspace_size > 0) { if (workspace->data.dptr == nullptr) { @@ -1228,6 +1239,7 @@ void fused_attn_arbitrary_seqlen_bwd( size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, size_t num_tokens_kv, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, + NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, @@ -1300,12 +1312,12 @@ void fused_attn_arbitrary_seqlen_bwd( fused_attn_arbitrary_seqlen_bwd_impl( batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, max_batch_size, max_tokens_q, max_tokens_kv, bias_b, bias_h, bias_sq, bias_skv, attn_scale, - p_dropout, qkv_layout, bias_type, mask_type, softmax_type, window_size_left, - window_size_right, bottom_right_diagonal, deterministic, devPtrQ, devPtrK, devPtrV, devPtrO, - devPtrSoftmaxStats, devPtrBias, devPtrSoftmaxOffset, devPtrdQ, devPtrdK, devPtrdV, devPtrdO, - devPtrdBias, devPtrdSoftmaxOffset, devPtrDropoutSeed, devPtrDropoutOffset, devPtrCuSeqlensQ, - devPtrCuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, get_cudnn_fe_dtype(QKV_type), - workspace->data.dptr, &workspace_size, stream, handle); + p_dropout, qkv_layout, o_format, do_format, dqkv_layout, bias_type, mask_type, softmax_type, + window_size_left, window_size_right, bottom_right_diagonal, deterministic, devPtrQ, devPtrK, + devPtrV, devPtrO, devPtrSoftmaxStats, devPtrBias, devPtrSoftmaxOffset, devPtrdQ, devPtrdK, + devPtrdV, devPtrdO, devPtrdBias, devPtrdSoftmaxOffset, devPtrDropoutSeed, devPtrDropoutOffset, + devPtrCuSeqlensQ, devPtrCuSeqlensKV, devPtrSeqOffsetsQ, devPtrSeqOffsetsKV, + get_cudnn_fe_dtype(QKV_type), workspace->data.dptr, &workspace_size, stream, handle); if (workspace_size > 0) { if (workspace->data.dptr == nullptr) { @@ -1322,4 +1334,3 @@ void fused_attn_arbitrary_seqlen_bwd( } } } // namespace transformer_engine -#endif // CUDNN_VERSION >= 8900 diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h index 4dd7f3d1d..8f79b5bb4 100644 --- a/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h +++ b/transformer_engine/common/fused_attn/fused_attn_f16_arbitrary_seqlen.h @@ -17,25 +17,26 @@ #include "transformer_engine/fused_attn.h" namespace transformer_engine { -#if (CUDNN_VERSION >= 8900) void fused_attn_arbitrary_seqlen_fwd( size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, size_t num_tokens_kv, size_t num_pages_k, size_t num_pages_v, size_t page_size_k, size_t page_size_v, size_t max_pages_per_seq_k, size_t max_pages_per_seq_v, bool is_training, bool return_max_logit, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, - const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, const Tensor *input_Bias, - const Tensor *input_SoftmaxOffset, Tensor *output_O, NVTETensorPack *Aux_CTX_Tensors, - const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *cu_seqlens_q_padded, - const Tensor *cu_seqlens_kv_padded, const Tensor *page_table_k, const Tensor *page_table_v, - const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); + NVTE_QKV_Format o_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, + const Tensor *input_Bias, const Tensor *input_SoftmaxOffset, Tensor *output_O, + NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, + const Tensor *cu_seqlens_q_padded, const Tensor *cu_seqlens_kv_padded, + const Tensor *page_table_k, const Tensor *page_table_v, const Tensor *rng_state, + Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); void fused_attn_arbitrary_seqlen_bwd( size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, size_t num_tokens_q, size_t num_tokens_kv, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, + NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, @@ -46,7 +47,6 @@ void fused_attn_arbitrary_seqlen_bwd( const Tensor *cu_seqlens_kv_padded, const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); -#endif // CUDNN_VERSION >= 8900 } // namespace transformer_engine #endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_ARBITRARY_SEQLEN_H_ diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu b/transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu index 336e3d538..d5151a51f 100644 --- a/transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu +++ b/transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu @@ -16,7 +16,6 @@ #include "fused_attn_f16_max512_seqlen.h" #include "utils.h" -#if (CUDNN_VERSION >= 8901) #define Q_ID 1 #define K_ID 2 #define V_ID 3 @@ -1342,4 +1341,3 @@ void fused_attn_max_512_bwd(size_t batch, size_t num_head, size_t q_max_seqlen, } } } // namespace transformer_engine -#endif // CUDNN_VERSION >= 8901 diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.h b/transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.h index 3b30c6e71..1e59d4dc8 100644 --- a/transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.h +++ b/transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.h @@ -17,7 +17,6 @@ #include "transformer_engine/fused_attn.h" namespace transformer_engine { -#if (CUDNN_VERSION >= 8901) void fused_attn_max_512_fwd(size_t batch, size_t num_head, size_t q_max_seqlen, size_t kv_max_seqlen, size_t head_dim, bool is_training, float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, @@ -37,7 +36,6 @@ void fused_attn_max_512_bwd(size_t batch, size_t num_head, size_t q_max_seqlen, Tensor *output_dBias, const Tensor *q_cu_seqlens, const Tensor *kv_cu_seqlens, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); -#endif // CUDNN_VERSION >= 8901 } // namespace transformer_engine #endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_MAX_512_H_ diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index 80e64370f..d97f38845 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -15,7 +15,6 @@ namespace fused_attn { using namespace transformer_engine; -#if (CUDNN_VERSION >= 8900) std::unordered_map tensor_name_to_uid = {{"Q", 1}, {"K", 2}, {"V", 3}, @@ -1652,16 +1651,20 @@ void fused_attn_fp8_bwd_impl( // fused attention FWD FP8 with FE 1.0+ void fused_attn_fp8_fwd_impl_v1( - int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d, bool is_training, - float scaling_factor, float dropout_probability, NVTE_QKV_Layout layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, void* devPtrQ, void* devPtrK, void* devPtrV, - void* devPtrM, void* devPtrZInv, void* devPtrO, void* devPtrDescaleQ, void* devPtrDescaleK, - void* devPtrDescaleV, void* devPtrDescaleS, void* devPtrScaleS, void* devPtrScaleO, - void* devPtrAmaxO, void* devPtrAmaxS, void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, - void* devPtrDropoutSeed, void* devPtrDropoutOffset, cudnn_frontend::DataType_t qkv_tensor_type, - cudnn_frontend::DataType_t o_tensor_type, void* workspace, size_t* workspace_size, - cudaStream_t stream, cudnnHandle_t handle) { + int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, + bool is_training, float scaling_factor, float dropout_probability, NVTE_QKV_Layout qkv_layout, + NVTE_QKV_Format o_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, + bool bottom_right_diagonal, void* devPtrQ, void* devPtrK, void* devPtrV, + void* devPtrSoftmaxOffset, void* devPtrM, void* devPtrZInv, void* devPtrO, void* devPtrDescaleQ, + void* devPtrDescaleK, void* devPtrDescaleV, void* devPtrDescaleS, void* devPtrScaleS, + void* devPtrScaleO, void* devPtrAmaxO, void* devPtrAmaxS, void* devPtrcuSeqlensQ, + void* devPtrcuSeqlensKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, + cudnn_frontend::DataType_t qkv_tensor_type, cudnn_frontend::DataType_t o_tensor_type, + NVTEScalingMode scaling_mode, NVTE_QKV_Format qkv_scale_inv_format, void* workspace, + size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; + const auto cudnn_runtime_version = cudnnGetVersion(); bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || @@ -1669,19 +1672,27 @@ void fused_attn_fp8_fwd_impl_v1( bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); bool is_dropout = (is_training && dropout_probability != 0.0f); + bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); auto bias_b = b; auto bias_h = h; auto bias_sq = s_q; auto bias_skv = s_kv; NVTE_CHECK(~is_bias, "FP8 fused attention does not support pre/post_scale_bias yet!"); NVTE_CHECK(~is_alibi, "FP8 fused attention does not support ALiBi yet!"); - bool is_current_scaling = (o_tensor_type == cudnn_frontend::DataType_t::HALF || - o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - bool is_delayed_scaling = (o_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || + bool is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && + (o_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || o_tensor_type == cudnn_frontend::DataType_t::FP8_E5M2); - NVTE_CHECK(is_current_scaling || is_delayed_scaling, - "FP8 fused attention only supports O tensor in kFloat16, kBFloat16, kFloat8E4M3 or " - "kFloat8E5M2!"); + bool is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && + (o_tensor_type == cudnn_frontend::DataType_t::HALF || + o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + bool is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING) && + (o_tensor_type == cudnn_frontend::DataType_t::HALF || + o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + NVTE_CHECK( + is_delayed_scaling || is_current_scaling || is_mxfp8, + "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 recipes!"); + NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, + "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); try { FADescriptor_v1 descriptor{b, @@ -1689,8 +1700,8 @@ void fused_attn_fp8_fwd_impl_v1( hg, s_q, s_kv, - d, - d, + d_qk, + d_v, 0, 0, 0, @@ -1704,13 +1715,18 @@ void fused_attn_fp8_fwd_impl_v1( scaling_factor, is_training, dropout_probability, - layout, + qkv_layout, + o_format, + NVTE_QKV_Format_NOT_SET, + NVTE_QKV_Layout_NOT_SET, + qkv_scale_inv_format, + NVTE_QKV_Format_NOT_SET, bias_type, mask_type, - NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX, - 0, - 0, - true, + softmax_type, + window_size_left, + window_size_right, + bottom_right_diagonal, true, qkv_tensor_type, o_tensor_type, @@ -1736,6 +1752,7 @@ void fused_attn_fp8_fwd_impl_v1( std::shared_ptr, // amax_o std::shared_ptr, // Stats std::shared_ptr, // bias + std::shared_ptr, // softmax_offset std::shared_ptr, // seq_q std::shared_ptr, // seq_kv std::shared_ptr, // dropout_seed @@ -1762,31 +1779,28 @@ void fused_attn_fp8_fwd_impl_v1( std::shared_ptr Q, K, V, attn_scale; std::shared_ptr descale_q, descale_k, descale_v; std::shared_ptr descale_s, scale_s, scale_o; - std::shared_ptr bias, seq_q, seq_kv; + std::shared_ptr bias, softmax_offset, seq_q, seq_kv; std::shared_ptr dropout_seed, dropout_offset; - std::vector q_stride(4); - std::vector k_stride(4); - std::vector v_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d, q_stride.data(), layout, - NVTE_QKV_Matrix::NVTE_Q_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d, k_stride.data(), layout, - NVTE_QKV_Matrix::NVTE_K_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d, v_stride.data(), layout, - NVTE_QKV_Matrix::NVTE_V_Matrix); + // Q, K, V, attn_scale + std::vector q_strides(4), k_strides(4), v_strides(4); + generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, q_strides.data(), + k_strides.data(), v_strides.data(), qkv_layout); Q = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("Q") - .set_dim({b, h, s_q, d}) - .set_stride(q_stride)); + .set_dim({b, h, s_q, d_qk}) + .set_stride(q_strides) + .set_data_type(qkv_tensor_type)); K = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("K") - .set_dim({b, hg, s_kv, d}) - .set_stride(k_stride)); + .set_dim({b, hg, s_kv, d_qk}) + .set_stride(k_strides) + .set_data_type(qkv_tensor_type)); V = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("V") - .set_dim({b, hg, s_kv, d}) - .set_stride(v_stride)); - + .set_dim({b, hg, s_kv, d_v}) + .set_stride(v_strides) + .set_data_type(qkv_tensor_type)); attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("attn_scale") .set_dim({1, 1, 1, 1}) @@ -1794,21 +1808,61 @@ void fused_attn_fp8_fwd_impl_v1( .set_is_pass_by_value(true) .set_data_type(fe::DataType_t::FLOAT)); - descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - descale_k = mha_graph->tensor_like(descale_q, "Descale_q"); - descale_v = mha_graph->tensor_like(descale_q, "Descale_V"); - descale_s = mha_graph->tensor_like(descale_q, "Descale_S"); - scale_s = mha_graph->tensor_like(descale_q, "Scale_S"); - - if (is_delayed_scaling) { - scale_o = mha_graph->tensor_like(descale_q, "Scale_O"); - } - if (is_current_scaling) { - scale_o = mha_graph->tensor(1.0f); + // Descale_q, Descale_k, Descale_v, Descale_s, Scale_s, Scale_o + if (is_delayed_scaling || is_current_scaling) { + descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + descale_k = mha_graph->tensor_like(descale_q, "Descale_q"); + descale_v = mha_graph->tensor_like(descale_q, "Descale_v"); + descale_s = mha_graph->tensor_like(descale_q, "Descale_s"); + scale_s = mha_graph->tensor_like(descale_q, "Scale_s"); + if (is_delayed_scaling) { + scale_o = mha_graph->tensor_like(descale_q, "Scale_o"); + } + if (is_current_scaling) { + scale_o = mha_graph->tensor(1.0f); + } + } else if (is_mxfp8) { + NVTE_QKV_Format q_scale_inv_format = (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) + ? qkv_scale_inv_format + : nvte_get_q_format(qkv_layout); + NVTE_QKV_Format kv_scale_inv_format = (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) + ? qkv_scale_inv_format + : nvte_get_kv_format(qkv_layout); + std::vector q_scale_strides(4); + std::vector k_scale_strides(4); + std::vector v_scale_strides(4); + auto padded = pad_s_d_for_mxfp8(s_q, s_kv, d_qk, d_v); + generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_qk_scale_padded, + q_scale_strides.data(), q_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_qk_scale_padded, + k_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_scale_padded, padded.d_v_padded, + v_scale_strides.data(), kv_scale_inv_format); + descale_q = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q") + .set_dim({b, h, padded.s_q_padded, padded.d_qk_scale_padded}) + .set_stride(q_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_k = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_k") + .set_dim({b, hg, padded.s_kv_padded, padded.d_qk_scale_padded}) + .set_stride(k_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_v = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_v") + .set_dim({b, hg, padded.s_kv_scale_padded, padded.d_v_padded}) + .set_stride(v_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); } fe::graph::SDPA_fp8_attributes sdpa_options; @@ -1818,6 +1872,20 @@ void fused_attn_fp8_fwd_impl_v1( .set_causal_mask(is_causal) .set_attn_scale(attn_scale); + fe::DiagonalAlignment_t const& diagonal_alignment = + bottom_right_diagonal ? fe::DiagonalAlignment_t::BOTTOM_RIGHT + : fe::DiagonalAlignment_t::TOP_LEFT; + sdpa_options.set_diagonal_alignment(diagonal_alignment); + + if (cudnn_runtime_version >= 92100) { + if (window_size_left != -1) { + sdpa_options.set_diagonal_band_left_bound(window_size_left + 1); + } + if (window_size_right != -1) { + sdpa_options.set_diagonal_band_right_bound(window_size_right); + } + } + // sdpa_options.set_alibi_mask(is_alibi); // if (is_bias) { // bias = mha_graph->tensor(fe::graph::Tensor_attributes() @@ -1855,19 +1923,41 @@ void fused_attn_fp8_fwd_impl_v1( sdpa_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); } - auto [O, Stats, amax_s, amax_o] = mha_graph->sdpa_fp8( - Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, scale_o, sdpa_options); + if (is_softmax_offset) { + softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_options.set_sink_token(softmax_offset); + } - std::vector o_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d, o_stride.data(), layout, - NVTE_QKV_Matrix::NVTE_O_Matrix); - O->set_output(true).set_dim({b, h, s_q, d}).set_stride(o_stride).set_data_type(o_tensor_type); - amax_o->set_output(true) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); + std::shared_ptr O, Stats, amax_s, amax_o; + if (is_delayed_scaling || is_current_scaling) { + auto outputs = mha_graph->sdpa_fp8(Q, K, V, descale_q, descale_k, descale_v, descale_s, + scale_s, scale_o, sdpa_options); + O = outputs[0]; + Stats = outputs[1]; + amax_s = outputs[2]; + amax_o = outputs[3]; + amax_s->set_output(true) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT); + } else if (is_mxfp8) { + auto outputs = mha_graph->sdpa_fp8(Q, K, V, descale_q, descale_k, descale_v, sdpa_options); + O = outputs[0]; + Stats = outputs[1]; + amax_o = outputs[2]; + } - amax_s->set_output(true) + std::vector o_strides(4); + generateMatrixStridesWithFormat(b, h, s_q, d_v, o_strides.data(), o_format); + O->set_output(true) + .set_dim({b, h, s_q, d_v}) + .set_stride(o_strides) + .set_data_type(o_tensor_type); + amax_o->set_output(!is_mxfp8) .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) .set_data_type(fe::DataType_t::FLOAT); @@ -1890,10 +1980,15 @@ void fused_attn_fp8_fwd_impl_v1( std::shared_ptr, // O std::shared_ptr, // amax_s std::shared_ptr> // amax_o - key_tensors_tuple = std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, descale_s, - scale_s, scale_o, attn_scale, O, amax_s, amax_o); + key_tensors_tuple = + is_mxfp8 ? std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, nullptr, nullptr, + nullptr, attn_scale, O, nullptr, amax_o) + : std::make_tuple(Q, K, V, descale_q, descale_k, descale_v, descale_s, + scale_s, scale_o, attn_scale, O, amax_s, amax_o); auto Stats_tuple = std::make_tuple(Stats); auto bias_tuple = is_bias ? std::make_tuple(bias) : std::make_tuple(nullptr); + auto softmax_offset_tuple = + is_softmax_offset ? std::make_tuple(softmax_offset) : std::make_tuple(nullptr); auto padding_tuple = is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) @@ -1904,17 +1999,17 @@ void fused_attn_fp8_fwd_impl_v1( NVTE_CHECK_CUDNN_FE(mha_graph->create_execution_plans({fe::HeurMode_t::A})); NVTE_CHECK_CUDNN_FE(mha_graph->check_support(handle)); NVTE_CHECK_CUDNN_FE(mha_graph->build_plans(handle)); - - auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, - bias_tuple, padding_tuple, dropout_tuple); + auto return_tuple = + std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, Stats_tuple, bias_tuple, + softmax_offset_tuple, padding_tuple, dropout_tuple); cache.insert({descriptor, return_tuple}); return return_tuple; }; auto [mha_graph, Q, K, V, descale_q, descale_k, descale_v, descale_s, scale_s, scale_o, - attn_scale, O, amax_s, amax_o, Stats, bias, seq_q, seq_kv, dropout_seed, dropout_offset] = - get_graph(sdpa_fp8_fprop_cache, descriptor); + attn_scale, O, amax_s, amax_o, Stats, bias, softmax_offset, seq_q, seq_kv, dropout_seed, + dropout_offset] = get_graph(sdpa_fp8_fprop_cache, descriptor); auto plan_workspace_size = mha_graph->get_workspace_size(); @@ -1937,17 +2032,19 @@ void fused_attn_fp8_fwd_impl_v1( {descale_q, devPtrDescaleQ}, {descale_k, devPtrDescaleK}, {descale_v, devPtrDescaleV}, - {descale_s, devPtrDescaleS}, - {scale_s, devPtrScaleS}, {attn_scale, &scaling_factor}, {O, devPtrO}, - {amax_s, devPtrAmaxS}, - {amax_o, devPtrAmaxO}, {Stats, devPtrM}}; if (is_delayed_scaling) { variant_pack[scale_o] = devPtrScaleO; } + if (is_delayed_scaling || is_current_scaling) { + variant_pack[descale_s] = devPtrDescaleS; + variant_pack[scale_s] = devPtrScaleS; + variant_pack[amax_s] = devPtrAmaxS; + variant_pack[amax_o] = devPtrAmaxO; + } /* if (is_bias) { variant_pack[bias] = devPtrBias; @@ -1972,6 +2069,10 @@ void fused_attn_fp8_fwd_impl_v1( variant_pack[dropout_offset] = devPtrDropoutOffset; } + if (is_softmax_offset) { + variant_pack[softmax_offset] = devPtrSoftmaxOffset; + } + NVTE_CHECK_CUDNN_FE(mha_graph->execute(handle, variant_pack, workspace)); } catch (cudnn_frontend::cudnnException& e) { NVTE_ERROR(e.what()); @@ -1980,20 +2081,27 @@ void fused_attn_fp8_fwd_impl_v1( // fused attention BWD FP8 with FE 1.0+ void fused_attn_fp8_bwd_impl_v1( - int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d, float scaling_factor, - float dropout_probability, NVTE_QKV_Layout layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, bool deterministic, void* devPtrQ, void* devPtrK, void* devPtrV, - void* devPtrM, void* devPtrZInv, void* devPtrO, void* devPtrdO, void* devPtrdQ, void* devPtrdK, - void* devPtrdV, void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, - void* devPtrDescaleO, void* devPtrDescaledO, void* devPtrDescaleS, void* devPtrDescaledP, - void* devPtrScaleS, void* devPtrScaledP, void* devPtrScaledQ, void* devPtrScaledK, - void* devPtrScaledV, void* devPtrAmaxdP, void* devPtrAmaxdQ, void* devPtrAmaxdK, - void* devPtrAmaxdV, void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, void* devPtrDropoutSeed, - void* devPtrDropoutOffset, cudnn_frontend::DataType_t qkv_tensor_type, + int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, + float scaling_factor, float dropout_probability, NVTE_QKV_Layout qkv_layout, + NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, + NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, + int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, + bool deterministic, void* devPtrQ, void* devPtrK, void* devPtrV, void* devPtrM, + void* devPtrZInv, void* devPtrO, void* devPtrdO, void* devPtrSoftmaxOffset, void* devPtrdQ, + void* devPtrdK, void* devPtrdV, void* devPtrdSoftmaxOffset, void* devPtrDescaleQ, + void* devPtrDescaleK, void* devPtrDescaleV, void* devPtrDescaleO, void* devPtrDescaledO, + void* devPtrDescaleS, void* devPtrDescaledP, void* devPtrScaleS, void* devPtrScaledP, + void* devPtrScaledQ, void* devPtrScaledK, void* devPtrScaledV, void* devPtrAmaxdP, + void* devPtrAmaxdQ, void* devPtrAmaxdK, void* devPtrAmaxdV, void* devPtrQ_t, void* devPtrK_t, + void* devPtrdO_f16, void* devPtrdO_t, void* devPtrDescaleQ_t, void* devPtrDescaleK_t, + void* devPtrDescaledO_t, void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, + void* devPtrDropoutSeed, void* devPtrDropoutOffset, cudnn_frontend::DataType_t qkv_tensor_type, cudnn_frontend::DataType_t o_tensor_type, cudnn_frontend::DataType_t do_tensor_type, - cudnn_frontend::DataType_t dqkv_tensor_type, void* workspace, size_t* workspace_size, - cudaStream_t stream, cudnnHandle_t handle) { + cudnn_frontend::DataType_t dqkv_tensor_type, NVTEScalingMode scaling_mode, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, void* workspace, + size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; + const auto cudnn_runtime_version = cudnnGetVersion(); bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); bool is_alibi = (bias_type == NVTE_Bias_Type::NVTE_ALIBI); bool is_causal = ((mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || @@ -2001,20 +2109,28 @@ void fused_attn_fp8_bwd_impl_v1( bool is_padding = ((mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)); bool is_dropout = (dropout_probability != 0.0f); + bool is_softmax_offset = (softmax_type != NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX); auto bias_b = b; auto bias_h = h; - const auto cudnn_runtime_version = cudnnGetVersion(); auto bias_sq = s_q; auto bias_skv = s_kv; NVTE_CHECK(~is_bias, "FP8 fused attention does not support pre/post_scale_bias yet!"); NVTE_CHECK(~is_alibi, "FP8 fused attention does not support ALiBi yet!"); - bool is_current_scaling = (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || - dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); - bool is_delayed_scaling = (dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || + bool is_delayed_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && + (dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E4M3 || dqkv_tensor_type == cudnn_frontend::DataType_t::FP8_E5M2); - NVTE_CHECK(is_current_scaling || is_delayed_scaling, - "FP8 fused attention only supports dQKV tensor in kFloat16, kBFloat16, kFloat8E4M3 or " - "kFloat8E5M2!"); + bool is_current_scaling = (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) && + (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || + dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + bool is_mxfp8 = (scaling_mode == NVTE_MXFP8_1D_SCALING) && + (dqkv_tensor_type == cudnn_frontend::DataType_t::HALF || + dqkv_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); + NVTE_CHECK( + is_delayed_scaling || is_current_scaling || is_mxfp8, + "FP8 fused attention only supports FP8DelayedScaling or FP8CurrentScaling or MXFP8 recipes!"); + NVTE_CHECK(!is_mxfp8 || cudnn_runtime_version >= 92100, + "MXFP8 fused attention requires cuDNN 9.21.0 or later!"); + bool is_O_in_F16 = (o_tensor_type == cudnn_frontend::DataType_t::HALF || o_tensor_type == cudnn_frontend::DataType_t::BFLOAT16); @@ -2024,8 +2140,8 @@ void fused_attn_fp8_bwd_impl_v1( hg, s_q, s_kv, - d, - d, + d_qk, + d_v, 0, 0, 0, @@ -2039,13 +2155,18 @@ void fused_attn_fp8_bwd_impl_v1( scaling_factor, true, dropout_probability, - layout, + qkv_layout, + o_format, + do_format, + dqkv_layout, + qkv_scale_inv_format, + do_scale_inv_format, bias_type, mask_type, - NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX, - 0, - 0, - true, + softmax_type, + window_size_left, + window_size_right, + bottom_right_diagonal, deterministic, qkv_tensor_type, o_tensor_type, @@ -2056,18 +2177,25 @@ void fused_attn_fp8_bwd_impl_v1( namespace fe = cudnn_frontend; using graph_and_tensors = std::tuple, - std::shared_ptr, // q - std::shared_ptr, // k - std::shared_ptr, // v - std::shared_ptr, // o - std::shared_ptr, // stats + std::shared_ptr, // Q + std::shared_ptr, // Q_t + std::shared_ptr, // K + std::shared_ptr, // K_t + std::shared_ptr, // V + std::shared_ptr, // O + std::shared_ptr, // Stats std::shared_ptr, // dO + std::shared_ptr, // dO_t + std::shared_ptr, // dO_f16 std::shared_ptr, // attn_scale std::shared_ptr, // descale_q + std::shared_ptr, // descale_q_t std::shared_ptr, // descale_k + std::shared_ptr, // descale_k_t std::shared_ptr, // descale_v std::shared_ptr, // descale_o std::shared_ptr, // descale_dO + std::shared_ptr, // descale_dO_t std::shared_ptr, // descale_s std::shared_ptr, // descale_dP std::shared_ptr, // scale_dQ @@ -2084,6 +2212,8 @@ void fused_attn_fp8_bwd_impl_v1( std::shared_ptr, // amax_dP std::shared_ptr, // bias std::shared_ptr, // dBias + std::shared_ptr, // softmax_offset + std::shared_ptr, // d_softmax_offset std::shared_ptr, // seq_q std::shared_ptr, // seq_kv std::shared_ptr, // dropout_seed @@ -2108,54 +2238,54 @@ void fused_attn_fp8_bwd_impl_v1( .set_intermediate_data_type(fe::DataType_t::FLOAT) .set_compute_data_type(fe::DataType_t::FLOAT); - std::shared_ptr q, k, v, o, dO, stats, attn_scale; - std::shared_ptr descale_q, descale_k, descale_v; + std::shared_ptr Q, Q_t, K, K_t, V, O, dO, dO_t, dO_f16, Stats, + attn_scale; + std::shared_ptr descale_q, descale_q_t, descale_k, descale_k_t, + descale_v; std::shared_ptr descale_s, descale_o; - std::shared_ptr descale_dP, descale_dO; + std::shared_ptr descale_dP, descale_dO, descale_dO_t; std::shared_ptr scale_s, scale_dP; std::shared_ptr scale_dQ, scale_dK, scale_dV; - std::shared_ptr bias, dBias, seq_q, seq_kv; + std::shared_ptr bias, dBias, softmax_offset, d_softmax_offset; + std::shared_ptr seq_q, seq_kv; std::shared_ptr dropout_seed, dropout_offset; - std::vector q_stride(4); - std::vector k_stride(4); - std::vector v_stride(4); - std::vector o_stride(4); - generateMatrixStrides(b, h, s_q, s_kv, d, q_stride.data(), layout, - NVTE_QKV_Matrix::NVTE_Q_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d, k_stride.data(), layout, - NVTE_QKV_Matrix::NVTE_K_Matrix); - generateMatrixStrides(b, hg, s_q, s_kv, d, v_stride.data(), layout, - NVTE_QKV_Matrix::NVTE_V_Matrix); - generateMatrixStrides(b, h, s_q, s_kv, d, o_stride.data(), layout, - NVTE_QKV_Matrix::NVTE_O_Matrix); - q = mha_graph->tensor(fe::graph::Tensor_attributes() + // Q, K, V, O, dO, stats, attn_scale + std::vector q_strides(4), k_strides(4), v_strides(4), o_strides(4), dO_strides(4); + generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, q_strides.data(), + k_strides.data(), v_strides.data(), qkv_layout); + generateMatrixStridesWithFormat(b, h, s_q, d_v, o_strides.data(), o_format); + generateMatrixStridesWithFormat(b, h, s_q, d_v, dO_strides.data(), do_format); + Q = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("Q") - .set_dim({b, h, s_q, d}) - .set_stride(q_stride)); - k = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_dim({b, h, s_q, d_qk}) + .set_stride(q_strides) + .set_data_type(qkv_tensor_type)); + K = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("K") - .set_dim({b, hg, s_kv, d}) - .set_stride(k_stride)); - v = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_dim({b, hg, s_kv, d_qk}) + .set_stride(k_strides) + .set_data_type(qkv_tensor_type)); + V = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("V") - .set_dim({b, hg, s_kv, d}) - .set_stride(v_stride)); - o = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_dim({b, hg, s_kv, d_v}) + .set_stride(v_strides) + .set_data_type(qkv_tensor_type)); + O = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("O") - .set_dim({b, h, s_q, d}) - .set_stride(o_stride) + .set_dim({b, h, s_q, d_v}) + .set_stride(o_strides) .set_data_type(o_tensor_type)); dO = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("dO") - .set_dim({b, h, s_q, d}) - .set_stride(o_stride)); - stats = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("stats") + .set_dim({b, h, s_q, d_v}) + .set_stride(dO_strides) + .set_data_type(do_tensor_type)); + Stats = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Stats") .set_dim({b, h, s_q, 1}) .set_stride({h * s_q, s_q, 1, 1}) .set_data_type(fe::DataType_t::FLOAT)); - attn_scale = mha_graph->tensor(fe::graph::Tensor_attributes() .set_name("attn_scale") .set_dim({1, 1, 1, 1}) @@ -2163,33 +2293,136 @@ void fused_attn_fp8_bwd_impl_v1( .set_is_pass_by_value(true) .set_data_type(fe::DataType_t::FLOAT)); - descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() - .set_name("Descale_q") - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT)); - descale_k = mha_graph->tensor_like(descale_q, "Descale_q"); - descale_v = mha_graph->tensor_like(descale_q, "Descale_V"); - descale_s = mha_graph->tensor_like(descale_q, "Descale_S"); - descale_dP = mha_graph->tensor_like(descale_q, "Descale_dP"); - if (is_O_in_F16) { - descale_o = mha_graph->tensor(1.0f); - } else { - descale_o = mha_graph->tensor_like(descale_q, "Descale_O"); - } - descale_dO = mha_graph->tensor_like(descale_q, "Descale_dO"); - scale_s = mha_graph->tensor_like(descale_q, "Scale_S"); - scale_dP = mha_graph->tensor_like(descale_q, "Scale_dP"); - - if (is_delayed_scaling) { - scale_dQ = mha_graph->tensor_like(descale_q, "Scale_dQ"); - scale_dK = mha_graph->tensor_like(descale_q, "Scale_dK"); - scale_dV = mha_graph->tensor_like(descale_q, "Scale_dV"); - } - if (is_current_scaling) { - scale_dQ = mha_graph->tensor(1.0f); - scale_dK = mha_graph->tensor(1.0f); - scale_dV = mha_graph->tensor(1.0f); + // Descale_q, Descale_k, Descale_v, Descale_s, Scale_s, Descale_dP, Scale_dP, Descale_o, Descale_dO, Scale_dQ, Scale_dK, Scale_dV + if (is_delayed_scaling || is_current_scaling) { + descale_q = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q") + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + descale_k = mha_graph->tensor_like(descale_q, "Descale_q"); + descale_v = mha_graph->tensor_like(descale_q, "Descale_v"); + descale_s = mha_graph->tensor_like(descale_q, "Descale_s"); + scale_s = mha_graph->tensor_like(descale_q, "Scale_s"); + descale_dP = mha_graph->tensor_like(descale_q, "Descale_dP"); + scale_dP = mha_graph->tensor_like(descale_q, "Scale_dP"); + if (is_current_scaling && is_O_in_F16) { + descale_o = mha_graph->tensor(1.0f); + } else { + descale_o = mha_graph->tensor_like(descale_q, "Descale_O"); + } + descale_dO = mha_graph->tensor_like(descale_q, "Descale_dO"); + if (is_delayed_scaling) { + scale_dQ = mha_graph->tensor_like(descale_q, "Scale_dQ"); + scale_dK = mha_graph->tensor_like(descale_q, "Scale_dK"); + scale_dV = mha_graph->tensor_like(descale_q, "Scale_dV"); + } + if (is_current_scaling) { + scale_dQ = mha_graph->tensor(1.0f); + scale_dK = mha_graph->tensor(1.0f); + scale_dV = mha_graph->tensor(1.0f); + } + } else if (is_mxfp8) { + NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + NVTE_QKV_Format q_scale_inv_format = + (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : q_format; + NVTE_QKV_Format kv_scale_inv_format = + (qkv_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? qkv_scale_inv_format : kv_format; + NVTE_QKV_Format do_scale_format_ = + (do_scale_inv_format != NVTE_QKV_Format_NOT_SET) ? do_scale_inv_format : do_format; + // Q_t, K_t, dO_t, dO_f16 + std::vector q_t_strides(4), k_t_strides(4), dO_t_strides(4); + generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_t_strides.data(), q_format); + generateMatrixStridesWithFormat(b, hg, s_kv, d_qk, k_t_strides.data(), kv_format); + generateMatrixStridesWithFormat(b, h, s_q, d_v, dO_t_strides.data(), do_format); + Q_t = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Q_t") + .set_dim({b, h, s_q, d_qk}) + .set_stride(q_t_strides) + .set_data_type(qkv_tensor_type)); + K_t = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("K_t") + .set_dim({b, hg, s_kv, d_qk}) + .set_stride(k_t_strides) + .set_data_type(qkv_tensor_type)); + dO_t = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("dO_t") + .set_dim({b, h, s_q, d_v}) + .set_stride(dO_t_strides) + .set_data_type(do_tensor_type)); + dO_f16 = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("dO_f16") + .set_dim({b, h, s_q, d_v}) + .set_stride(dO_strides) + .set_data_type(o_tensor_type)); + // Descale_q, Descale_q_t, Descale_k, Descale_k_t, Descale_v, Descale_dO, Descale_dO_t + auto padded = pad_s_d_for_mxfp8(s_q, s_kv, d_qk, d_v); + std::vector q_scale_strides(4), q_t_scale_strides(4), k_scale_strides(4), + k_t_scale_strides(4), v_scale_strides(4), dO_scale_strides(4), dO_t_scale_strides(4); + generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_qk_scale_padded, + q_scale_strides.data(), q_scale_inv_format); + generateMatrixStridesWithFormat(b, h, padded.s_q_scale_padded, padded.d_qk_padded, + q_t_scale_strides.data(), q_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_qk_scale_padded, + k_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_scale_padded, padded.d_qk_padded, + k_t_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, hg, padded.s_kv_padded, padded.d_v_scale_padded, + v_scale_strides.data(), kv_scale_inv_format); + generateMatrixStridesWithFormat(b, h, padded.s_q_padded, padded.d_v_scale_padded, + dO_scale_strides.data(), do_scale_format_); + generateMatrixStridesWithFormat(b, h, padded.s_q_scale_padded, padded.d_v_padded, + dO_t_scale_strides.data(), do_scale_format_); + descale_q = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q") + .set_dim({b, h, padded.s_q_padded, padded.d_qk_scale_padded}) + .set_stride(q_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_q_t = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_q_t") + .set_dim({b, h, padded.s_q_scale_padded, padded.d_qk_padded}) + .set_stride(q_t_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_k = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_k") + .set_dim({b, hg, padded.s_kv_padded, padded.d_qk_scale_padded}) + .set_stride(k_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_k_t = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_k_t") + .set_dim({b, hg, padded.s_kv_scale_padded, padded.d_qk_padded}) + .set_stride(k_t_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_v = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_v") + .set_dim({b, hg, padded.s_kv_padded, padded.d_v_scale_padded}) + .set_stride(v_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_dO = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_dO") + .set_dim({b, h, padded.s_q_padded, padded.d_v_scale_padded}) + .set_stride(dO_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); + descale_dO_t = + mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("Descale_dO_t") + .set_dim({b, h, padded.s_q_scale_padded, padded.d_v_padded}) + .set_stride(dO_t_scale_strides) + .set_data_type(fe::DataType_t::FP8_E8M0) + .set_reordering_type(fe::TensorReordering_t::F8_128x4)); } fe::graph::SDPA_fp8_backward_attributes sdpa_backward_options; @@ -2198,6 +2431,20 @@ void fused_attn_fp8_bwd_impl_v1( .set_causal_mask(is_causal) .set_attn_scale(attn_scale); + fe::DiagonalAlignment_t const& diagonal_alignment = + bottom_right_diagonal ? fe::DiagonalAlignment_t::BOTTOM_RIGHT + : fe::DiagonalAlignment_t::TOP_LEFT; + sdpa_backward_options.set_diagonal_alignment(diagonal_alignment); + + if (cudnn_runtime_version >= 92100) { + if (window_size_left != -1) { + sdpa_backward_options.set_diagonal_band_left_bound(window_size_left + 1); + } + if (window_size_right != -1) { + sdpa_backward_options.set_diagonal_band_right_bound(window_size_right); + } + } + // sdpa_backward_options.set_alibi_mask(is_alibi); // if (is_bias) { @@ -2251,40 +2498,75 @@ void fused_attn_fp8_bwd_impl_v1( sdpa_backward_options.set_dropout(dropout_probability, dropout_seed, dropout_offset); } - auto [dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP] = mha_graph->sdpa_fp8_backward( - q, k, v, o, dO, stats, descale_q, descale_k, descale_v, descale_o, descale_dO, descale_s, - descale_dP, scale_s, scale_dQ, scale_dK, scale_dV, scale_dP, sdpa_backward_options); + if (is_softmax_offset) { + softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_backward_options.set_sink_token(softmax_offset); + d_softmax_offset = mha_graph->tensor(fe::graph::Tensor_attributes() + .set_name("d_softmax_offset") + .set_dim({1, h, 1, 1}) + .set_stride({h, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT)); + sdpa_backward_options.set_dsink_token(d_softmax_offset); + } - dQ->set_output(true).set_dim({b, h, s_q, d}).set_stride(q_stride); - dK->set_output(true).set_dim({b, hg, s_kv, d}).set_stride(k_stride); - dV->set_output(true).set_dim({b, hg, s_kv, d}).set_stride(v_stride); - amax_dQ->set_output(true) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_data_type(fe::DataType_t::FLOAT); - amax_dK->set_output(true) + std::shared_ptr dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP; + if (is_delayed_scaling || is_current_scaling) { + std::tie(dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP) = + std::apply([](const auto&... elems) { return std::make_tuple(elems...); }, + mha_graph->sdpa_fp8_backward(Q, K, V, O, dO, Stats, descale_q, descale_k, + descale_v, descale_o, descale_dO, descale_s, + descale_dP, scale_s, scale_dQ, scale_dK, + scale_dV, scale_dP, sdpa_backward_options)); + } else if (is_mxfp8) { + std::tie(dQ, dK, dV, amax_dQ, amax_dK, amax_dV) = std::apply( + [](const auto&... elems) { return std::make_tuple(elems...); }, + mha_graph->sdpa_fp8_backward(Q, Q_t, K, K_t, V, O, dO_f16, dO, dO_t, Stats, descale_q, + descale_q_t, descale_k, descale_k_t, descale_v, descale_dO, + descale_dO_t, sdpa_backward_options)); + } + std::vector dq_strides(4), dk_strides(4), dv_strides(4); + generateMatrixStridesWithLayout(b, h, hg, s_q, s_kv, d_qk, d_v, dq_strides.data(), + dk_strides.data(), dv_strides.data(), dqkv_layout); + dQ->set_output(true) + .set_dim({b, h, s_q, d_qk}) + .set_stride(dq_strides) + .set_data_type(dqkv_tensor_type); + dK->set_output(true) + .set_dim({b, hg, s_kv, d_qk}) + .set_stride(dk_strides) + .set_data_type(dqkv_tensor_type); + dV->set_output(true) + .set_dim({b, hg, s_kv, d_v}) + .set_stride(dv_strides) + .set_data_type(dqkv_tensor_type); + amax_dQ->set_output(!is_mxfp8) .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) .set_data_type(fe::DataType_t::FLOAT); - amax_dV->set_output(true) + amax_dK->set_output(!is_mxfp8) .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) .set_data_type(fe::DataType_t::FLOAT); - amax_dP->set_output(true) + amax_dV->set_output(!is_mxfp8) .set_dim({1, 1, 1, 1}) .set_stride({1, 1, 1, 1}) .set_data_type(fe::DataType_t::FLOAT); + if (is_delayed_scaling || is_current_scaling) { + amax_dP->set_output(true) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_data_type(fe::DataType_t::FLOAT); + } - dO->set_data_type(do_tensor_type); - dQ->set_data_type(dqkv_tensor_type); - dK->set_data_type(dqkv_tensor_type); - dV->set_data_type(dqkv_tensor_type); - - std::tuple, // q - std::shared_ptr, // k - std::shared_ptr, // v - std::shared_ptr, // o - std::shared_ptr, // stats + std::tuple, // Q + std::shared_ptr, // K + std::shared_ptr, // V + std::shared_ptr, // O + std::shared_ptr, // Stats std::shared_ptr, // dO std::shared_ptr, // attn_scale std::shared_ptr, // descale_q @@ -2307,10 +2589,16 @@ void fused_attn_fp8_bwd_impl_v1( std::shared_ptr, // amax_dV std::shared_ptr> // amax_dP key_tensors_tuple = std::make_tuple( - q, k, v, o, stats, dO, attn_scale, descale_q, descale_k, descale_v, descale_o, + Q, K, V, O, Stats, dO, attn_scale, descale_q, descale_k, descale_v, descale_o, descale_dO, descale_s, descale_dP, scale_s, scale_dQ, scale_dK, scale_dV, scale_dP, dQ, dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP); + auto mxfp8_tensors_tuple = + is_mxfp8 ? std::make_tuple(Q_t, K_t, dO_f16, dO_t, descale_q_t, descale_k_t, descale_dO_t) + : std::make_tuple(nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr); auto bias_tuple = is_bias ? std::make_tuple(bias, dBias) : std::make_tuple(nullptr, nullptr); + auto softmax_offset_tuple = is_softmax_offset + ? std::make_tuple(softmax_offset, d_softmax_offset) + : std::make_tuple(nullptr, nullptr); auto padding_tuple = is_padding ? std::make_tuple(seq_q, seq_kv) : std::make_tuple(nullptr, nullptr); auto dropout_tuple = is_dropout ? std::make_tuple(dropout_seed, dropout_offset) @@ -2322,17 +2610,18 @@ void fused_attn_fp8_bwd_impl_v1( NVTE_CHECK_CUDNN_FE(mha_graph->check_support(handle)); NVTE_CHECK_CUDNN_FE(mha_graph->build_plans(handle)); - auto return_tuple = std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, bias_tuple, - padding_tuple, dropout_tuple); + auto return_tuple = + std::tuple_cat(std::make_tuple(mha_graph), key_tensors_tuple, mxfp8_tensors_tuple, + bias_tuple, softmax_offset_tuple, padding_tuple, dropout_tuple); cache.insert({descriptor, return_tuple}); return return_tuple; }; - - auto [mha_graph, q, k, v, o, stats, dO, attn_scale, descale_q, descale_k, descale_v, descale_o, + auto [mha_graph, Q, K, V, O, Stats, dO, attn_scale, descale_q, descale_k, descale_v, descale_o, descale_dO, descale_s, descale_dP, scale_s, scale_dQ, scale_dK, scale_dV, scale_dP, dQ, - dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP, bias, dBias, seq_q, seq_kv, dropout_seed, - dropout_offset] = get_graph(sdpa_fp8_bprop_cache, descriptor); + dK, dV, amax_dQ, amax_dK, amax_dV, amax_dP, Q_t, K_t, dO_f16, dO_t, descale_q_t, + descale_k_t, descale_dO_t, bias, dBias, softmax_offset, d_softmax_offset, seq_q, seq_kv, + dropout_seed, dropout_offset] = get_graph(sdpa_fp8_bprop_cache, descriptor); auto plan_workspace_size = mha_graph->get_workspace_size(); @@ -2349,37 +2638,47 @@ void fused_attn_fp8_bwd_impl_v1( // build variant pack std::unordered_map, void*> variant_pack = { - {q, devPtrQ}, - {k, devPtrK}, - {v, devPtrV}, - {o, devPtrO}, - {stats, devPtrM}, + {Q, devPtrQ}, + {K, devPtrK}, + {V, devPtrV}, + {O, devPtrO}, + {Stats, devPtrM}, {dO, devPtrdO}, {attn_scale, &scaling_factor}, {descale_q, devPtrDescaleQ}, {descale_k, devPtrDescaleK}, {descale_v, devPtrDescaleV}, {descale_dO, devPtrDescaledO}, - {descale_s, devPtrDescaleS}, - {descale_dP, devPtrDescaledP}, - {scale_s, devPtrScaleS}, - {scale_dP, devPtrScaledP}, {dQ, devPtrdQ}, {dK, devPtrdK}, {dV, devPtrdV}, - {amax_dQ, devPtrAmaxdQ}, - {amax_dK, devPtrAmaxdK}, - {amax_dV, devPtrAmaxdV}, - {amax_dP, devPtrAmaxdP}, }; - + if (is_delayed_scaling || is_current_scaling) { + variant_pack[descale_s] = devPtrDescaleS; + variant_pack[descale_dP] = devPtrDescaledP; + variant_pack[scale_s] = devPtrScaleS; + variant_pack[scale_dP] = devPtrScaledP; + variant_pack[amax_dP] = devPtrAmaxdP; + variant_pack[amax_dQ] = devPtrAmaxdQ; + variant_pack[amax_dK] = devPtrAmaxdK; + variant_pack[amax_dV] = devPtrAmaxdV; + } + if (is_delayed_scaling || (is_current_scaling && !is_O_in_F16)) { + variant_pack[descale_o] = devPtrDescaleO; + } if (is_delayed_scaling) { variant_pack[scale_dQ] = devPtrScaledQ; variant_pack[scale_dK] = devPtrScaledK; variant_pack[scale_dV] = devPtrScaledV; } - if (!is_O_in_F16) { - variant_pack[descale_o] = devPtrDescaleO; + if (is_mxfp8) { + variant_pack[Q_t] = devPtrQ_t; + variant_pack[K_t] = devPtrK_t; + variant_pack[dO_f16] = devPtrdO_f16; + variant_pack[dO_t] = devPtrdO_t; + variant_pack[descale_q_t] = devPtrDescaleQ_t; + variant_pack[descale_k_t] = devPtrDescaleK_t; + variant_pack[descale_dO_t] = devPtrDescaledO_t; } /* if (is_bias) { @@ -2410,70 +2709,100 @@ void fused_attn_fp8_bwd_impl_v1( variant_pack[dropout_offset] = devPtrDropoutOffset; } + if (is_softmax_offset) { + variant_pack[softmax_offset] = devPtrSoftmaxOffset; + variant_pack[d_softmax_offset] = devPtrdSoftmaxOffset; + } + NVTE_CHECK_CUDNN_FE(mha_graph->execute(handle, variant_pack, workspace)); } catch (cudnn_frontend::cudnnException& e) { NVTE_ERROR(e.what()); } -} - -#endif +} // NOLINT(readability/fn_size) } // namespace fused_attn -#if (CUDNN_VERSION >= 8900) // fused attention FWD FP8 with separate Q, K, V -void fused_attn_fp8_fwd(size_t batch, size_t num_attn_heads, size_t num_gqa_groups, - size_t max_seqlen_q, size_t max_seqlen_kv, size_t head_dim, - bool is_training, float attn_scale, float p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, const Tensor* input_Q, const Tensor* input_K, - const Tensor* input_V, Tensor* input_output_S, Tensor* output_O, - NVTETensorPack* Aux_CTX_Tensors, const Tensor* cu_seqlens_q, - const Tensor* cu_seqlens_kv, const Tensor* rng_state, Tensor* workspace, - cudaStream_t stream, cudnnHandle_t handle) { +void fused_attn_fp8_fwd( + size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, + size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, bool is_training, float attn_scale, + float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, + bool bottom_right_diagonal, const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, + const Tensor* input_SoftmaxOffset, Tensor* input_output_S, Tensor* output_O, + NVTETensorPack* Aux_CTX_Tensors, const Tensor* cu_seqlens_q, const Tensor* cu_seqlens_kv, + const Tensor* rng_state, Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; - void* devPtrQ = input_Q->data.dptr; - void* devPtrK = input_K->data.dptr; - void* devPtrV = input_V->data.dptr; - void* devPtrDescaleQ = input_Q->scale_inv.dptr; - void* devPtrDescaleK = input_Q->scale_inv.dptr; - void* devPtrDescaleV = input_Q->scale_inv.dptr; - - void* devPtrO = output_O->data.dptr; - void* devPtrAmaxO = output_O->amax.dptr; - void* devPtrScaleO = output_O->scale.dptr; - + void *devPtrQ = nullptr, *devPtrK = nullptr, *devPtrV = nullptr; + void *devPtrDescaleQ = nullptr, *devPtrDescaleK = nullptr, *devPtrDescaleV = nullptr; + void *devPtrO = nullptr, *devPtrAmaxO = nullptr, *devPtrScaleO = nullptr; + void *devPtrAmaxS = nullptr, *devPtrScaleS = nullptr, *devPtrDescaleS = nullptr; + devPtrQ = input_Q->data.dptr; + devPtrDescaleQ = input_Q->scale_inv.dptr; + devPtrK = input_K->data.dptr; + devPtrDescaleK = input_K->scale_inv.dptr; + devPtrO = output_O->data.dptr; + if (input_Q->scaling_mode == NVTE_DELAYED_TENSOR_SCALING) { + devPtrV = input_V->data.dptr; + devPtrDescaleV = input_V->scale_inv.dptr; + devPtrScaleO = output_O->scale.dptr; + devPtrAmaxS = input_output_S->amax.dptr; + devPtrScaleS = input_output_S->scale.dptr; + devPtrDescaleS = input_output_S->scale_inv.dptr; + devPtrAmaxO = output_O->amax.dptr; + } else if (input_Q->scaling_mode == NVTE_MXFP8_1D_SCALING) { + devPtrV = input_V->columnwise_data.dptr; + devPtrDescaleV = input_V->columnwise_scale_inv.dptr; + } + void* devPtrSoftmaxOffset = nullptr; + if (softmax_type != NVTE_VANILLA_SOFTMAX) { + devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; + } void* devPtrM = nullptr; void* devPtrZInv = nullptr; if (Aux_CTX_Tensors->size == 0) { - Aux_CTX_Tensors->size = 3; - Tensor* output_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[0]); - Tensor* output_ZInv = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[1]); - Tensor* output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[2]); + int i = 0; + Tensor* output_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_M->data.dptr = nullptr; output_M->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; output_M->data.dtype = DType::kFloat32; - output_ZInv->data.dptr = nullptr; - output_ZInv->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; - output_ZInv->data.dtype = DType::kFloat32; + if (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD) { + Tensor* output_ZInv = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + output_ZInv->data.dptr = nullptr; + output_ZInv->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; + output_ZInv->data.dtype = DType::kFloat32; + } + Tensor* output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_rng_state->data.dptr = nullptr; output_rng_state->data.shape = {2}; output_rng_state->data.dtype = DType::kInt64; - } else if (Aux_CTX_Tensors->size == 3) { - Tensor* output_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[0]); - Tensor* output_ZInv = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[1]); - Tensor* output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[2]); + if (softmax_type != NVTE_VANILLA_SOFTMAX) { + Tensor* output_softmax_offset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + output_softmax_offset->data.dptr = nullptr; + output_softmax_offset->data.shape = {1, num_attn_heads, 1, 1}; + output_softmax_offset->data.dtype = DType::kFloat32; + } + Aux_CTX_Tensors->size = i; + } else if (Aux_CTX_Tensors->size >= 2) { + int i = 0; + Tensor* output_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); devPtrM = output_M->data.dptr; - devPtrZInv = output_ZInv->data.dptr; + devPtrZInv = nullptr; + if (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD) { + Tensor* output_ZInv = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + devPtrZInv = output_ZInv->data.dptr; + } + Tensor* output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_rng_state->data.dptr = rng_state->data.dptr; + if (softmax_type != NVTE_VANILLA_SOFTMAX) { + Tensor* output_softmax_offset = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); + output_softmax_offset->data.dptr = devPtrSoftmaxOffset; + } } else { NVTE_ERROR("Unexpected Aux_CTX_Tensors->size."); } - void* devPtrAmaxS = input_output_S->amax.dptr; - void* devPtrScaleS = input_output_S->scale.dptr; - void* devPtrDescaleS = input_output_S->scale_inv.dptr; - void* devPtrcuSeqlensQ = reinterpret_cast(reinterpret_cast(cu_seqlens_q->data.dptr)); void* devPtrcuSeqlensKV = @@ -2488,17 +2817,20 @@ void fused_attn_fp8_fwd(size_t batch, size_t num_attn_heads, size_t num_gqa_grou size_t workspace_size = 0; NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); - if ((qkv_format == NVTE_QKV_Format::NVTE_BSHD) || (qkv_format == NVTE_QKV_Format::NVTE_SBHD)) { + if ((qkv_format == NVTE_QKV_Format::NVTE_BSHD) || (qkv_format == NVTE_QKV_Format::NVTE_SBHD) || + (qkv_format == NVTE_QKV_Format::NVTE_BHSD)) { fused_attn::fused_attn_fp8_fwd_impl_v1( - batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim, is_training, - attn_scale, p_dropout, qkv_layout, bias_type, mask_type, devPtrQ, devPtrK, devPtrV, devPtrM, - devPtrZInv, devPtrO, devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleS, - devPtrScaleS, devPtrScaleO, devPtrAmaxO, devPtrAmaxS, devPtrcuSeqlensQ, devPtrcuSeqlensKV, - devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), - get_cudnn_fe_dtype(O_type), workspace->data.dptr, &workspace_size, stream, handle); + batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, + is_training, attn_scale, p_dropout, qkv_layout, o_format, bias_type, mask_type, + softmax_type, window_size_left, window_size_right, bottom_right_diagonal, devPtrQ, devPtrK, + devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrZInv, devPtrO, devPtrDescaleQ, devPtrDescaleK, + devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, devPtrAmaxO, devPtrAmaxS, + devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, devPtrDropoutOffset, + get_cudnn_fe_dtype(QKV_type), get_cudnn_fe_dtype(O_type), input_Q->scaling_mode, + qkv_scale_inv_format, workspace->data.dptr, &workspace_size, stream, handle); } else if (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD) { fused_attn::fused_attn_fp8_fwd_impl( - batch, num_attn_heads, max_seqlen_q, max_seqlen_kv, head_dim, is_training, attn_scale, + batch, num_attn_heads, max_seqlen_q, max_seqlen_kv, head_dim_qk, is_training, attn_scale, p_dropout, qkv_layout, devPtrQ, devPtrK, devPtrV, devPtrM, devPtrZInv, devPtrO, devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, devPtrAmaxO, devPtrAmaxS, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, @@ -2521,24 +2853,35 @@ void fused_attn_fp8_fwd(size_t batch, size_t num_attn_heads, size_t num_gqa_grou } } // fused attention BWD FP8 with separate Q, K, V -void fused_attn_fp8_bwd(size_t batch, size_t num_attn_heads, size_t num_gqa_groups, - size_t max_seqlen_q, size_t max_seqlen_kv, size_t head_dim, - float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, bool deterministic, - const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, - const Tensor* input_O, const Tensor* input_dO, const Tensor* input_M, - const Tensor* input_ZInv, const Tensor* input_S, Tensor* input_output_dP, - const Tensor* output_dQ, const Tensor* output_dK, const Tensor* output_dV, - const Tensor* cu_seqlens_q, const Tensor* cu_seqlens_kv, - const Tensor* rng_state, Tensor* workspace, cudaStream_t stream, - cudnnHandle_t handle) { +void fused_attn_fp8_bwd( + size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, + size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, float attn_scale, float p_dropout, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, + NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format, + NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, + bool bottom_right_diagonal, bool deterministic, const Tensor* input_Q, const Tensor* input_K, + const Tensor* input_V, const Tensor* input_O, const Tensor* input_dO, + const Tensor* input_dO_f16, const Tensor* input_M, const Tensor* input_ZInv, + const Tensor* input_S, const Tensor* input_SoftmaxOffset, Tensor* input_output_dP, + const Tensor* output_dQ, const Tensor* output_dK, const Tensor* output_dV, + Tensor* output_dSoftmaxOffset, const Tensor* cu_seqlens_q, const Tensor* cu_seqlens_kv, + const Tensor* rng_state, Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; void* devPtrQ = input_Q->data.dptr; void* devPtrK = input_K->data.dptr; void* devPtrV = input_V->data.dptr; void* devPtrDescaleQ = input_Q->scale_inv.dptr; - void* devPtrDescaleK = input_Q->scale_inv.dptr; - void* devPtrDescaleV = input_Q->scale_inv.dptr; + void* devPtrDescaleK = input_K->scale_inv.dptr; + void* devPtrDescaleV = input_V->scale_inv.dptr; + void *devPtrQ_t = nullptr, *devPtrK_t = nullptr, *devPtrDescaleQ_t = nullptr, + *devPtrDescaleK_t = nullptr; + if (input_Q->scaling_mode == NVTE_MXFP8_1D_SCALING) { + devPtrQ_t = input_Q->columnwise_data.dptr; + devPtrDescaleQ_t = input_Q->columnwise_scale_inv.dptr; + devPtrK_t = input_K->columnwise_data.dptr; + devPtrDescaleK_t = input_K->columnwise_scale_inv.dptr; + } void* devPtrO = input_O->data.dptr; const DType O_type = input_O->data.dtype; @@ -2548,25 +2891,46 @@ void fused_attn_fp8_bwd(size_t batch, size_t num_attn_heads, size_t num_gqa_grou } void* devPtrdO = input_dO->data.dptr; void* devPtrDescaledO = input_dO->scale_inv.dptr; + void *devPtrdO_t = nullptr, *devPtrdO_f16 = nullptr, *devPtrDescaledO_t = nullptr; + if (input_dO->scaling_mode == NVTE_MXFP8_1D_SCALING) { + devPtrdO_t = input_dO->columnwise_data.dptr; + devPtrdO_f16 = input_dO_f16->data.dptr; + devPtrDescaledO_t = input_dO->columnwise_scale_inv.dptr; + } void* devPtrM = input_M->data.dptr; - void* devPtrZInv = input_ZInv->data.dptr; + void* devPtrZInv = (input_ZInv != nullptr) ? input_ZInv->data.dptr : nullptr; + + void *devPtrScaleS = nullptr, *devPtrDescaleS = nullptr, *devPtrAmaxdP = nullptr, + *devPtrScaledP = nullptr, *devPtrDescaledP = nullptr; + if (input_Q->scaling_mode == NVTE_DELAYED_TENSOR_SCALING) { + devPtrScaleS = input_S->scale.dptr; + devPtrDescaleS = input_S->scale_inv.dptr; + devPtrAmaxdP = input_output_dP->amax.dptr; + devPtrScaledP = input_output_dP->scale.dptr; + devPtrDescaledP = input_output_dP->scale_inv.dptr; + } - void* devPtrScaleS = input_S->scale.dptr; - void* devPtrDescaleS = input_S->scale_inv.dptr; - void* devPtrAmaxdP = input_output_dP->amax.dptr; - void* devPtrScaledP = input_output_dP->scale.dptr; - void* devPtrDescaledP = input_output_dP->scale_inv.dptr; + void* devPtrSoftmaxOffset = nullptr; + void* devPtrdSoftmaxOffset = nullptr; + if (softmax_type != NVTE_VANILLA_SOFTMAX) { + devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; + devPtrdSoftmaxOffset = output_dSoftmaxOffset->data.dptr; + } void* devPtrdQ = output_dQ->data.dptr; void* devPtrdK = output_dK->data.dptr; void* devPtrdV = output_dV->data.dptr; - void* devPtrAmaxdQ = output_dQ->amax.dptr; - void* devPtrAmaxdK = output_dQ->amax.dptr; - void* devPtrAmaxdV = output_dQ->amax.dptr; - void* devPtrScaledQ = output_dQ->scale.dptr; - void* devPtrScaledK = output_dQ->scale.dptr; - void* devPtrScaledV = output_dQ->scale.dptr; + void *devPtrAmaxdQ = nullptr, *devPtrAmaxdK = nullptr, *devPtrAmaxdV = nullptr, + *devPtrScaledQ = nullptr, *devPtrScaledK = nullptr, *devPtrScaledV = nullptr; + if (input_Q->scaling_mode == NVTE_DELAYED_TENSOR_SCALING) { + devPtrAmaxdQ = output_dQ->amax.dptr; + devPtrAmaxdK = output_dK->amax.dptr; + devPtrAmaxdV = output_dV->amax.dptr; + devPtrScaledQ = output_dQ->scale.dptr; + devPtrScaledK = output_dK->scale.dptr; + devPtrScaledV = output_dV->scale.dptr; + } void* devPtrcuSeqlensQ = reinterpret_cast(reinterpret_cast(cu_seqlens_q->data.dptr)); @@ -2582,21 +2946,29 @@ void fused_attn_fp8_bwd(size_t batch, size_t num_attn_heads, size_t num_gqa_grou const DType dQKV_type = output_dQ->data.dtype; size_t workspace_size = 0; - NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); - if ((qkv_format == NVTE_QKV_Format::NVTE_BSHD) || (qkv_format == NVTE_QKV_Format::NVTE_SBHD)) { + NVTE_QKV_Format dqkv_format = nvte_get_qkv_format(dqkv_layout); + if ((dqkv_format == NVTE_QKV_Format::NVTE_BSHD) || (dqkv_format == NVTE_QKV_Format::NVTE_SBHD) || + (dqkv_format == NVTE_QKV_Format::NVTE_BHSD)) { fused_attn::fused_attn_fp8_bwd_impl_v1( - batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim, attn_scale, - p_dropout, qkv_layout, bias_type, mask_type, deterministic, devPtrQ, devPtrK, devPtrV, - devPtrM, devPtrZInv, devPtrO, devPtrdO, devPtrdQ, devPtrdK, devPtrdV, devPtrDescaleQ, - devPtrDescaleK, devPtrDescaleV, devPtrDescaleO, devPtrDescaledO, devPtrDescaleS, - devPtrDescaledP, devPtrScaleS, devPtrScaledP, devPtrScaledQ, devPtrScaledK, devPtrScaledV, - devPtrAmaxdP, devPtrAmaxdQ, devPtrAmaxdK, devPtrAmaxdV, devPtrcuSeqlensQ, devPtrcuSeqlensKV, + batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, + attn_scale, p_dropout, qkv_layout, o_format, do_format, dqkv_layout, bias_type, mask_type, + softmax_type, window_size_left, window_size_right, bottom_right_diagonal, deterministic, + devPtrQ, devPtrK, devPtrV, devPtrM, devPtrZInv, devPtrO, devPtrdO, devPtrSoftmaxOffset, + devPtrdQ, devPtrdK, devPtrdV, devPtrdSoftmaxOffset, devPtrDescaleQ, devPtrDescaleK, + devPtrDescaleV, devPtrDescaleO, devPtrDescaledO, devPtrDescaleS, devPtrDescaledP, + devPtrScaleS, devPtrScaledP, devPtrScaledQ, devPtrScaledK, devPtrScaledV, devPtrAmaxdP, + devPtrAmaxdQ, devPtrAmaxdK, devPtrAmaxdV, devPtrQ_t, devPtrK_t, devPtrdO_f16, devPtrdO_t, + devPtrDescaleQ_t, devPtrDescaleK_t, devPtrDescaledO_t, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), get_cudnn_fe_dtype(O_type), get_cudnn_fe_dtype(dO_type), get_cudnn_fe_dtype(dQKV_type), - workspace->data.dptr, &workspace_size, stream, handle); - } else if (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD) { + input_dO->scaling_mode, qkv_scale_inv_format, do_scale_inv_format, workspace->data.dptr, + &workspace_size, stream, handle); + } else if (dqkv_layout == NVTE_QKV_Layout::NVTE_T3HD) { + // remove this when cuDNN FE supports FP8 + THD + NVTE_CHECK(input_ZInv != nullptr && input_ZInv->data.dptr != nullptr, + "ZInv tensor required for FP8 fused attention backward with T3HD layout."); fused_attn::fused_attn_fp8_bwd_impl( - batch, num_attn_heads, max_seqlen_q, max_seqlen_kv, head_dim, attn_scale, p_dropout, + batch, num_attn_heads, max_seqlen_q, max_seqlen_kv, head_dim_qk, attn_scale, p_dropout, qkv_layout, devPtrQ, devPtrK, devPtrV, devPtrM, devPtrZInv, devPtrO, devPtrdO, devPtrdQ, devPtrdK, devPtrdV, devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleO, devPtrDescaledO, devPtrDescaleS, devPtrDescaledP, devPtrScaleS, devPtrScaledP, @@ -2619,5 +2991,4 @@ void fused_attn_fp8_bwd(size_t batch, size_t num_attn_heads, size_t num_gqa_grou return; } } -#endif // end of CUDNN>=8900 } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index 225e700ef..aaf5039ee 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -12,29 +12,31 @@ #include "transformer_engine/transformer_engine.h" namespace transformer_engine { -#if (CUDNN_VERSION >= 8900) // fused attention FWD FP8 with separate Q, K, V -void fused_attn_fp8_fwd(size_t batch, size_t num_attn_heads, size_t num_gqa_groups, - size_t max_seqlen_q, size_t max_seqlen_kv, size_t head_dim, - bool is_training, float attn_scale, float p_dropout, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, const Tensor *input_Q, const Tensor *input_K, - const Tensor *input_V, Tensor *input_output_S, Tensor *output_O, - NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens_q, - const Tensor *cu_seqlens_kv, const Tensor *rng_state, Tensor *workspace, - cudaStream_t stream, cudnnHandle_t handle); +void fused_attn_fp8_fwd( + size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, + size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, bool is_training, float attn_scale, + float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, + bool bottom_right_diagonal, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, + const Tensor *input_SoftmaxOffset, Tensor *input_output_S, Tensor *output_O, + NVTETensorPack *Aux_CTX_Tensors, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, + const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); // fused attention BWD FP8 with separate Q, K, V -void fused_attn_fp8_bwd(size_t batch, size_t num_attn_heads, size_t num_gqa_groups, - size_t max_seqlen_q, size_t max_seqlen_kv, size_t head_dim, - float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, bool deterministic, - const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_O, const Tensor *input_dO, const Tensor *input_M, - const Tensor *input_ZInv, const Tensor *input_S, Tensor *input_output_dP, - const Tensor *output_dQ, const Tensor *output_dK, const Tensor *output_dV, - const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, - const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, - cudnnHandle_t handle); -#endif // end of CUDNN>=8900 +void fused_attn_fp8_bwd( + size_t batch, size_t num_attn_heads, size_t num_gqa_groups, size_t max_seqlen_q, + size_t max_seqlen_kv, size_t head_dim_qk, size_t head_dim_v, float attn_scale, float p_dropout, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, + NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format, + NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, + NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, + bool bottom_right_diagonal, bool deterministic, const Tensor *input_Q, const Tensor *input_K, + const Tensor *input_V, const Tensor *input_O, const Tensor *input_dO, + const Tensor *input_dO_f16, const Tensor *input_M, const Tensor *input_ZInv, + const Tensor *input_S, const Tensor *input_SoftmaxOffset, Tensor *input_output_dP, + const Tensor *output_dQ, const Tensor *output_dK, const Tensor *output_dV, + Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, + const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu index a897b0933..f37eeb0c6 100644 --- a/transformer_engine/common/fused_attn/utils.cu +++ b/transformer_engine/common/fused_attn/utils.cu @@ -293,6 +293,27 @@ void generateMatrixStrides(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int6 strideA[hidden_dim_idx] = 1; } break; + case NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD: + if ((matrix == NVTE_QKV_Matrix::NVTE_Q_Matrix) || + (matrix == NVTE_QKV_Matrix::NVTE_O_Matrix)) { + strideA[batch_dim_idx] = h * s_q * d; + strideA[head_dim_idx] = s_q * d; + strideA[seqlen_dim_idx] = d; + strideA[hidden_dim_idx] = 1; + } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix) || + (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix)) { + strideA[batch_dim_idx] = h * s_kv * d; + strideA[head_dim_idx] = s_kv * d; + strideA[seqlen_dim_idx] = d; + strideA[hidden_dim_idx] = 1; + } else if ((matrix == NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose) || + (matrix == NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose)) { + strideA[batch_dim_idx] = h * s_kv * d; + strideA[head_dim_idx] = s_kv * d; + strideA[seqlen_transpose_dim_idx] = d; + strideA[hidden_transpose_dim_idx] = 1; + } + break; } if (matrix == NVTE_QKV_Matrix::NVTE_S_Matrix) { diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index 1ec1616c4..c3736a6c6 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -14,6 +14,7 @@ #include #include +#include "../common.h" #include "transformer_engine/fused_attn.h" #include "transformer_engine/transformer_engine.h" @@ -27,11 +28,198 @@ enum NVTE_QKV_Matrix { NVTE_K_Matrix = 1, // keys NVTE_K_Matrix_Transpose = 2, // keys transposed NVTE_V_Matrix = 3, // values - NVTE_V_Matrix_Transpose = 4, // value matrix transposed + NVTE_V_Matrix_Transpose = 4, // values transposed NVTE_S_Matrix = 5, // output of GEMM1 NVTE_O_Matrix = 6, // final output }; +// Padded sizes for MXFP8 layout (s_q/s_kv/d_qk/d_v and their scaled dimensions) +struct MXFP8PaddedSizes { + int64_t s_q_padded; + int64_t s_kv_padded; + int64_t s_q_scale; + int64_t s_kv_scale; + int64_t s_q_scale_padded; + int64_t s_kv_scale_padded; + int64_t d_qk_padded; + int64_t d_v_padded; + int64_t d_qk_scale; + int64_t d_v_scale; + int64_t d_qk_scale_padded; + int64_t d_v_scale_padded; +}; + +// Pad s and d for MXFP8 quantization +inline MXFP8PaddedSizes pad_s_d_for_mxfp8(int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v) { + constexpr int64_t block_size = 32; + MXFP8PaddedSizes p; + p.s_q_padded = DIVUP_TO_MULTIPLE(s_q, 128); + p.s_kv_padded = DIVUP_TO_MULTIPLE(s_kv, 128); + p.s_q_scale = DIVUP(s_q, block_size); + p.s_kv_scale = DIVUP(s_kv, block_size); + p.s_q_scale_padded = DIVUP_TO_MULTIPLE(p.s_q_scale, 4); + p.s_kv_scale_padded = DIVUP_TO_MULTIPLE(p.s_kv_scale, 4); + p.d_qk_padded = DIVUP_TO_MULTIPLE(d_qk, 128); + p.d_v_padded = DIVUP_TO_MULTIPLE(d_v, 128); + p.d_qk_scale = DIVUP(d_qk, block_size); + p.d_v_scale = DIVUP(d_v, block_size); + p.d_qk_scale_padded = DIVUP_TO_MULTIPLE(p.d_qk_scale, 4); + p.d_v_scale_padded = DIVUP_TO_MULTIPLE(p.d_v_scale, 4); + return p; +} + +// Get matrix strides for a 4D tensor [batch_size, num_heads, sequence_len, head_dim] given a QKV format. +// strides must point to at least 4 int64_t elements. +inline void generateMatrixStridesWithFormat(int64_t b, int64_t h, int64_t s, int64_t d, + int64_t *strides, NVTE_QKV_Format format) { + constexpr int b_dim = 0; + constexpr int h_dim = 1; + constexpr int s_dim = 2; + constexpr int d_dim = 3; + + switch (format) { + case NVTE_QKV_Format::NVTE_BSHD: + case NVTE_QKV_Format::NVTE_THD: + strides[b_dim] = s * h * d; + strides[h_dim] = d; + strides[s_dim] = h * d; + strides[d_dim] = 1; + break; + case NVTE_QKV_Format::NVTE_SBHD: + strides[b_dim] = h * d; + strides[h_dim] = d; + strides[s_dim] = b * h * d; + strides[d_dim] = 1; + break; + case NVTE_QKV_Format::NVTE_BHSD: + strides[b_dim] = h * s * d; + strides[h_dim] = s * d; + strides[s_dim] = d; + strides[d_dim] = 1; + break; + default: + NVTE_CHECK(false, "Invalid format."); + break; + } +} + +// get matrix strides based on layout and matrix type +inline void generateMatrixStridesWithLayout(int64_t b, int64_t h, int64_t hg, int64_t s_q, + int64_t s_kv, int64_t d_qk, int64_t d_v, + int64_t *q_strides, int64_t *k_strides, + int64_t *v_strides, NVTE_QKV_Layout layout) { + constexpr int b_dim = 0; + constexpr int h_dim = 1; + constexpr int s_dim = 2; + constexpr int d_dim = 3; + const NVTE_QKV_Format q_format = nvte_get_q_format(layout); + const NVTE_QKV_Format kv_format = nvte_get_kv_format(layout); + + switch (layout) { + case NVTE_QKV_Layout::NVTE_SB3HD: + q_strides[b_dim] = 3 * h * d_qk; + q_strides[h_dim] = d_qk; + q_strides[s_dim] = b * 3 * h * d_qk; + q_strides[d_dim] = 1; + for (int i = 0; i < 4; i++) { + k_strides[i] = v_strides[i] = q_strides[i]; + } + break; + case NVTE_QKV_Layout::NVTE_SBH3D: + q_strides[b_dim] = 3 * h * d_qk; + q_strides[h_dim] = 3 * d_qk; + q_strides[s_dim] = b * 3 * h * d_qk; + q_strides[d_dim] = 1; + for (int i = 0; i < 4; i++) { + k_strides[i] = v_strides[i] = q_strides[i]; + } + break; + case NVTE_QKV_Layout::NVTE_SBHD_SB2HD: + generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_strides, q_format); + k_strides[b_dim] = 2 * hg * d_qk; + k_strides[h_dim] = d_qk; + k_strides[s_dim] = b * 2 * hg * d_qk; + k_strides[d_dim] = 1; + for (int i = 0; i < 4; i++) { + v_strides[i] = k_strides[i]; + } + break; + case NVTE_QKV_Layout::NVTE_SBHD_SBH2D: + generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_strides, q_format); + k_strides[b_dim] = 2 * hg * d_qk; + k_strides[h_dim] = 2 * d_qk; + k_strides[s_dim] = b * 2 * hg * d_qk; + k_strides[d_dim] = 1; + for (int i = 0; i < 4; i++) { + v_strides[i] = k_strides[i]; + } + break; + case NVTE_QKV_Layout::NVTE_BS3HD: + case NVTE_QKV_Layout::NVTE_T3HD: + q_strides[b_dim] = s_q * 3 * h * d_qk; + q_strides[h_dim] = d_qk; + q_strides[s_dim] = 3 * h * d_qk; + q_strides[d_dim] = 1; + for (int i = 0; i < 4; i++) { + k_strides[i] = v_strides[i] = q_strides[i]; + } + break; + case NVTE_QKV_Layout::NVTE_BSH3D: + case NVTE_QKV_Layout::NVTE_TH3D: + q_strides[b_dim] = s_q * 3 * h * d_qk; + q_strides[h_dim] = 3 * d_qk; + q_strides[s_dim] = 3 * h * d_qk; + q_strides[d_dim] = 1; + for (int i = 0; i < 4; i++) { + k_strides[i] = v_strides[i] = q_strides[i]; + } + break; + case NVTE_QKV_Layout::NVTE_BSHD_BS2HD: + case NVTE_QKV_Layout::NVTE_THD_T2HD: + generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_strides, q_format); + k_strides[b_dim] = s_kv * 2 * hg * d_qk; + k_strides[h_dim] = d_qk; + k_strides[s_dim] = 2 * hg * d_qk; + k_strides[d_dim] = 1; + for (int i = 0; i < 4; i++) { + v_strides[i] = k_strides[i]; + } + break; + case NVTE_QKV_Layout::NVTE_BSHD_BSH2D: + case NVTE_QKV_Layout::NVTE_THD_TH2D: + generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_strides, q_format); + k_strides[b_dim] = s_kv * 2 * hg * d_qk; + k_strides[h_dim] = 2 * d_qk; + k_strides[s_dim] = 2 * hg * d_qk; + k_strides[d_dim] = 1; + for (int i = 0; i < 4; i++) { + v_strides[i] = k_strides[i]; + } + break; + case NVTE_QKV_Layout::NVTE_SBHD_SBHD_SBHD: + case NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_SBHD_SBHD: + case NVTE_QKV_Layout::NVTE_BSHD_BSHD_BSHD: + case NVTE_QKV_Layout::NVTE_THD_THD_THD: + case NVTE_QKV_Layout::NVTE_THD_BSHD_BSHD: + case NVTE_QKV_Layout::NVTE_Paged_KV_BSHD_BSHD_BSHD: + case NVTE_QKV_Layout::NVTE_Paged_KV_THD_BSHD_BSHD: + case NVTE_QKV_Layout::NVTE_SBHD_BSHD_BSHD: + case NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_BSHD_BSHD: + case NVTE_QKV_Layout::NVTE_BSHD_SBHD_SBHD: + case NVTE_QKV_Layout::NVTE_THD_SBHD_SBHD: + case NVTE_QKV_Layout::NVTE_Paged_KV_BSHD_SBHD_SBHD: + case NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD: + case NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD: + generateMatrixStridesWithFormat(b, h, s_q, d_qk, q_strides, q_format); + generateMatrixStridesWithFormat(b, hg, s_kv, d_qk, k_strides, kv_format); + generateMatrixStridesWithFormat(b, hg, s_kv, d_v, v_strides, kv_format); + break; + default: + NVTE_CHECK(false, "Invalid layout."); + break; + } +} + void generateMatrixStrides(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, int64_t *strideA, NVTE_QKV_Layout layout, NVTE_QKV_Matrix matrix); @@ -106,7 +294,12 @@ struct FADescriptor_v1 { float attnScale; bool isTraining; float dropoutProbability; - NVTE_QKV_Layout layout; + NVTE_QKV_Layout qkv_layout; + NVTE_QKV_Format o_format; + NVTE_QKV_Format do_format; + NVTE_QKV_Layout dqkv_layout; + NVTE_QKV_Format qkv_scale_inv_format; + NVTE_QKV_Format do_scale_inv_format; NVTE_Bias_Type bias_type; NVTE_Mask_Type mask_type; NVTE_Softmax_Type softmax_type; @@ -123,17 +316,19 @@ struct FADescriptor_v1 { bool operator<(const FADescriptor_v1 &rhs) const { return std::tie(b, h, hg, s_q, s_kv, d_qk, d_v, num_pages_k, num_pages_v, page_size_k, page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, bias_b, bias_h, bias_sq, - bias_skv, attnScale, isTraining, dropoutProbability, layout, mask_type, + bias_skv, attnScale, isTraining, dropoutProbability, qkv_layout, o_format, + do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, deterministic, bias_type, qkv_tensor_type, o_tensor_type, do_tensor_type, dqkv_tensor_type, return_max_logit) < std::tie(rhs.b, rhs.h, rhs.hg, rhs.s_q, rhs.s_kv, rhs.d_qk, rhs.d_v, rhs.num_pages_k, rhs.num_pages_v, rhs.page_size_k, rhs.page_size_v, rhs.max_pages_per_seq_k, rhs.max_pages_per_seq_v, rhs.bias_b, rhs.bias_h, rhs.bias_sq, rhs.bias_skv, - rhs.attnScale, rhs.isTraining, rhs.dropoutProbability, rhs.layout, - rhs.mask_type, rhs.softmax_type, rhs.window_size_left, rhs.window_size_right, - rhs.bottom_right_diagonal, rhs.deterministic, rhs.bias_type, - rhs.qkv_tensor_type, rhs.o_tensor_type, rhs.do_tensor_type, + rhs.attnScale, rhs.isTraining, rhs.dropoutProbability, rhs.qkv_layout, + rhs.o_format, rhs.do_format, rhs.dqkv_layout, rhs.qkv_scale_inv_format, + rhs.do_scale_inv_format, rhs.mask_type, rhs.softmax_type, rhs.window_size_left, + rhs.window_size_right, rhs.bottom_right_diagonal, rhs.deterministic, + rhs.bias_type, rhs.qkv_tensor_type, rhs.o_tensor_type, rhs.do_tensor_type, rhs.dqkv_tensor_type, rhs.return_max_logit); } }; diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 8d9adeb62..912dc32d3 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -52,6 +52,8 @@ enum NVTE_QKV_Layout { NVTE_Paged_KV_SBHD_SBHD_SBHD = 22, /*!< Paged_KV_SBHD_SBHD_SBHD layout */ NVTE_Paged_KV_THD_BSHD_BSHD = 23, /*!< Paged_KV_THD_BSHD_BSHD layout */ NVTE_Paged_KV_THD_SBHD_SBHD = 24, /*!< Paged_KV_THD_SBHD_SBHD layout */ + NVTE_BHSD_BHSD_BHSD = 25, /*!< BHSD_BHSD_BHSD layout */ + NVTE_QKV_Layout_NOT_SET, /*!< Not set */ }; /*! \enum NVTE_QKV_Layout_Group @@ -70,6 +72,8 @@ enum NVTE_QKV_Layout_Group { NVTE_HD_HD_HD = 4, /*! Paged_KV_HD_HD_HD QKV layouts, e.g. Paged_KV_BSHD_BSHD_BSHD, Paged_KV_THD_SBHD_SBHD */ NVTE_Paged_KV_HD_HD_HD = 5, + /*! SD_SD_SD QKV layouts, e.g. BHSD_BHSD_BHSD */ + NVTE_SD_SD_SD = 6, }; /*! \enum NVTE_QKV_Format @@ -90,6 +94,10 @@ enum NVTE_QKV_Format { NVTE_THD_2BSHD = 5, /*! THD format for Q and SBHD format for KV, i.e. THD_SBHD_SBHD, Paged_KV_THD_SBHD_SBHD */ NVTE_THD_2SBHD = 6, + /*! BHSD QKV format, e.g. BHSD_BHSD_BHSD */ + NVTE_BHSD = 7, + /*! Not set */ + NVTE_QKV_Format_NOT_SET, }; /*! \enum NVTE_Bias_Type @@ -274,6 +282,9 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * \param[in] attn_scale Scaling factor for Q * K.T. * \param[in] dropout Dropout probability. * \param[in] qkv_layout QKV tensors' layout. + * \param[in] o_format Output format. + * \param[in] qkv_scale_inv_format Format of scale-inverse tensors for QKV; + * if NVTE_QKV_Format_NOT_SET, inferred from qkv_layout. * \param[in] bias_type Bias type. * \param[in] attn_mask_type Attention mask type. * \param[in] softmax_type Attention softmax type. @@ -292,7 +303,8 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso const NVTETensor page_table_v, const NVTETensor rng_state, size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, bool return_max_logit, bool cuda_graph, float attn_scale, float dropout, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, NVTETensor workspace, cudaStream_t stream); @@ -347,6 +359,13 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso * \param[in] attn_scale Scaling factor for Q * K.T. * \param[in] dropout Dropout probability. * \param[in] qkv_layout QKV tensors' layout. + * \param[in] o_format Output format. + * \param[in] do_format Output gradient's format. + * \param[in] dqkv_layout QKV gradient tensors' layout. + * \param[in] qkv_scale_inv_format Format of scale-inverse tensors for QKV; + * if NVTE_QKV_Format_NOT_SET, inferred from qkv_layout. + * \param[in] do_scale_inv_format Format of scale-inverse tensors for dO; + * if NVTE_QKV_Format_NOT_SET, inferred from the output layout. * \param[in] bias_type Bias type. * \param[in] attn_mask_type Attention mask type. * \param[in] softmax_type Attention softmax type. @@ -366,11 +385,13 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso const NVTETensor cu_seqlens_q_padded, const NVTETensor cu_seqlens_kv_padded, size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float dropout, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - int64_t window_size_left, int64_t window_size_right, - bool bottom_right_diagonal, bool deterministic, bool cuda_graph, - NVTETensor workspace, cudaStream_t stream); + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, + NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, + NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, + NVTE_Softmax_Type softmax_type, int64_t window_size_left, + int64_t window_size_right, bool bottom_right_diagonal, bool deterministic, + bool cuda_graph, NVTETensor workspace, cudaStream_t stream); /*! \brief Update the RNG state with the seed and calculated offset. * @@ -584,8 +605,81 @@ void nvte_prepare_flash_attn_fwd(NVTETensor qkvi, NVTETensor qkv, cudaStream_t s void nvte_prepare_flash_attn_bwd(NVTETensor q, NVTETensor k, NVTETensor v, NVTETensor qkv, cudaStream_t stream); +/*! \brief Transpose multiple tensors from BSHD/SBHD to BHSD. + * + * Each input tensor is 4D in BSHD or SBHD layout, and the corresponding output tensor + * is 4D in BHSD layout. Output tensors are pre-allocated and may have a larger last dimension. + * + * \param[in] inputs List of input tensors. + * \param[in,out] outputs List of output tensors. + * \param[in] num_tensors Number of tensors in the list. + * \param[in] original_format Original QKV format (NVTE_BSHD or NVTE_SBHD). + * \param[in] stream CUDA stream. + */ +void nvte_multi_tensor_transpose_to_bhsd(NVTETensor *inputs, NVTETensor *outputs, + size_t num_tensors, NVTE_QKV_Format original_format, + cudaStream_t stream); + +/*! \brief Pad the last dimension of multiple 2D tensors with zeros in one kernel launch. + * + * Each tensor copies a row-major (rows, in_cols) input to a (rows, out_cols) output, + * zero-filling the region [in_cols, out_cols) in every row. + * Outputs must be pre-allocated with out_cols >= in_cols and matching dtype. + * + * \param[in] inputs List of input tensors. + * \param[in,out] outputs List of output tensors. + * \param[in] num_tensors Number of tensors in the list. + * \param[in] stream CUDA stream. + */ +void nvte_multi_tensor_pad_last_dim(NVTETensor *inputs, NVTETensor *outputs, size_t num_tensors, + cudaStream_t stream); + #ifdef __cplusplus } // extern "C" -#endif + +#include +#include +#include + +/*! \brief Parses a QKV tensor shape into canonical (b, h, s, d, t) dimensions + * and converts between QKV formats. + */ +class AttentionShape { + public: + inline AttentionShape(NVTE_QKV_Format fmt, const size_t *shape) : canonical_{} { + auto [ndim, order] = dim_order(fmt); + for (size_t i = 0; i < ndim; ++i) canonical_[order[i]] = shape[i]; + } + + size_t b() const { return canonical_[0]; } + size_t h() const { return canonical_[1]; } + size_t s() const { return canonical_[2]; } + size_t d() const { return canonical_[3]; } + size_t t() const { return canonical_[4]; } + + inline void to_format(NVTE_QKV_Format dst_fmt, size_t *dst_shape) const { + auto [ndim, order] = dim_order(dst_fmt); + for (size_t i = 0; i < ndim; ++i) dst_shape[i] = canonical_[order[i]]; + } + + private: + static inline std::pair> dim_order(NVTE_QKV_Format fmt) { + switch (fmt) { + case NVTE_QKV_Format::NVTE_BSHD: + return {4, {0, 2, 1, 3}}; // b s h d + case NVTE_QKV_Format::NVTE_SBHD: + return {4, {2, 0, 1, 3}}; // s b h d + case NVTE_QKV_Format::NVTE_BHSD: + return {4, {0, 1, 2, 3}}; // b h s d + case NVTE_QKV_Format::NVTE_THD: + return {3, {4, 1, 3, -1}}; // t h d + default: + return {0, {}}; + } + } + size_t canonical_[5] = {}; +}; + +#endif // __cplusplus #endif diff --git a/transformer_engine/common/include/transformer_engine/swizzle.h b/transformer_engine/common/include/transformer_engine/swizzle.h index 4e28de3be..396093b54 100644 --- a/transformer_engine/common/include/transformer_engine/swizzle.h +++ b/transformer_engine/common/include/transformer_engine/swizzle.h @@ -32,10 +32,10 @@ void nvte_swizzle_scaling_factors(const NVTETensor input, NVTETensor output, cud /*! \brief Swizzling scaling factors into the required interleaved layout for GEMM * - * \param[in] inputs Input tensors with non-swizzled scale_inv. - * \param[in,out] outputs Output tensors which hosts swizzled scale_inv. - * \param[in] num_tensors Number of input and output tensors. - * \param[in] stream CUDA stream used for the operation. + * \param[in] inputs Input tensors with non-swizzled scale_inv. + * \param[in,out] outputs Output tensors which hosts swizzled scale_inv. + * \param[in] num_tensors Number of input and output tensors. + * \param[in] stream CUDA stream used for the operation. * * Requirements: * - scale_inv is stored in row-major. @@ -45,6 +45,17 @@ void nvte_swizzle_scaling_factors(const NVTETensor input, NVTETensor output, cud void nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETensor* outputs, const size_t num_tensors, cudaStream_t stream); +/*! \brief Same as nvte_multi_tensor_swizzle_scaling_factors, but skips + * scale_inv shape/padding validation. + * + * Use this variant when the data and scale_inv tensors intentionally have + * different shapes, e.g. when scale_invs have been transposed for attention. + */ +void nvte_multi_tensor_swizzle_scaling_factors_unchecked(const NVTETensor* inputs, + NVTETensor* outputs, + const size_t num_tensors, + cudaStream_t stream); + /*! \brief Unswizzling scaling factors from the interleaved layout used by GEMM back to row-major * * \param[in] input Input tensor with swizzled scale_inv. diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index 6c5977624..de4fdbb04 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -21,6 +21,17 @@ namespace { constexpr int MXFP8_BLOCK_SIZE = 32; constexpr int NVFP4_BLOCK_SIZE = 16; +int get_max_dynamic_smem() { + static int max_smem = -1; + if (max_smem < 0) { + int device; + NVTE_CHECK_CUDA(cudaGetDevice(&device)); + NVTE_CHECK_CUDA( + cudaDeviceGetAttribute(&max_smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); + } + return max_smem; +} + constexpr __device__ __host__ int TB_DIM = 32; constexpr __device__ __host__ int NEW_SF_TILE_DIM_K = 16; constexpr __device__ __host__ int N_SF_PER_TD_PER_TILE = 4; @@ -282,6 +293,171 @@ __device__ void swizzle_row_scaling_kernel_impl(const void* input, void* output, } } +template +__global__ void __launch_bounds__(TB_DIM* TB_DIM) + swizzle_row_scaling_kernel(const void* input, void* output, const int M, const int K, + const int original_M, const int original_K) { + swizzle_row_scaling_kernel_impl( + input, output, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y); +} + +// Narrow-K specialization for row scaling swizzle. +// When K is small (num_tiles_k < TB_DIM), the standard kernel wastes threadIdx.x +// because there aren't enough K-tiles to distribute across threads. +// This kernel repurposes the thread dimensions: threadIdx.x iterates rows within +// an M-tile, threadIdx.y indexes M-tiles within the block, processing TB_DIM +// M-tiles per block with full thread utilization. +template +__device__ void swizzle_row_scaling_narrow_k_kernel_impl(const void* input, void* output, + const int M, const int K, + const int original_M, const int original_K, + const int bid, const int grid_dim) { + constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4; + const int K_i32 = K / 4; + const int num_tiles_m = M / SF_TILE_DIM_M; + + const int m_tile = bid * blockDim.y + threadIdx.y; + const bool active = (m_tile < num_tiles_m); + + extern __shared__ int4 slm_v4i[]; + const int slm_tile_v4i = K_i32 * (SF_TILE_SIZE_I32 / 4); + + if (active) { + const bool padding_m = (m_tile == num_tiles_m - 1) && (original_M < M); + const bool padding_k = (original_K < K); + + int4* my_slm = slm_v4i + threadIdx.y * slm_tile_v4i; + + for (int k = 0; k < K_i32; k++) { + const int input_base = m_tile * SF_TILE_DIM_M * K_i32 + k; + const int* input_i32 = reinterpret_cast(input) + input_base; + + int regs[N_SF_PER_TD_PER_TILE]; +#pragma unroll + for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) { + const int row = i * TB_DIM + threadIdx.x; + regs[i] = __ldg(input_i32 + row * K_i32); + if (padding_m || padding_k) { + for (int j = 0; j < 4; j++) { + const int byte_row = m_tile * SF_TILE_DIM_M + row; + const int byte_col = k * 4 + j; + if (byte_row >= original_M || byte_col >= original_K) { + reinterpret_cast(®s[i])[j] = 0; + } + } + } + } + + my_slm[k * (SF_TILE_SIZE_I32 / 4) + threadIdx.x] = *reinterpret_cast(regs); + } + } + + __syncthreads(); + + if (active) { + int4* my_slm = slm_v4i + threadIdx.y * slm_tile_v4i; + int4* out_v4i = + reinterpret_cast(reinterpret_cast(output) + m_tile * SF_TILE_DIM_M * K_i32); + + for (int i = threadIdx.x; i < slm_tile_v4i; i += blockDim.x) { + out_v4i[i] = my_slm[i]; + } + } +} + +template +__global__ void __launch_bounds__(TB_DIM* TB_DIM) + swizzle_row_scaling_narrow_k_kernel(const void* input, void* output, const int M, const int K, + const int original_M, const int original_K) { + swizzle_row_scaling_narrow_k_kernel_impl( + input, output, M, K, original_M, original_K, blockIdx.x, gridDim.x); +} + +// Narrow-M variant of the column scaling swizzle kernel, for when num_tiles_m < TB_DIM. +// Analogous to the narrow-K row kernel: when the M dimension is small, the normal +// col kernel underutilizes threads in the load phase because threadIdx.x covers M +// positions with vectorized loads, leaving many threads idle. This kernel repurposes +// thread dimensions: threadIdx.y indexes K-tiles within the block, threadIdx.x covers +// one int32 column of an M-tile, and M-tiles are iterated serially. +template +__device__ void swizzle_col_scaling_narrow_m_kernel_impl(const void* input, void* output, + const int M, const int K, + const int original_M, const int original_K, + const int bid, const int grid_dim) { + constexpr int SF_TILE_SIZE_I32 = SF_TILE_DIM_M * SF_TILE_DIM_K / 4; + constexpr int SF_TILE_DIM_M_I32 = SF_TILE_DIM_M / 4; + constexpr int SF_TILE_DIM_K_I32 = SF_TILE_DIM_K; + + const int M_i32 = M / 4; + const int K_i32 = K; + const int num_tiles_m = M / SF_TILE_DIM_M; + const int num_tiles_k = K / SF_TILE_DIM_K; + + const int k_tile = bid * blockDim.y + threadIdx.y; + const bool active = (k_tile < num_tiles_k); + const int remaining = num_tiles_k - bid * static_cast(blockDim.y); + const int k_tiles_in_block = remaining <= 0 ? 0 : (remaining < TB_DIM ? remaining : TB_DIM); + + extern __shared__ int slm_narrow_m[]; + + if (active) { + const bool padding_k = (k_tile == num_tiles_k - 1) && (original_K < K); + const int32_t* input_i32 = reinterpret_cast(input); + + for (int m_tile = 0; m_tile < num_tiles_m; m_tile++) { + const bool padding_m = (m_tile == num_tiles_m - 1) && (original_M < M); + + int regs[N_SF_PER_TD_PER_TILE]; +#pragma unroll + for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) { + const int k_row = k_tile * SF_TILE_DIM_K_I32 + i; + const int m_col = m_tile * SF_TILE_DIM_M_I32 + threadIdx.x; + regs[i] = __ldg(input_i32 + k_row * M_i32 + m_col); + if (padding_m || padding_k) { + for (int j = 0; j < 4; j++) { + if (m_col * 4 + j >= original_M || k_row >= original_K) { + reinterpret_cast(®s[i])[j] = 0; + } + } + } + } + + regs_shuffle_with_bit_shifts(regs); + + int tM = threadIdx.x * N_SF_PER_TD_PER_TILE; + int* slm_tile = + slm_narrow_m + m_tile * TB_DIM * SF_TILE_SIZE_I32 + threadIdx.y * SF_TILE_SIZE_I32; +#pragma unroll + for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) { + slm_tile[(tM % SF_TILE_DIM_M) / NEW_SF_TILE_DIM_M_I32 + + ((tM + i) % NEW_SF_TILE_DIM_M_I32) * NEW_SF_TILE_DIM_K_I32] = regs[i]; + } + } + } + + __syncthreads(); + + const int linear_id = threadIdx.y * blockDim.x + threadIdx.x; + for (int m_tile = 0; m_tile < num_tiles_m; m_tile++) { + int4* out_v4i = reinterpret_cast(reinterpret_cast(output) + + m_tile * SF_TILE_DIM_M_I32 * K_i32 + + bid * TB_DIM * SF_TILE_SIZE_I32); + int4* slm_v4i = reinterpret_cast(slm_narrow_m + m_tile * TB_DIM * SF_TILE_SIZE_I32); + const int n_v4i = k_tiles_in_block * SF_TILE_SIZE_I32 / 4; + for (int j = linear_id; j < n_v4i; j += blockDim.x * blockDim.y) { + out_v4i[j] = slm_v4i[j]; + } + } +} + +template +__global__ void __launch_bounds__(TB_DIM* TB_DIM) + swizzle_col_scaling_narrow_m_kernel(const void* input, void* output, const int M, const int K, + const int original_M, const int original_K) { + swizzle_col_scaling_narrow_m_kernel_impl( + input, output, M, K, original_M, original_K, blockIdx.x, gridDim.x); +} + template __device__ void unswizzle_row_scaling_kernel_impl(const void* input, void* output, const int M, const int K, const int bid_x, const int bid_y, @@ -422,14 +598,6 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) } } -template -__global__ void __launch_bounds__(TB_DIM* TB_DIM) - swizzle_row_scaling_kernel(const void* input, void* output, const int M, const int K, - const int original_M, const int original_K) { - swizzle_row_scaling_kernel_impl( - input, output, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y); -} - constexpr int kMaxTensorsPerKernel = 64; // Args must be <4 KB struct MultiSwizzleArgs { // (input) Data buffers for input scaling factors @@ -617,6 +785,50 @@ __global__ void multi_tensor_swizzle_col_scaling_kernel(MultiSwizzleArgs kernel_ input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y); } +template +__global__ void __launch_bounds__(TB_DIM* TB_DIM) + multi_tensor_swizzle_row_scaling_narrow_k_kernel(MultiSwizzleArgs kernel_args) { + const int bid = blockIdx.x; + int tensor_id = 0; + while (kernel_args.block_range[tensor_id + 1] <= bid) { + ++tensor_id; + } + const void* input = kernel_args.input_list[tensor_id]; + void* output = kernel_args.output_list[tensor_id]; + const int M = kernel_args.m_list[tensor_id]; + const int K = kernel_args.k_list[tensor_id]; + const int original_M = kernel_args.original_m_list[tensor_id]; + const int original_K = kernel_args.original_k_list[tensor_id]; + const int flat_bid = bid - kernel_args.block_range[tensor_id]; + const int num_tiles_m = M / SF_TILE_DIM_M; + const int grid_dim = DIVUP(num_tiles_m, TB_DIM); + + swizzle_row_scaling_narrow_k_kernel_impl( + input, output, M, K, original_M, original_K, flat_bid, grid_dim); +} + +template +__global__ void __launch_bounds__(TB_DIM* TB_DIM) + multi_tensor_swizzle_col_scaling_narrow_m_kernel(MultiSwizzleArgs kernel_args) { + const int bid = blockIdx.x; + int tensor_id = 0; + while (kernel_args.block_range[tensor_id + 1] <= bid) { + ++tensor_id; + } + const void* input = kernel_args.input_list[tensor_id]; + void* output = kernel_args.output_list[tensor_id]; + const int M = kernel_args.m_list[tensor_id]; + const int K = kernel_args.k_list[tensor_id]; + const int original_M = kernel_args.original_m_list[tensor_id]; + const int original_K = kernel_args.original_k_list[tensor_id]; + const int flat_bid = bid - kernel_args.block_range[tensor_id]; + const int num_tiles_k = K / SF_TILE_DIM_K; + const int grid_dim = DIVUP(num_tiles_k, TB_DIM); + + swizzle_col_scaling_narrow_m_kernel_impl( + input, output, M, K, original_M, original_K, flat_bid, grid_dim); +} + } // namespace void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t stream) { @@ -737,13 +949,6 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s // Perform row-wise swizzle if (rowwise_swizzle) { - int vec_load_size = (num_tiles_k - 1) % 4 + 1; - /* there is no int3 and misaligned if using int4/int2 */ - if (vec_load_size == 3) vec_load_size = 1; - int n_tiles_in_tb = TB_DIM * vec_load_size; - dim3 num_blocks(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m); - int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); - int original_M{0}, original_K{0}; void *input_scale_inv_ptr{nullptr}, *output_scale_inv_ptr{nullptr}; switch (scaling_mode) { @@ -772,79 +977,114 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s NVTE_ERROR("Invalid scaling mode"); } - switch (vec_load_size) { - case 4: - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_row_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - swizzle_row_scaling_kernel - <<>>( - input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K); - break; - case 2: - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_row_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - swizzle_row_scaling_kernel - <<>>( - input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K); - break; - case 1: - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_row_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - swizzle_row_scaling_kernel - <<>>( - input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K); - break; - default: - NVTE_ERROR("Not valid vec_load_size."); - break; + const int narrow_k_slm_size = + TB_DIM * num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t)); + if (num_tiles_k < TB_DIM && narrow_k_slm_size <= get_max_dynamic_smem()) { + // Narrow-K: batch TB_DIM M-tiles per block, fully utilizing all threads. + dim3 num_blocks_narrow(DIVUP(num_tiles_m, TB_DIM)); + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(swizzle_row_scaling_narrow_k_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, narrow_k_slm_size)); + swizzle_row_scaling_narrow_k_kernel + <<>>( + input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K); + } else { + int vec_load_size = (num_tiles_k - 1) % 4 + 1; + /* there is no int3 and misaligned if using int4/int2 */ + if (vec_load_size == 3) vec_load_size = 1; + int n_tiles_in_tb = TB_DIM * vec_load_size; + dim3 num_blocks(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m); + int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); + + switch (vec_load_size) { + case 4: + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(swizzle_row_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + swizzle_row_scaling_kernel + <<>>( + input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K); + break; + case 2: + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(swizzle_row_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + swizzle_row_scaling_kernel + <<>>( + input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K); + break; + case 1: + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(swizzle_row_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + swizzle_row_scaling_kernel + <<>>( + input_scale_inv_ptr, output_scale_inv_ptr, m, k, original_M, original_K); + break; + default: + NVTE_ERROR("Not valid vec_load_size."); + break; + } } NVTE_CHECK_CUDA(cudaGetLastError()); } // Perform column-wise swizzle if (columnwise_swizzle) { - int vec_load_size = (num_tiles_m - 1) % 4 + 1; - if (vec_load_size == 3) vec_load_size = 1; /* no int3 and misaligned if using int4/int2 */ - int n_tiles_in_tb = TB_DIM * vec_load_size; - dim3 num_blocks(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size)); - int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); const int original_M = input->flat_last_dim(); const int original_K = input->flat_first_dim() / MXFP8_BLOCK_SIZE; - switch (vec_load_size) { - case 4: - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_col_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - swizzle_col_scaling_kernel - <<>>(input->columnwise_scale_inv.dptr, - output->columnwise_scale_inv.dptr, m, k, - original_M, original_K); - break; - case 2: - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_col_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - swizzle_col_scaling_kernel - <<>>(input->columnwise_scale_inv.dptr, - output->columnwise_scale_inv.dptr, m, k, - original_M, original_K); - break; - case 1: - NVTE_CHECK_CUDA( - cudaFuncSetAttribute(swizzle_col_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - swizzle_col_scaling_kernel - <<>>(input->columnwise_scale_inv.dptr, - output->columnwise_scale_inv.dptr, m, k, - original_M, original_K); - break; - default: - NVTE_ERROR("Not valid vec_load_size."); - break; + const int narrow_m_slm_size = + TB_DIM * num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t)); + if (num_tiles_m < TB_DIM && narrow_m_slm_size <= get_max_dynamic_smem()) { + // Narrow-M: batch TB_DIM K-tiles per block, fully utilizing all threads. + dim3 num_blocks_narrow(DIVUP(num_tiles_k, TB_DIM)); + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(swizzle_col_scaling_narrow_m_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, narrow_m_slm_size)); + swizzle_col_scaling_narrow_m_kernel + <<>>( + input->columnwise_scale_inv.dptr, output->columnwise_scale_inv.dptr, m, k, original_M, + original_K); + } else { + int vec_load_size = (num_tiles_m - 1) % 4 + 1; + if (vec_load_size == 3) vec_load_size = 1; /* no int3 and misaligned if using int4/int2 */ + int n_tiles_in_tb = TB_DIM * vec_load_size; + dim3 num_blocks(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size)); + int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); + + switch (vec_load_size) { + case 4: + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(swizzle_col_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + swizzle_col_scaling_kernel + <<>>(input->columnwise_scale_inv.dptr, + output->columnwise_scale_inv.dptr, m, + k, original_M, original_K); + break; + case 2: + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(swizzle_col_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + swizzle_col_scaling_kernel + <<>>(input->columnwise_scale_inv.dptr, + output->columnwise_scale_inv.dptr, m, + k, original_M, original_K); + break; + case 1: + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(swizzle_col_scaling_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + swizzle_col_scaling_kernel + <<>>(input->columnwise_scale_inv.dptr, + output->columnwise_scale_inv.dptr, m, + k, original_M, original_K); + break; + default: + NVTE_ERROR("Not valid vec_load_size."); + break; + } } NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -853,83 +1093,138 @@ void swizzle_scaling_factors(const Tensor* input, Tensor* output, cudaStream_t s template void launch_multi_tensor_swizzle_scaling_factors(MultiSwizzleArgs& kernel_args, const int vec_load_size, const bool is_rowwise, + const bool use_narrow_k, const bool use_narrow_m, cudaStream_t stream) { - int n_tiles_in_tb = TB_DIM * vec_load_size; - int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); - /* Calculate number of CUDA blocks needed for each tensor. - * We have to do it here because we have to iterate over all tensors in this batch to - * get the minimum vec_load_size. - */ - for (size_t j = 0; j < kernel_args.num_tensors; j++) { - const int m = kernel_args.m_list[j]; - const int k = kernel_args.k_list[j]; - int num_tiles_m = m / SF_TILE_DIM_M; - int num_tiles_k = k / SF_TILE_DIM_K; - if (is_rowwise) { - kernel_args.block_range[j + 1] = - kernel_args.block_range[j] + DIVUP(num_tiles_k, n_tiles_in_tb) * num_tiles_m; - } else { - kernel_args.block_range[j + 1] = - kernel_args.block_range[j] + - DIVUP(num_tiles_k, TB_DIM) * DIVUP(num_tiles_m, vec_load_size); + // cudaFuncSetAttribute is a host-synchronous driver call; cache the max shared memory + // setting per kernel variant so we only pay the cost when slm_size actually increases. + auto set_smem_if_needed = [](auto kernel_fn, int slm, int& cached) { + if (cached < slm) { + NVTE_CHECK_CUDA( + cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, slm)); + cached = slm; } - } - // Launch kernel - const int num_blocks = kernel_args.block_range[kernel_args.num_tensors]; + }; + dim3 block_size(TB_DIM, TB_DIM); - if (is_rowwise) { - switch (vec_load_size) { - case 4: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - multi_tensor_swizzle_row_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - multi_tensor_swizzle_row_scaling_kernel - <<>>(kernel_args); - break; - case 2: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - multi_tensor_swizzle_row_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - multi_tensor_swizzle_row_scaling_kernel - <<>>(kernel_args); - break; - case 1: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - multi_tensor_swizzle_row_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - multi_tensor_swizzle_row_scaling_kernel - <<>>(kernel_args); - break; - default: - NVTE_ERROR("Not valid vec_load_size."); - break; + + if (is_rowwise && use_narrow_k) { + // Narrow-K path: each block handles TB_DIM M-tiles with full thread utilization. + // slm_size depends on num_tiles_k, which can vary per tensor — use the max. + int max_num_tiles_k = 0; + for (size_t j = 0; j < kernel_args.num_tensors; j++) { + const int num_tiles_m = kernel_args.m_list[j] / SF_TILE_DIM_M; + const int num_tiles_k = kernel_args.k_list[j] / SF_TILE_DIM_K; + max_num_tiles_k = std::max(max_num_tiles_k, num_tiles_k); + kernel_args.block_range[j + 1] = kernel_args.block_range[j] + DIVUP(num_tiles_m, TB_DIM); + } + int slm_size = TB_DIM * max_num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); + const int num_blocks = kernel_args.block_range[kernel_args.num_tensors]; + + static int cached_narrow_k = -1; + set_smem_if_needed( + multi_tensor_swizzle_row_scaling_narrow_k_kernel, slm_size, + cached_narrow_k); + multi_tensor_swizzle_row_scaling_narrow_k_kernel + <<>>(kernel_args); + } else if (!is_rowwise && use_narrow_m) { + // Narrow-M path: each block handles TB_DIM K-tiles with full thread utilization. + // slm_size depends on num_tiles_m, which can vary per tensor — use the max. + int max_num_tiles_m = 0; + for (size_t j = 0; j < kernel_args.num_tensors; j++) { + const int num_tiles_m = kernel_args.m_list[j] / SF_TILE_DIM_M; + const int num_tiles_k = kernel_args.k_list[j] / SF_TILE_DIM_K; + max_num_tiles_m = std::max(max_num_tiles_m, num_tiles_m); + kernel_args.block_range[j + 1] = kernel_args.block_range[j] + DIVUP(num_tiles_k, TB_DIM); } + int slm_size = TB_DIM * max_num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); + const int num_blocks = kernel_args.block_range[kernel_args.num_tensors]; + + static int cached_narrow_m = -1; + set_smem_if_needed( + multi_tensor_swizzle_col_scaling_narrow_m_kernel, slm_size, + cached_narrow_m); + multi_tensor_swizzle_col_scaling_narrow_m_kernel + <<>>(kernel_args); } else { - switch (vec_load_size) { - case 4: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - multi_tensor_swizzle_col_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - multi_tensor_swizzle_col_scaling_kernel - <<>>(kernel_args); - break; - case 2: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - multi_tensor_swizzle_col_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - multi_tensor_swizzle_col_scaling_kernel - <<>>(kernel_args); - break; - case 1: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - multi_tensor_swizzle_col_scaling_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - multi_tensor_swizzle_col_scaling_kernel - <<>>(kernel_args); - break; - default: - NVTE_ERROR("Not valid vec_load_size."); - break; + int n_tiles_in_tb = TB_DIM * vec_load_size; + int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); + /* Calculate number of CUDA blocks needed for each tensor. + * We have to do it here because we have to iterate over all tensors in this batch to + * get the minimum vec_load_size. + */ + for (size_t j = 0; j < kernel_args.num_tensors; j++) { + const int m = kernel_args.m_list[j]; + const int k = kernel_args.k_list[j]; + int num_tiles_m = m / SF_TILE_DIM_M; + int num_tiles_k = k / SF_TILE_DIM_K; + if (is_rowwise) { + kernel_args.block_range[j + 1] = + kernel_args.block_range[j] + DIVUP(num_tiles_k, n_tiles_in_tb) * num_tiles_m; + } else { + kernel_args.block_range[j + 1] = + kernel_args.block_range[j] + + DIVUP(num_tiles_k, TB_DIM) * DIVUP(num_tiles_m, vec_load_size); + } + } + const int num_blocks = kernel_args.block_range[kernel_args.num_tensors]; + + static int cached_row_int4 = -1, cached_row_int2 = -1, cached_row_int1 = -1; + static int cached_col_int4 = -1, cached_col_int2 = -1, cached_col_int1 = -1; + + if (is_rowwise) { + switch (vec_load_size) { + case 4: + set_smem_if_needed( + multi_tensor_swizzle_row_scaling_kernel, slm_size, + cached_row_int4); + multi_tensor_swizzle_row_scaling_kernel + <<>>(kernel_args); + break; + case 2: + set_smem_if_needed( + multi_tensor_swizzle_row_scaling_kernel, slm_size, + cached_row_int2); + multi_tensor_swizzle_row_scaling_kernel + <<>>(kernel_args); + break; + case 1: + set_smem_if_needed( + multi_tensor_swizzle_row_scaling_kernel, slm_size, + cached_row_int1); + multi_tensor_swizzle_row_scaling_kernel + <<>>(kernel_args); + break; + default: + NVTE_ERROR("Not valid vec_load_size."); + break; + } + } else { + switch (vec_load_size) { + case 4: + set_smem_if_needed( + multi_tensor_swizzle_col_scaling_kernel, slm_size, + cached_col_int4); + multi_tensor_swizzle_col_scaling_kernel + <<>>(kernel_args); + break; + case 2: + set_smem_if_needed( + multi_tensor_swizzle_col_scaling_kernel, slm_size, + cached_col_int2); + multi_tensor_swizzle_col_scaling_kernel + <<>>(kernel_args); + break; + case 1: + set_smem_if_needed( + multi_tensor_swizzle_col_scaling_kernel, slm_size, + cached_col_int1); + multi_tensor_swizzle_col_scaling_kernel + <<>>(kernel_args); + break; + default: + NVTE_ERROR("Not valid vec_load_size."); + break; + } } } NVTE_CHECK_CUDA(cudaGetLastError()); @@ -1019,7 +1314,8 @@ void launch_multi_tensor_unswizzle_scaling_factors(MultiSwizzleArgs& kernel_args } void multi_tensor_swizzle_scaling_factors(const std::vector& input, - std::vector& output, cudaStream_t stream) { + std::vector& output, cudaStream_t stream, + bool check_scale_inv_shapes) { auto num_tensors = input.size(); bool all_has_data = true; bool all_has_columnwise_data = true; @@ -1038,8 +1334,10 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, // We don't allow empty tensors. They should be filtered out before calling this function. NVTE_CHECK(input[i]->numel() != 0, "Tensor input[", i, "] is empty."); - CheckInputTensor(*input[i], "scaling_factor_input[" + std::to_string(i) + "]"); - CheckInputTensor(*output[i], "scaling_factor_output[" + std::to_string(i) + "]"); + CheckInputTensor(*input[i], "scaling_factor_input[" + std::to_string(i) + "]", + check_scale_inv_shapes); + CheckInputTensor(*output[i], "scaling_factor_output[" + std::to_string(i) + "]", + check_scale_inv_shapes); all_has_data = all_has_data && input[i]->scale_inv.has_data(); all_has_columnwise_data = (all_has_columnwise_data && input[i]->columnwise_scale_inv.has_data()); @@ -1060,16 +1358,18 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, kernel_args.num_tensors = 0; kernel_args.block_range[0] = 0; int vec_load_size = 4; + bool all_narrow_k = true; for (size_t i = 0; i < num_tensors; i++) { //Launch kernel if argument struct is full if (kernel_args.num_tensors == kMaxTensorsPerKernel) { // There is no int3 and misaligned if using int4/int2. if (vec_load_size == 3) vec_load_size = 1; launch_multi_tensor_swizzle_scaling_factors( - kernel_args, vec_load_size, true, stream); + kernel_args, vec_load_size, true, all_narrow_k, false, stream); // Reset the argument struct and vec_load_size kernel_args.num_tensors = 0; vec_load_size = 4; + all_narrow_k = true; } int m, k; @@ -1103,6 +1403,10 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, } int num_tiles_k = k / SF_TILE_DIM_K; + const int narrow_k_slm = + TB_DIM * num_tiles_k * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t)); + all_narrow_k = + all_narrow_k && (num_tiles_k < TB_DIM) && (narrow_k_slm <= get_max_dynamic_smem()); int vec_load_size_i = (num_tiles_k - 1) % 4 + 1; // We use the minimum vec_load_size across all tensors. // TODO(zhongbo): fix vec_load_size for NVFP4 @@ -1132,7 +1436,7 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, // There is no int3 and misaligned if using int4/int2. if (vec_load_size == 3) vec_load_size = 1; launch_multi_tensor_swizzle_scaling_factors( - kernel_args, vec_load_size, true, stream); + kernel_args, vec_load_size, true, all_narrow_k, false, stream); } if (columnwise_swizzle) { @@ -1143,16 +1447,18 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, kernel_args.num_tensors = 0; kernel_args.block_range[0] = 0; int vec_load_size = 4; + bool all_narrow_m = true; for (size_t i = 0; i < num_tensors; i++) { //Launch kernel if argument struct is full if (kernel_args.num_tensors == kMaxTensorsPerKernel) { // There is no int3 and misaligned if using int4/int2. if (vec_load_size == 3) vec_load_size = 1; launch_multi_tensor_swizzle_scaling_factors( - kernel_args, vec_load_size, false, stream); + kernel_args, vec_load_size, false, false, all_narrow_m, stream); // Reset the argument struct and vec_load_size kernel_args.num_tensors = 0; vec_load_size = 4; + all_narrow_m = true; } const int m = input[i]->columnwise_scale_inv.shape[1]; const int k = input[i]->columnwise_scale_inv.shape[0]; @@ -1166,7 +1472,12 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, "Input.columnwise_scale_inv size is not equal to " "Output.columnwise_scale_inv size!"); + int num_tiles_m = m / SF_TILE_DIM_M; int num_tiles_k = k / SF_TILE_DIM_K; + const int narrow_m_slm = + TB_DIM * num_tiles_m * SF_TILE_DIM_M * SF_TILE_DIM_K * static_cast(sizeof(int8_t)); + all_narrow_m = + all_narrow_m && (num_tiles_m < TB_DIM) && (narrow_m_slm <= get_max_dynamic_smem()); int vec_load_size_i = (num_tiles_k - 1) % 4 + 1; // We use the minimum vec_load_size across all tensors. vec_load_size = std::min(vec_load_size, vec_load_size_i); @@ -1184,7 +1495,7 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, // There is no int3 and misaligned if using int4/int2. if (vec_load_size == 3) vec_load_size = 1; launch_multi_tensor_swizzle_scaling_factors( - kernel_args, vec_load_size, false, stream); + kernel_args, vec_load_size, false, false, all_narrow_m, stream); } } @@ -1529,7 +1840,24 @@ void nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETen input_list.push_back(convertNVTETensorCheck(inputs[i])); output_list.push_back(convertNVTETensorCheck(outputs[i])); } - multi_tensor_swizzle_scaling_factors(input_list, output_list, stream); + multi_tensor_swizzle_scaling_factors(input_list, output_list, stream, + /*check_scale_inv_shapes=*/true); +} + +void nvte_multi_tensor_swizzle_scaling_factors_unchecked(const NVTETensor* inputs, + NVTETensor* outputs, + const size_t num_tensors, + cudaStream_t stream) { + NVTE_API_CALL(nvte_multi_tensor_swizzle_scaling_factors_unchecked); + using namespace transformer_engine; + NVTE_CHECK(num_tensors > 0, "Number of tensors should be greater than 0."); + std::vector input_list, output_list; + for (size_t i = 0; i < num_tensors; i++) { + input_list.push_back(convertNVTETensorCheck(inputs[i])); + output_list.push_back(convertNVTETensorCheck(outputs[i])); + } + multi_tensor_swizzle_scaling_factors(input_list, output_list, stream, + /*check_scale_inv_shapes=*/false); } void nvte_unswizzle_scaling_factors(const NVTETensor input, NVTETensor output, diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index eacd10eb3..1261879a8 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -120,7 +120,7 @@ void CheckScaleTensorShape(const Tensor &t, const std::string &name) { const auto &expected = std::vector{expected_x, expected_y}; NVTE_CHECK(t.columnwise_scale_inv.shape == expected, "Tensor \"", name, - "\" has invalid columnwise_scale_inv shape (expected ", expected, ", got ", + "\" has invalid columnwise_scale_inv shape (expected ", expected, ", got ", t.columnwise_scale_inv.shape, ")"); } } else if (t.scaling_mode == NVTE_NVFP4_1D_SCALING) { @@ -144,7 +144,7 @@ void CheckScaleTensorShape(const Tensor &t, const std::string &name) { } } -void CheckInputTensor(const Tensor &t, const std::string &name) { +void CheckInputTensor(const Tensor &t, const std::string &name, bool check_scale_inv_shapes) { const DType type = t.dtype(); if (is_fp8_dtype(type)) { // FP8 input needs to have scale_inv @@ -195,7 +195,9 @@ void CheckInputTensor(const Tensor &t, const std::string &name) { } NVTE_CHECK(t.has_data() || t.has_columnwise_data(), "Input ", name, " is not allocated!"); - CheckScaleTensorShape(t, name); + if (check_scale_inv_shapes) { + CheckScaleTensorShape(t, name); + } } void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empty) { diff --git a/transformer_engine/common/util/pybind_helper.h b/transformer_engine/common/util/pybind_helper.h index 6adba23a8..fdfa47da8 100644 --- a/transformer_engine/common/util/pybind_helper.h +++ b/transformer_engine/common/util/pybind_helper.h @@ -48,7 +48,9 @@ .value("NVTE_SBHD_2BSHD", NVTE_QKV_Format::NVTE_SBHD_2BSHD) \ .value("NVTE_BSHD_2SBHD", NVTE_QKV_Format::NVTE_BSHD_2SBHD) \ .value("NVTE_THD_2BSHD", NVTE_QKV_Format::NVTE_THD_2BSHD) \ - .value("NVTE_THD_2SBHD", NVTE_QKV_Format::NVTE_THD_2SBHD); \ + .value("NVTE_THD_2SBHD", NVTE_QKV_Format::NVTE_THD_2SBHD) \ + .value("NVTE_BHSD", NVTE_QKV_Format::NVTE_BHSD) \ + .value("NVTE_QKV_Format_NOT_SET", NVTE_QKV_Format::NVTE_QKV_Format_NOT_SET); \ pybind11::enum_(m, "NVTE_QKV_Layout", pybind11::module_local()) \ .value("NVTE_SB3HD", NVTE_QKV_Layout::NVTE_SB3HD) \ .value("NVTE_SBH3D", NVTE_QKV_Layout::NVTE_SBH3D) \ @@ -74,7 +76,8 @@ .value("NVTE_Paged_KV_SBHD_BSHD_BSHD", NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_BSHD_BSHD) \ .value("NVTE_Paged_KV_SBHD_SBHD_SBHD", NVTE_QKV_Layout::NVTE_Paged_KV_SBHD_SBHD_SBHD) \ .value("NVTE_Paged_KV_THD_BSHD_BSHD", NVTE_QKV_Layout::NVTE_Paged_KV_THD_BSHD_BSHD) \ - .value("NVTE_Paged_KV_THD_SBHD_SBHD", NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD); \ + .value("NVTE_Paged_KV_THD_SBHD_SBHD", NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD) \ + .value("NVTE_BHSD_BHSD_BHSD", NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD); \ pybind11::enum_(m, "NVTE_Fused_Attn_Backend", pybind11::module_local()) \ .value("NVTE_F16_max512_seqlen", NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) \ .value("NVTE_F16_arbitrary_seqlen", NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) \ diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 92e67ac19..76f2d9289 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -145,19 +145,28 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, DType dtype, bool is_training, size_t max_segments_per_seq, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal) { - auto q_shape = std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim}; + auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; + auto q_shape = is_ragged + ? std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim} + : std::vector{input_batch, q_max_seqlen, attn_heads, qk_head_dim}; auto q_tensor = TensorWrapper(nullptr, q_shape, dtype); - auto k_shape = std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim}; + auto k_shape = is_ragged + ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim} + : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, qk_head_dim}; auto k_tensor = TensorWrapper(nullptr, k_shape, dtype); - auto v_shape = std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim}; + auto v_shape = is_ragged + ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim} + : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, v_head_dim}; auto v_tensor = TensorWrapper(nullptr, v_shape, dtype); + auto o_shape = is_ragged ? std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim} + : std::vector{input_batch, q_max_seqlen, attn_heads, v_head_dim}; auto bias_shape = std::vector{bias_batch, bias_heads, q_max_seqlen, kv_max_seqlen}; auto bias_tensor = TensorWrapper(nullptr, bias_shape, dtype); // F16 doesn't use this tensor auto s_tensor = TensorWrapper(nullptr, std::vector{1}, dtype); - auto o_tensor = TensorWrapper(nullptr, q_shape, dtype); + auto o_tensor = TensorWrapper(nullptr, o_shape, dtype); auto dummy_rng_state_tensor = TensorWrapper(nullptr, std::vector{2}, DType::kInt64); auto dummy_page_table_tensor = TensorWrapper(nullptr, std::vector{1}, DType::kInt32); @@ -168,7 +177,6 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( nvte_tensor_pack_create(&aux_output_tensors); TensorWrapper query_workspace_tensor; - auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; // It is a WAR to pre-create all possible cuDNN graph at the JIT compile time size_t max_num_segments = is_ragged ? input_batch * max_segments_per_seq : input_batch; size_t min_num_segments = input_batch; @@ -191,9 +199,9 @@ pybind11::tuple GetFusedAttnForwardWorkspaceSizes( q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), ragged_offset_tensor.data(), ragged_offset_tensor.data(), dummy_page_table_tensor.data(), dummy_page_table_tensor.data(), dummy_rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, false, false, - scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, query_workspace_tensor.data(), - nullptr); + scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), + NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, window_size_left, + window_size_right, bottom_right_diagonal, query_workspace_tensor.data(), nullptr); } nvte_tensor_pack_destroy(&aux_output_tensors); @@ -257,7 +265,8 @@ static void FusedAttnForwardImpl( /* Output tensors */ auto s_tensor = TensorWrapper(nullptr, std::vector{1}, dtype); // not used in F16 - auto o_shape = std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim}; + auto o_shape = is_ragged ? std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim} + : std::vector{input_batch, q_max_seqlen, attn_heads, v_head_dim}; auto o_tensor = TensorWrapper(output, o_shape, dtype); /* Prepare RNG state */ @@ -285,9 +294,15 @@ static void FusedAttnForwardImpl( void *q_ptr = q; void *k_ptr = k; void *v_ptr = v; - auto q_shape = std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim}; - auto k_shape = std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim}; - auto v_shape = std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim}; + auto q_shape = is_ragged + ? std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim} + : std::vector{input_batch, q_max_seqlen, attn_heads, qk_head_dim}; + auto k_shape = is_ragged + ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim} + : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, qk_head_dim}; + auto v_shape = is_ragged + ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim} + : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, v_head_dim}; if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { // QKV packed in q: [batch*seqlen, 3, heads, dim] @@ -328,8 +343,9 @@ static void FusedAttnForwardImpl( q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), q_seq_offsets_tensor.data(), k_seq_offsets_tensor.data(), dummy_page_table_tensor.data(), dummy_page_table_tensor.data(), rng_state_tensor.data(), q_max_seqlen, kv_max_seqlen, is_training, false, false, - scaling_factor, dropout_probability, qkv_layout, bias_type, mask_type, softmax_type, - window_size_left, window_size_right, bottom_right_diagonal, workspace_tensor.data(), stream); + scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), + NVTE_QKV_Format_NOT_SET, bias_type, mask_type, softmax_type, window_size_left, + window_size_right, bottom_right_diagonal, workspace_tensor.data(), stream); nvte_tensor_pack_destroy(&aux_output_tensors); } @@ -418,17 +434,26 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, NVTE_QKV_Layout qkv_layout, DType dtype, bool is_training, bool deterministic, size_t max_segments_per_seq, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal) { - auto q_shape = std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim}; + auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; + auto q_shape = is_ragged + ? std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim} + : std::vector{input_batch, q_max_seqlen, attn_heads, qk_head_dim}; auto q_tensor = TensorWrapper(nullptr, q_shape, dtype); auto dq_tensor = TensorWrapper(nullptr, q_shape, dtype); - auto k_shape = std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim}; + auto k_shape = is_ragged + ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim} + : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, qk_head_dim}; auto k_tensor = TensorWrapper(nullptr, k_shape, dtype); auto dk_tensor = TensorWrapper(nullptr, k_shape, dtype); - auto v_shape = std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim}; + auto v_shape = is_ragged + ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim} + : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, v_head_dim}; auto v_tensor = TensorWrapper(nullptr, v_shape, dtype); auto dv_tensor = TensorWrapper(nullptr, v_shape, dtype); - auto output_shape = std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim}; + auto output_shape = is_ragged + ? std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim} + : std::vector{input_batch, q_max_seqlen, attn_heads, v_head_dim}; auto doutput_tensor = TensorWrapper(nullptr, output_shape, dtype); auto output_tensor = TensorWrapper(nullptr, output_shape, dtype); @@ -443,7 +468,6 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( TensorWrapper query_workspace_tensor; - auto is_ragged = nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD; // It is a WAR to pre-create all possible cuDNN graph at the JIT compile time size_t max_num_segments = is_ragged ? input_batch * max_segments_per_seq : input_batch; size_t min_num_segments = input_batch; @@ -469,18 +493,19 @@ pybind11::tuple GetFusedAttnBackwardWorkspaceSizes( auto dummy_ragged_offset_tensor = TensorWrapper(nullptr, std::vector{num_segments + 1}, DType::kInt32); - nvte_fused_attn_bwd(q_tensor.data(), k_tensor.data(), v_tensor.data(), output_tensor.data(), - doutput_tensor.data(), - s_tensor.data(), // not used for F16 - s_tensor.data(), // not used for F16 - &aux_input_tensors, dq_tensor.data(), dk_tensor.data(), dv_tensor.data(), - dbias_tensor.data(), dummy_d_softmax_offset_tensor.data(), - q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), - dummy_ragged_offset_tensor.data(), dummy_ragged_offset_tensor.data(), - q_max_seqlen, kv_max_seqlen, scaling_factor, dropout_probability, - qkv_layout, bias_type, mask_type, softmax_type, window_size_left, - window_size_right, bottom_right_diagonal, deterministic, false, - query_workspace_tensor.data(), nullptr); + nvte_fused_attn_bwd( + q_tensor.data(), k_tensor.data(), v_tensor.data(), output_tensor.data(), + doutput_tensor.data(), + s_tensor.data(), // not used for F16 + s_tensor.data(), // not used for F16 + &aux_input_tensors, dq_tensor.data(), dk_tensor.data(), dv_tensor.data(), + dbias_tensor.data(), dummy_d_softmax_offset_tensor.data(), q_cu_seqlens_tensor.data(), + kv_cu_seqlens_tensor.data(), dummy_ragged_offset_tensor.data(), + dummy_ragged_offset_tensor.data(), q_max_seqlen, kv_max_seqlen, scaling_factor, + dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), + nvte_get_q_format(qkv_layout), qkv_layout, NVTE_QKV_Format_NOT_SET, NVTE_QKV_Format_NOT_SET, + bias_type, mask_type, softmax_type, window_size_left, window_size_right, + bottom_right_diagonal, deterministic, false, query_workspace_tensor.data(), nullptr); } nvte_tensor_pack_destroy(&aux_input_tensors); @@ -503,7 +528,9 @@ static void FusedAttnBackwardImpl( FUSED_ATTN_IMPL_COMMON_BLOCK; /* Input tensors */ - auto output_shape = std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim}; + auto output_shape = is_ragged + ? std::vector{input_batch * q_max_seqlen, attn_heads, v_head_dim} + : std::vector{input_batch, q_max_seqlen, attn_heads, v_head_dim}; auto output_tensor = TensorWrapper(output, output_shape, dtype); auto doutput_tensor = TensorWrapper(doutput, output_shape, dtype); @@ -530,7 +557,7 @@ static void FusedAttnBackwardImpl( bias_heads, q_max_seqlen, kv_max_seqlen, dtype, backend, softmax_aux, rng_state, bias, softmax_offset); - /* Call the underly NVTE API */ + /* Call the underlying NVTE API */ // Prepare Q, K, V pointers and shapes based on layout void *q_ptr = q; void *k_ptr = k; @@ -538,9 +565,15 @@ static void FusedAttnBackwardImpl( void *dq_ptr = dq; void *dk_ptr = dk; void *dv_ptr = dv; - auto q_shape = std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim}; - auto k_shape = std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim}; - auto v_shape = std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim}; + auto q_shape = is_ragged + ? std::vector{input_batch * q_max_seqlen, attn_heads, qk_head_dim} + : std::vector{input_batch, q_max_seqlen, attn_heads, qk_head_dim}; + auto k_shape = is_ragged + ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, qk_head_dim} + : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, qk_head_dim}; + auto v_shape = is_ragged + ? std::vector{input_batch * kv_max_seqlen, num_gqa_groups, v_head_dim} + : std::vector{input_batch, kv_max_seqlen, num_gqa_groups, v_head_dim}; if (layout_group == NVTE_QKV_Layout_Group::NVTE_3HD) { // QKV packed in q: [batch*seqlen, 3, heads, dim] @@ -596,17 +629,18 @@ static void FusedAttnBackwardImpl( } } - nvte_fused_attn_bwd(q_tensor.data(), k_tensor.data(), v_tensor.data(), output_tensor.data(), - doutput_tensor.data(), - s_tensor.data(), // not used for F16 - s_tensor.data(), // not used for F16 - &aux_input_tensors, dq_tensor.data(), dk_tensor.data(), dv_tensor.data(), - dbias_tensor.data(), dsoftmax_offset_tensor.data(), - q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), - q_seq_offsets_tensor.data(), k_seq_offsets_tensor.data(), q_max_seqlen, - kv_max_seqlen, scaling_factor, dropout_probability, qkv_layout, bias_type, - mask_type, softmax_type, window_size_left, window_size_right, - bottom_right_diagonal, deterministic, false, workspace_tensor.data(), stream); + nvte_fused_attn_bwd( + q_tensor.data(), k_tensor.data(), v_tensor.data(), output_tensor.data(), + doutput_tensor.data(), + s_tensor.data(), // not used for F16 + s_tensor.data(), // not used for F16 + &aux_input_tensors, dq_tensor.data(), dk_tensor.data(), dv_tensor.data(), dbias_tensor.data(), + dsoftmax_offset_tensor.data(), q_cu_seqlens_tensor.data(), kv_cu_seqlens_tensor.data(), + q_seq_offsets_tensor.data(), k_seq_offsets_tensor.data(), q_max_seqlen, kv_max_seqlen, + scaling_factor, dropout_probability, qkv_layout, nvte_get_q_format(qkv_layout), + nvte_get_q_format(qkv_layout), qkv_layout, NVTE_QKV_Format_NOT_SET, NVTE_QKV_Format_NOT_SET, + bias_type, mask_type, softmax_type, window_size_left, window_size_right, + bottom_right_diagonal, deterministic, false, workspace_tensor.data(), stream); nvte_tensor_pack_destroy(&aux_input_tensors); } diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index ecf3af2bf..60a6f655b 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -29,6 +29,7 @@ Float8Quantizer, Float8CurrentScalingQuantizer, ) +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer from transformer_engine.pytorch.quantized_tensor import ( QuantizedTensorStorage, prepare_for_saving, @@ -36,7 +37,6 @@ ) from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.constants import ( - TE_DType, QKVLayouts, dist_group_type, ) @@ -72,6 +72,7 @@ print_quantizers, ConvertTHDtoBSHD, ConvertBSHDtoTHD, + mxfp8_quantize_fast_path, ) from transformer_engine.pytorch.attention.dot_product_attention.utils import ( AttentionLogging as attn_log, @@ -193,15 +194,27 @@ def forward(ctx, tensor1, tensor2, tensor3, quantizer, quantizer_name, qkv_layou query_layer, key_layer, value_layer = [ x.contiguous() for x in [tensor1, tensor2, tensor3] ] - q_fp8, k_fp8, v_fp8 = combine_and_quantize( - qkv_layout, query_layer, key_layer, value_layer, quantizer + # always in sbhd_sbhd_sbhd shape at this point + q_fp8, k_fp8, v_fp8, qkv_layout, _ = combine_and_quantize( + qkv_layout, + query_layer, + key_layer, + value_layer, + quantizer, + keep_same_data_and_scale_inv_format=True, ) tensors = combine_and_dequantize( qkv_layout, q_fp8, k_fp8, v_fp8, src_nominal_dtype=query_layer.dtype ) + if isinstance(quantizer, MXFP8Quantizer): + # bhsd_bhsd_bhsd after combine_and_quantize; permute back to sbhd_sbhd_sbhd + tensors = [x.permute(2, 0, 1, 3).contiguous() for x in tensors] elif quantizer_name in ["S_quantizer", "O_quantizer"]: - t_fp8 = quantizer(tensor1) - tensors = (t_fp8.dequantize(dtype=tensor1.dtype), tensor2, tensor3) + if quantizer is not None: + t_fp8 = quantizer(tensor1) + tensors = (t_fp8.dequantize(dtype=tensor1.dtype), tensor2, tensor3) + else: + tensors = (tensor1, tensor2, tensor3) else: tensors = (tensor1, tensor2, tensor3) ctx.quantizer = quantizer @@ -213,16 +226,28 @@ def forward(ctx, tensor1, tensor2, tensor3, quantizer, quantizer_name, qkv_layou def backward(ctx, grad1, grad2, grad3): # pylint: disable=missing-function-docstring if ctx.quantizer_name in ["dO_quantizer", "dP_quantizer"]: - dt_fp8 = ctx.quantizer(grad1) - tensors = dt_fp8.dequantize(dtype=grad1.dtype), grad2, grad3 + if ctx.quantizer is not None: + dt_fp8 = ctx.quantizer(grad1) + tensors = dt_fp8.dequantize(dtype=grad1.dtype), grad2, grad3 + else: + tensors = grad1, grad2, grad3 elif ctx.quantizer_name == "dQKV_quantizer": query_grad, key_grad, value_grad = [x.contiguous() for x in [grad1, grad2, grad3]] - dq_fp8, dk_fp8, dv_fp8 = combine_and_quantize( - ctx.qkv_layout, query_grad, key_grad, value_grad, ctx.quantizer + # always in sbhd_sbhd_sbhd shape at this point + dq_fp8, dk_fp8, dv_fp8, new_qkv_layout, _ = combine_and_quantize( + ctx.qkv_layout, + query_grad, + key_grad, + value_grad, + ctx.quantizer, + keep_same_data_and_scale_inv_format=True, ) tensors = combine_and_dequantize( - ctx.qkv_layout, dq_fp8, dk_fp8, dv_fp8, src_nominal_dtype=query_grad.dtype + new_qkv_layout, dq_fp8, dk_fp8, dv_fp8, src_nominal_dtype=query_grad.dtype ) + if isinstance(ctx.quantizer, MXFP8Quantizer): + # bhsd_bhsd_bhsd after combine_and_quantize; permute back to sbhd_sbhd_sbhd + tensors = [x.permute(2, 0, 1, 3).contiguous() for x in tensors] else: tensors = grad1, grad2, grad3 return tensors[0], tensors[1], tensors[2], None, None, None @@ -425,10 +450,9 @@ def forward( ) ) - batch_size, seqlen = query_layer.shape[1], query_layer.shape[0] apply_qk_layer_scaling = self.apply_qk_layer_scaling and key_layer.dtype == torch.float16 - # [b, np, sq, sk] + # [b, h, sq, sk] output_size = ( query_layer.size(1), query_layer.size(2), @@ -447,12 +471,7 @@ def forward( int(query_layer.shape[2] / value_layer.shape[2]), dim=2 ) - # [sq, b, np, hn] -> [sq, b * np, hn] - query_layer = query_layer.reshape(output_size[2], output_size[0] * output_size[1], -1) - # [sk, b, np, hn] -> [sk, b * np, hn] - key_layer = key_layer.reshape(output_size[3], output_size[0] * output_size[1], -1) - - # preallocting result tensor: [b * np, sq, sk] + # preallocting result tensor: [b * h, sq, sk] matmul_result = torch.empty( output_size[0] * output_size[1], output_size[2], @@ -466,14 +485,15 @@ def forward( scale /= self.layer_number if fp8: + # get fp8 recipe for DPA + fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() + if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None: + fp8_recipe = fp8_meta["local_recipes"][0] # get quantizers from DPA; all Nones if not fp8 QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer = ( - dpa_utils.get_attention_quantizers(fp8, quantizers) + dpa_utils.get_attention_quantizers(fp8, fp8_recipe, quantizers) ) # S/dP are forced to use DS quantizers in DPA.init_fp8_metadata; revert them here for true CS emulation - fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() - if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None: - fp8_recipe = fp8_meta["local_recipes"][0] if fp8_recipe.float8_current_scaling(): S_quantizer = Float8CurrentScalingQuantizer( fp8_dtype=S_quantizer.dtype, device="cuda" @@ -481,25 +501,50 @@ def forward( dP_quantizer = Float8CurrentScalingQuantizer( fp8_dtype=dP_quantizer.dtype, device="cuda" ) + # disable swizzle for MXFP8Quantizer + for quantizer in [ + QKV_quantizer, + O_quantizer, + S_quantizer, + dQKV_quantizer, + dO_quantizer, + dP_quantizer, + ]: + if isinstance(quantizer, MXFP8Quantizer): + quantizer.optimize_for_gemm = False + quantizer.internal = False - if "2" in qkv_layout or "3" in qkv_layout: - qkv_format, *_ = dpa_utils.get_qkv_format(qkv_layout) - qkv_layout = "_".join([qkv_format] * 3) + # q, k, v are in sbhd after previous reshaping # quantize and dequantize QKV to emulate FP8 query_layer, key_layer, value_layer = FP8EmulationFunc.apply( - query_layer, key_layer, value_layer, QKV_quantizer, "QKV_quantizer", qkv_layout + query_layer, + key_layer, + value_layer, + QKV_quantizer, + "QKV_quantizer", + "sbhd_sbhd_sbhd", ) # quantize and dequantize dQKV to emulate FP8 query_layer, key_layer, value_layer = FP8EmulationFunc.apply( - query_layer, key_layer, value_layer, dQKV_quantizer, "dQKV_quantizer", qkv_layout + query_layer, + key_layer, + value_layer, + dQKV_quantizer, + "dQKV_quantizer", + "sbhd_sbhd_sbhd", ) - # Raw attention scores. [b * np, sq, sk] + # [sq, b, h, d] -> [sq, b * h, d] + query_layer = query_layer.reshape(output_size[2], output_size[0] * output_size[1], -1) + # [sk, b, h, d] -> [sk, b * h, d] + key_layer = key_layer.reshape(output_size[3], output_size[0] * output_size[1], -1) + + # Raw attention scores. [b * h, sq, sk] if core_attention_bias_type == "no_bias": matmul_result = torch.baddbmm( matmul_result, - query_layer.transpose(0, 1), # [b * np, sq, hn] - key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk] + query_layer.transpose(0, 1), # [b * h, sq, d] + key_layer.transpose(0, 1).transpose(1, 2), # [b * h, d, sk] beta=0.0, alpha=scale, ).view(*output_size) @@ -507,8 +552,8 @@ def forward( elif core_attention_bias_type == "pre_scale_bias": assert core_attention_bias is not None, "core_attention_bias should not be None!" matmul_result = torch.bmm( - query_layer.transpose(0, 1), # [b * np, sq, hn] - key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk] + query_layer.transpose(0, 1), # [b * h, sq, d] + key_layer.transpose(0, 1).transpose(1, 2), # [b * h, d, sk] ) matmul_result = matmul_result.view(*output_size) + core_attention_bias matmul_result *= scale @@ -533,8 +578,8 @@ def forward( ) matmul_result = torch.baddbmm( matmul_result, - query_layer.transpose(0, 1), # [b * np, sq, hn] - key_layer.transpose(0, 1).transpose(1, 2), # [b * np, hn, sk] + query_layer.transpose(0, 1), # [b * h, sq, d] + key_layer.transpose(0, 1).transpose(1, 2), # [b * h, d, sk] beta=0.0, alpha=scale, ) @@ -551,13 +596,13 @@ def forward( # max attention score max_logit = None if self.return_max_logit: - # matmul_result [b, np, sq, dk], max_logit [np] + # matmul_result [b, h, sq, dk], max_logit [h] max_logit = matmul_result if attn_mask_type != "no_mask": max_logit = self.mask_func(matmul_result, attention_mask) max_logit = torch.amax(max_logit, dim=(0, 2, 3)) - # add attention sink to the last column: [b, np, sq, sk+1] + # add attention sink to the last column: [b, h, sq, sk+1] if self.softmax_type != "vanilla": matmul_result = torch.cat( [ @@ -582,7 +627,7 @@ def forward( if "padding" in attn_mask_type: attention_probs = attention_probs.masked_fill(attention_mask, 0) - # remove attention sink: [b, np, sq, sk] + # remove attention sink: [b, h, sq, sk] if self.softmax_type != "vanilla": attention_probs = attention_probs[..., :-1] @@ -592,7 +637,7 @@ def forward( attention_probs = self.attention_dropout(attention_probs) # value_layer -> context layer. - # [sk, b, np, hn] --> [b, np, sq, hn] + # [sk, b, h, d] --> [b, h, sq, d] output_size = ( value_layer.size(1), value_layer.size(2), @@ -600,10 +645,10 @@ def forward( value_layer.size(3), ) - # change view [sk, b * np, hn] + # change view [sk, b * h, d] value_layer = value_layer.reshape(value_layer.size(0), output_size[0] * output_size[1], -1) - # change view [b * np, sq, sk] + # change view [b * h, sq, sk] attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1) if fp8: @@ -612,37 +657,37 @@ def forward( attention_probs, None, None, S_quantizer, "S_quantizer", None ) - # matmul: [b * np, sq, hn] + # matmul: [b * h, sq, d] context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1)) - # change view [b, np, sq, hn] + # change view [b, h, sq, d] context_layer = context_layer.view(*output_size) if q_format == "sbhd": - # [b, np, sq, hn] --> [sq, b, np, hn] + # [b, h, sq, d] --> [sq, b, h, d] context_layer = context_layer.permute(2, 0, 1, 3).contiguous() - # [sq, b, np, hn] --> [sq, b, hp] - context_layer = context_layer.view(seqlen, batch_size, -1) + # [sq, b, h, d] --> [sq, b, hd] + context_layer = context_layer.view(max_seqlen_q, batch_size, -1) if q_format == "bshd": - # [b, np, sq, hn] --> [b, sq, np, hn] + # [b, h, sq, d] --> [b, sq, h, d] context_layer = context_layer.permute(0, 2, 1, 3).contiguous() - # [b, sq, np, hn] --> [b, sq, hp] - context_layer = context_layer.view(batch_size, seqlen, -1) + # [b, sq, h, d] --> [b, sq, hd] + context_layer = context_layer.view(batch_size, max_seqlen_q, -1) if q_format == "thd": - # [b, np, sq, hn] --> [b, sq, np, hn] + # [b, h, sq, d] --> [b, sq, h, d] context_layer = context_layer.permute(0, 2, 1, 3).contiguous() - # [b, sq, np, hn] --> [tq, np, hn] + # [b, sq, h, d] --> [tq, h, d] context_layer = ConvertBSHDtoTHD.apply( context_layer, cu_seqlens_q, ) - # [tq, np, hn] --> [tq, hp] + # [tq, h, d] --> [tq, hd] context_layer = context_layer.view(context_layer.shape[0], -1) if fp8: @@ -1254,21 +1299,26 @@ def forward( if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None: fp8_recipe = fp8_meta["local_recipes"][0] - # input types are inferred from the real data while output types are controlled by fp8_output - # fp8_output should be set upstream as (DPA.fp8 and DPA.fp8_meta["recipe"].fp8_mha) + # qkv_layout may change due to MXFP8 quantization + # o_format should stay the same as original q_format + original_qkv_layout = qkv_layout + _, o_format, _ = dpa_utils.get_qkv_format(qkv_layout) + + # input types are inferred from real data while output types are controlled by fp8_output + # fp8_output should be set upstream assert isinstance(k, q.__class__) and isinstance( v, q.__class__ - ), "q, k, v must be of the same class, e.g. torch.Tensor or Float8Tensor." - is_input_fp8 = isinstance(q, Float8Tensor) + ), "q, k, v must be of the same class, e.g. torch.Tensor or QuantizedTensorStorage." + is_input_fp8 = isinstance(q, QuantizedTensorStorage) is_output_fp8 = fp8_output - # whether fwd kernel in FP8: fp8 = (DPA.fp8 and DPA.fp8_meta["recipe"].fp8_dpa) - # whether bwd kernel in FP8: + # whether fwd kernel will be run in FP8: fp8 = (DPA.fp8 and DPA.fp8_meta["recipe"].fp8_dpa) + # whether bwd kernel will be run in FP8: is_bwd_fp8 = fp8 and int(os.getenv("NVTE_FP8_DPA_BWD", "1")) # get quantizers from DPA; all Nones if not fp8 QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer = ( - dpa_utils.get_attention_quantizers(fp8, quantizers) + dpa_utils.get_attention_quantizers(fp8, fp8_recipe, quantizers) ) # get nominal data type for out @@ -1277,16 +1327,20 @@ def forward( out_nominal_dtype = q.dtype max_logit = None + qkv_scale_inv_format = None if fp8: fused_attention_backend = FusedAttnBackend["FP8"] # q, k, v: torch.Tensor; dtype = torch.float16 or torch.bfloat16 - # q_fp8, k_fp8, v_fp8: Float8Tensor; dtype = torch.float16 or torch.bfloat16 - # fp8_dtype = tex.DType.kFloat8E4M3 + # q_fp8, k_fp8, v_fp8: Float8Tensor/MXFP8Tensor; + # dtype = torch.float16 or torch.bfloat16 + # fp8_dtype = tex.DType.kFloat8E4M3 if is_input_fp8: q_fp8, k_fp8, v_fp8 = q, k, v else: - q_fp8, k_fp8, v_fp8 = combine_and_quantize(qkv_layout, q, k, v, QKV_quantizer) + q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format = combine_and_quantize( + qkv_layout, q, k, v, QKV_quantizer, used_in_backward=is_training + ) # print quantizers print_quantizers( @@ -1304,6 +1358,7 @@ def forward( # DelayedScaling: Float8Tensor; dtype = torch.float16 or torch.bfloat16 # fp8_dtype = tex.DType.kFloat8E4M3 # Float8CurrentScaling: torch.Tensor; dtype = torch.float16 or torch.bfloat16 + # MXFP8BlockScaling: torch.Tensor; dtype = torch.float16 or torch.bfloat16 out_, aux_ctx_tensors, *_ = fused_attn_fwd( is_training, max_seqlen_q, @@ -1326,6 +1381,8 @@ def forward( dropout_p, fast_zero_fill, qkv_layout, + o_format, + qkv_scale_inv_format, attn_bias_type, attn_mask_type, softmax_type, @@ -1336,20 +1393,34 @@ def forward( cuda_graph=is_graph_capturing(), ) - # out_fp8: Float8Tensor; dtype = torch.float16 or torch.bfloat16 + # out_fp8: Float8Tensor/MXFP8Tensor; dtype = torch.float16 or torch.bfloat16 # fp8_dtype = tex.DType.kFloat8E4M3 # out: torch.Tensor; dtype = torch.float16 or torch.bfloat16 out_fp8 = out_ - out = out_ - - if isinstance(out_, Float8Tensor): - if not is_output_fp8 or not is_bwd_fp8: - out = out_.dequantize().view(out_.shape) - else: - if is_output_fp8 or ( + out_f16 = out_ + bwd_requires_o_f16 = is_training and ( + not is_bwd_fp8 + or ( is_bwd_fp8 - and not (fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16) - ): + and ( + (fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16) + or fp8_recipe.mxfp8() + ) + ) + ) + bwd_requires_o_fp8 = ( + is_training + and is_bwd_fp8 + and ( + fp8_recipe.delayed() + or (fp8_recipe.float8_current_scaling() and not _dpa_fp8_cs_o_in_f16) + ) + ) + if isinstance(out_, QuantizedTensorStorage): + if not is_output_fp8 or bwd_requires_o_f16: + out_f16 = out_.dequantize().view(out_.shape) + else: + if is_output_fp8 or bwd_requires_o_fp8: out_fp8 = O_quantizer(out_) # print quantizers @@ -1365,21 +1436,25 @@ def forward( ) # return appropriate tensors - out_ret = out_fp8 if is_output_fp8 else out + out_ret = out_fp8 if is_output_fp8 else out_f16 - # save appropriate tensors + # save q, k, v, o tensors fp8_tensors = (None, None, None, None) - qkvo_tensors = (None, None, None, None) + f16_tensors = (None, None, None, None) if is_bwd_fp8: - if fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16: + if ( + fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16 + ) or fp8_recipe.mxfp8(): fp8_tensors = (q_fp8, k_fp8, v_fp8, None) - qkvo_tensors = (None, None, None, out) - else: + f16_tensors = (None, None, None, out_f16) + elif fp8_recipe.delayed() or ( + fp8_recipe.float8_current_scaling() and not _dpa_fp8_cs_o_in_f16 + ): fp8_tensors = (q_fp8, k_fp8, v_fp8, out_fp8) else: if is_input_fp8: q, k, v = combine_and_dequantize(qkv_layout, q_fp8, k_fp8, v_fp8) - qkvo_tensors = (q, k, v, out) + f16_tensors = (q, k, v, out_f16) else: # q, k, v, out_: torch.Tensor; dtype = torch.float16 or torch.bfloat16 out_, aux_ctx_tensors, *max_logit = fused_attn_fwd( @@ -1404,6 +1479,8 @@ def forward( dropout_p, fast_zero_fill, qkv_layout, + o_format, + None, attn_bias_type, attn_mask_type, softmax_type, @@ -1414,10 +1491,10 @@ def forward( return_max_logit, is_graph_capturing(), ) - out = out_ + out_f16 = out_ out_ret = out_ fp8_tensors = (None, None, None, None) - qkvo_tensors = (q, k, v, out) + f16_tensors = (q, k, v, out_f16) nvtx_range_pop(f"{nvtx_label}") @@ -1431,7 +1508,7 @@ def forward( if ctx.fp8: tensor_list = fp8_tensors else: - tensor_list = [q, k, v, out] + tensor_list = [q, k, v, out_f16] mark_activation_offload(*tensor_list) mark_activation_offload(*aux_ctx_tensors) @@ -1441,7 +1518,7 @@ def forward( tensors_to_save, tensor_objects = prepare_for_saving( *fp8_tensors, - *qkvo_tensors, + *f16_tensors, cu_seqlens_q, cu_seqlens_kv, cu_seqlens_q_padded, @@ -1489,9 +1566,17 @@ def forward( ctx.qkv_layout = reload_layout[:-1] else: ctx.qkv_layout = qkv_layout + if fp8 and not ctx.fp8: + ctx.qkv_layout = original_qkv_layout else: ctx.qkv_layout = qkv_layout + if fp8 and not ctx.fp8: + ctx.qkv_layout = original_qkv_layout + ctx.o_format = o_format + ctx.qkv_scale_inv_format = qkv_scale_inv_format + # dqkv should have the same layout as the original qkv + ctx.dqkv_layout = original_qkv_layout ctx.attn_bias_type = attn_bias_type ctx.attn_mask_type = attn_mask_type ctx.softmax_type = softmax_type @@ -1511,14 +1596,24 @@ def forward( def backward(ctx, d_out, *_args): # pylint: disable=missing-function-docstring - # d_out is expected to be in FP8 if is_output_fp8=True, - # but in the case it's not, convert it to FP8 before any operation - if ctx.fp8 and ctx.is_output_fp8 and not isinstance(d_out, QuantizedTensorStorage): - d_out = ctx.dO_quantizer(d_out) - if not ctx.use_FAv2_bwd: - d_out._data = d_out._data.contiguous() - elif not ctx.use_FAv2_bwd: + # d_out: torch.Tensor; dtype = torch.float16 or torch.bfloat16 + # d_out_fp8: Float8Tensor; dtype = torch.float16 or torch.bfloat16 + # fp8_dtype = tex.DType.kFloat8E5M2 + if not isinstance(d_out, QuantizedTensorStorage) and not ctx.use_FAv2_bwd: d_out = d_out.contiguous() + d_out_fp8 = None + do_format = ctx.o_format + do_scale_inv_format = None + if ctx.fp8: + if isinstance(d_out, QuantizedTensorStorage): + d_out_fp8 = d_out + elif isinstance(ctx.dO_quantizer, MXFP8Quantizer): + (d_out_fp8,), do_scale_inv_format = mxfp8_quantize_fast_path( + [(d_out, ctx.dO_quantizer)], + do_format, + ) + else: + d_out_fp8 = ctx.dO_quantizer(d_out) ( q_fp8, k_fp8, @@ -1579,14 +1674,6 @@ def backward(ctx, d_out, *_args): dqkv_nominal_dtype = ctx.nominal_dtype if ctx.fp8: - # d_out: torch.Tensor; dtype = torch.float16 or torch.bfloat16 - # d_out_fp8: Float8Tensor; dtype = torch.float16 or torch.bfloat16 - # fp8_dtype = tex.DType.kFloat8E5M2 - if ctx.is_output_fp8: - d_out_fp8 = d_out - else: - d_out_fp8 = ctx.dO_quantizer(d_out) - # print quantizers print_quantizers( "FusedAttnFunc.backward >> before: ", @@ -1599,27 +1686,31 @@ def backward(ctx, d_out, *_args): ctx.dP_quantizer, ) - # get tex.DType for dq, dk, dv data - dqkv_te_dtype = d_out_fp8._fp8_dtype - - # q_fp8, k_fp8, v_fp8, out_fp8: Float8Tensor; dtype = torch.float16 or torch.bfloat16, + # DelayedScaling/Float8CurrentScaling/MXFP8BlockScaling: + # q_fp8, k_fp8, v_fp8: Float8Tensor/MXFP8Tensor; dtype = torch.float16 or torch.bfloat16, # fp8_dtype = tex.DType.kFloat8E4M3 - # d_out_fp8: Float8Tensor; dtype = torch.float16 or torch.bfloat16 + # d_out_fp8: Float8Tensor/MXFP8Tensor; dtype = torch.float16 or torch.bfloat16 # fp8_dtype = tex.DType.kFloat8E5M2 - # out_: - # DelayedScaling: Float8Tensor; dtype = torch.float16 or torch.bfloat16 + # DelayedScaling: + # out_: Float8Tensor; dtype = torch.float16 or torch.bfloat16 # fp8_dtype = tex.DType.kFloat8E4M3 - # Float8CurrentScaling: torch.Tensor; dtype = torch.float16 or torch.bfloat16 - # - # dq_, dk_, dv_: - # DelayedScaling: Float8Tensor; dtype = torch.float16 or torch.bfloat16 + # dq_, dk_, dv_: Float8Tensor; dtype = torch.float16 or torch.bfloat16 # fp8_dtype = tex.DType.kFloat8E5M2 - # Float8CurrentScaling: torch.Tensor; dtype = torch.float16 or torch.bfloat16 - out_ = ( - out - if ctx.fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16 - else out_fp8 - ) + # Float8CurrentScaling: + # out_: NVTE_DPA_FP8CS_O_in_F16=1: + # torch.Tensor; dtype = torch.float16 or torch.bfloat16 + # NVTE_DPA_FP8CS_O_in_F16=0: + # Float8Tensor; dtype = torch.float16 or torch.bfloat16 + # fp8_dtype = tex.DType.kFloat8E4M3 + # dq_, dk_, dv_: torch.Tensor; dtype = torch.float16 or torch.bfloat16 + # MXFP8BlockScaling: + # out_, dq_, dk_, dv_, d_out: torch.Tensor; dtype = torch.float16 or torch.bfloat16 + out_ = out_fp8 + if ctx.fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16: + out_ = out + if ctx.fp8_recipe.mxfp8(): + out_ = out + aux_ctx_tensors.append(d_out) dq_, dk_, dv_, *rest = fused_attn_bwd( ctx.max_seqlen_q, ctx.max_seqlen_kv, @@ -1631,7 +1722,6 @@ def backward(ctx, d_out, *_args): out_, d_out_fp8, dqkv_nominal_dtype, - dqkv_te_dtype, aux_ctx_tensors, ctx.fused_attention_backend, cu_seqlens_q_padded, @@ -1643,6 +1733,11 @@ def backward(ctx, d_out, *_args): ctx.dropout_p, ctx.fast_zero_fill, ctx.qkv_layout, + ctx.o_format, + do_format, + ctx.dqkv_layout, + ctx.qkv_scale_inv_format, + do_scale_inv_format, ctx.attn_bias_type, ctx.attn_mask_type, ctx.softmax_type, @@ -1651,23 +1746,22 @@ def backward(ctx, d_out, *_args): ctx.deterministic, is_graph_capturing(), ) - # dq, dk, dv: torch.Tensor; dtype = torch.float16 or torch.bfloat16 dq, dk, dv = dq_, dk_, dv_ - is_float8tensor = isinstance(dq_, Float8Tensor) - if is_float8tensor and not ctx.is_input_fp8: + is_quantized_tensor = isinstance(dq_, QuantizedTensorStorage) + if is_quantized_tensor and not ctx.is_input_fp8: # return in F16 dq, dk, dv = combine_and_dequantize( - ctx.qkv_layout, + ctx.dqkv_layout, dq_, dk_, dv_, src_nominal_dtype=dq_.dtype, ) - if not is_float8tensor and ctx.is_input_fp8: + if not is_quantized_tensor and ctx.is_input_fp8: # return in FP8 - dq, dk, dv = combine_and_quantize( - ctx.qkv_layout, dq_, dk_, dv_, ctx.dQKV_quantizer + dq, dk, dv, _, _ = combine_and_quantize( + ctx.dqkv_layout, dq_, dk_, dv_, ctx.dQKV_quantizer ) # print quantizers @@ -1684,7 +1778,6 @@ def backward(ctx, d_out, *_args): else: if isinstance(d_out, QuantizedTensorStorage): d_out = d_out.dequantize(dtype=ctx.nominal_dtype) - dqkv_te_dtype = TE_DType[d_out.dtype] # q, k, v, out, d_out, dq, dk, dv: torch.Tensor; torch.float16 or torch.bfloat16 dq, dk, dv, *rest = fused_attn_bwd( ctx.max_seqlen_q, @@ -1697,7 +1790,6 @@ def backward(ctx, d_out, *_args): out, d_out, dqkv_nominal_dtype, - dqkv_te_dtype, aux_ctx_tensors, ctx.fused_attention_backend, cu_seqlens_q_padded, @@ -1709,6 +1801,11 @@ def backward(ctx, d_out, *_args): ctx.dropout_p, ctx.fast_zero_fill, ctx.qkv_layout, + ctx.o_format, + do_format, + ctx.dqkv_layout, + None, + None, ctx.attn_bias_type, ctx.attn_mask_type, ctx.softmax_type, @@ -1873,9 +1970,9 @@ def forward( fused_attention_backend != tex.NVTE_Fused_Attn_Backend.NVTE_No_Backend ), "No fused attention backend supports this input combination!" assert all( - x.dtype in [torch.float16, torch.bfloat16] or isinstance(x, Float8Tensor) + x.dtype in [torch.float16, torch.bfloat16] or isinstance(x, QuantizedTensorStorage) for x in [query_layer, key_layer, value_layer] - ), "FusedAttention only supports FP16 and BF16 data types, or Float8Tensors." + ), "FusedAttention only supports FP16 and BF16 data types, or QuantizedTensors." assert ( query_layer.is_cuda and key_layer.is_cuda and value_layer.is_cuda ), "FusedAttention only supports CUDA tensors." @@ -1981,7 +2078,7 @@ def forward( " with FP8!" ) if fp8_recipe.float8_current_scaling() and context_parallel: - all_quantizers = dpa_utils.get_attention_quantizers(fp8, quantizers) + all_quantizers = dpa_utils.get_attention_quantizers(fp8, fp8_recipe, quantizers) for q in all_quantizers: if isinstance(q, Float8CurrentScalingQuantizer): q.with_amax_reduction = True diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 64cccaac6..dfc15cc6c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -22,13 +22,11 @@ ) from transformer_engine.pytorch.quantization import FP8GlobalStateManager from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor +from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.jit import jit_fuser from transformer_engine.pytorch.graph import is_graph_capturing -from transformer_engine.pytorch.constants import ( - dist_group_type, - TE_DType, -) +from transformer_engine.pytorch.constants import dist_group_type from transformer_engine.pytorch.distributed import ( get_distributed_world_size, get_distributed_rank, @@ -48,6 +46,7 @@ combine_and_quantize, combine_and_dequantize, print_quantizers, + mxfp8_quantize_fast_path, ) _cu_seqlens_info_with_cp_cache = {} @@ -59,6 +58,18 @@ _dpa_fp8_cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1" +def get_bsh_dims(tensor_format): + """Get batch dimension and sequence dimension from tensor format""" + if tensor_format in ["bshd", "sbhd", "bhsd"]: + batch_dim = tensor_format.index("b") + seq_dim = tensor_format.index("s") + head_dim = tensor_format.index("h") + else: # tensor_format == "thd" + batch_dim = seq_dim = tensor_format.index("t") + head_dim = tensor_format.index("h") + return batch_dim, seq_dim, head_dim + + def flash_attn_p2p_communicate( rank, send_tensor, send_dst, recv_tensor, recv_src, cp_group, batch_p2p_comm ): @@ -237,10 +248,10 @@ def get_seq_chunk_ids_for_reordering_after_attn(cp_size, device): def reorder_seq_chunks_for_a2a_before_attn(x, chunk_ids_for_a2a, seq_dim, cp_size): """Reorder sequence chunk for A2A communication before attention compute.""" # [cp, b, s, h//cp, d] -> [b, cp, s, h//cp, d] - # or [cp, s, b, h//cp, d] -> [cp, s, b, h//cp, d] + # [cp, s, b, h//cp, d] -> [cp, s, b, h//cp, d] x = x.movedim(0, seq_dim).contiguous() # [b, cp, s, h//cp, d] -> [b, cp*2, s//2, h//cp, d] - # or [cp, s, b, h//cp, d] -> [cp*2, s//2, b, h//cp, d] + # [cp, s, b, h//cp, d] -> [cp*2, s//2, b, h//cp, d] x = x.view(*x.shape[:seq_dim], cp_size * 2, -1, *x.shape[(seq_dim + 2) :]) # reorder the sequence chunks x = torch.index_select(x, dim=seq_dim, index=chunk_ids_for_a2a) @@ -251,12 +262,12 @@ def reorder_seq_chunks_for_a2a_before_attn(x, chunk_ids_for_a2a, seq_dim, cp_siz def reorder_seq_chunks_for_a2a_after_attn(x, chunk_ids_for_a2a, seq_dim, cp_size): """Reorder sequence chunk for A2A communication after attention compute.""" # [b, cp*2, s//2, h//cp, d] -> [cp*2, b, s//2, h//cp, d] - # or [cp*2, s//2, b, h//cp, d] -> [cp*2, s//2, b, h//cp, d] + # [cp*2, s//2, b, h//cp, d] -> [cp*2, s//2, b, h//cp, d] x = x.movedim(seq_dim, 0).contiguous() # reorder the sequence chunks x = torch.index_select(x, dim=0, index=chunk_ids_for_a2a) # [cp*2, b, s//2, h//cp, d] -> [cp, 2, b, s//2, h//cp, d] - # or [cp*2, s//2, b, h//cp, d] -> [cp, 2, s//2, b, h//cp, d] + # [cp*2, s//2, b, h//cp, d] -> [cp, 2, s//2, b, h//cp, d] x = x.view(cp_size, 2, *x.shape[1:]) return x @@ -410,15 +421,32 @@ def flash_attn_a2a_communicate( cp_stream: torch.cuda.Stream, before_attn: bool, qkv_format: str = "bshd", - cu_seqlens_padded: torch.Tensor = None, + cu_seqlens_q_padded: torch.Tensor = None, + cu_seqlens_kv_padded: torch.Tensor = None, + a2a_input_names: List[str] = None, ) -> Union[torch.Tensor, List[torch.Tensor]]: """A2A communication for context parallelism.""" - - assert ( - qkv_format != "thd" or cu_seqlens_padded is not None - ), "cu_seqlens_padded is required for THD format!" + assert a2a_input_names in [ + ["q", "k", "v"], + ["out"], + ["dout"], + ["dq", "dk", "dv"], + ], "a2a_input_names must be one of ['q', 'k', 'v'], ['out'], ['dout'], ['dq', 'dk', 'dv']!" + if a2a_input_names in [["out"], ["dout"]]: + assert qkv_format != "thd" or cu_seqlens_q_padded is not None, ( + f"flash_attn_a2a_communicate requires cu_seqlens_q_padded for {a2a_input_names} with" + " THD format!" + ) + if a2a_input_names in [["q", "k", "v"], ["dq", "dk", "dv"]]: + assert qkv_format != "thd" or ( + cu_seqlens_q_padded is not None and cu_seqlens_kv_padded is not None + ), ( + "flash_attn_a2a_communicate requires cu_seqlens_q_padded and cu_seqlens_kv_padded for" + f" {a2a_input_names} with THD format!" + ) a2a_inputs = [a2a_inputs] if not isinstance(a2a_inputs, list) else a2a_inputs a2a_outputs, a2a_reqs = [None] * len(a2a_inputs), [None] * len(a2a_inputs) + _, _, head_dim = get_bsh_dims(qkv_format) if before_attn: for i in range(len(a2a_inputs) + 2): if 0 < i < len(a2a_inputs) + 1: @@ -430,18 +458,24 @@ def flash_attn_a2a_communicate( with torch.cuda.stream(cp_stream): a2a_reqs[i - 2].wait() x = a2a_outputs[i - 2] - if qkv_format in ["bshd", "sbhd"]: + if qkv_format in ["bshd", "sbhd", "bhsd"]: # reorder the sequence chunks x = reorder_seq_chunks_for_a2a_before_attn( x, chunk_ids_for_a2a, seq_dim, cp_size ) - # [b, cp*2, s//2, np//cp, hn] -> [b, cp*s, np//cp, hn] - # or [cp*2, s//2, b, np//cp, hn] -> [cp*s, b, np//cp, hn] + # [b, cp*2, s//2, h//cp, d] -> [b, cp*s, h//cp, d] + # [cp*2, s//2, b, h//cp, d] -> [cp*s, b, h//cp, d] + # [b, h//cp, cp*2, s//2, d] -> [b, h//cp, cp*s, d] a2a_outputs[i - 2] = x.view( *x.shape[:seq_dim], -1, *x.shape[(seq_dim + 2) :] ) else: # qkv_format == "thd" - # [cp, t, np//cp, hn] -> [cp*t, np//cp, hn] + cu_seqlens_padded = ( + cu_seqlens_q_padded + if a2a_input_names[i - 2] in ["q", "out", "dout", "dq"] + else cu_seqlens_kv_padded + ) + # [cp, t, h//cp, d] -> [cp*t, h//cp, d] x = x.view(-1, *x.shape[2:]) # reorder the sequence chunks a2a_outputs[i - 2] = reorder_seq_chunks_after_a2a_before_attn_thd( @@ -450,14 +484,21 @@ def flash_attn_a2a_communicate( if i < len(a2a_inputs): x = a2a_inputs[i] - # [b, s, np, hn] -> [b, s, cp, np//cp, hn] - # or [s, b, np, hn] -> [s, b, cp, np//cp, hn] - # or [t, np, hn] -> [t, cp, np//cp, hn] - x = x.view(*x.shape[:-2], cp_size, x.shape[-2] // cp_size, x.shape[-1]) - # [b, s, cp, np//cp, hn] -> [cp, b, s, np//cp, hn] - # or [s, b, cp, np//cp, hn] -> [cp, s, b, np//cp, hn] - # or [t, cp, np//cp, hn] -> [cp, t, np//cp, hn] - a2a_inputs[i] = x.movedim(-3, 0).contiguous() + # [b, s, h, d] -> [b, s, cp, h//cp, d] + # [s, b, h, d] -> [s, b, cp, h//cp, d] + # [b, h, s, d] -> [b, cp, h//cp, s, d] + # [t, h, d] -> [t, cp, h//cp, d] + x = x.view( + *x.shape[:head_dim], + cp_size, + x.shape[head_dim] // cp_size, + *x.shape[head_dim + 1 :], + ) + # [b, s, cp, h//cp, d] -> [cp, b, s, h//cp, d] + # [s, b, cp, h//cp, d] -> [cp, s, b, h//cp, d] + # [b, cp, h//cp, s, d] -> [cp, b, h//cp, s, d] + # [t, cp, h//cp, d] -> [cp, t, h//cp, d] + a2a_inputs[i] = x.movedim(head_dim, 0).contiguous() else: for i in range(len(a2a_inputs) + 2): if 0 < i < len(a2a_inputs) + 1: @@ -467,30 +508,57 @@ def flash_attn_a2a_communicate( ) if i < len(a2a_inputs): x = a2a_inputs[i] - if qkv_format in ["bshd", "sbhd"]: - # [b, cp*s, np//cp, hn] -> [b, cp*2, s//2, np//cp, hn] - # or [cp*s, b, np//cp, hn] -> [cp*2, s//2, b, np//cp, hn] + if qkv_format in ["bshd", "sbhd", "bhsd"]: + # [b, cp*s, h//cp, d] -> [b, cp*2, s//2, h//cp, d] + # [cp*s, b, h//cp, d] -> [cp*2, s//2, b, h//cp, d] + # [b, h//cp, cp*s, d] -> [b, h//cp, cp*2, s//2, d] x = x.view(*x.shape[:seq_dim], cp_size * 2, -1, *x.shape[(seq_dim + 1) :]) # reorder the sequence chunks a2a_inputs[i] = reorder_seq_chunks_for_a2a_after_attn( x, chunk_ids_for_a2a, seq_dim, cp_size ) else: # qkv_format == "thd" + cu_seqlens_padded = ( + cu_seqlens_q_padded + if a2a_input_names[i] in ["q", "out", "dout", "dq"] + else cu_seqlens_kv_padded + ) # reorder the sequence chunks x = reorder_seq_chunks_before_a2a_after_attn_thd(x, cu_seqlens_padded, cp_size) - # [cp*t, np//cp, hn] -> [cp, t, np//cp, hn] + # [cp*t, h//cp, d] -> [cp, t, h//cp, d] a2a_inputs[i] = x.view(cp_size, -1, *x.shape[-2:]) if i > 1: with torch.cuda.stream(cp_stream): a2a_reqs[i - 2].wait() x = a2a_outputs[i - 2] - # [cp, 2, b, s//2, np//cp, hn] -> [b, 2, s//2, cp, np//cp, hn] - # or [cp, 2, s//2, b, np//cp, hn] -> [2, s//2, b, cp, np//cp, hn] - # or [cp, t, np//cp, hn] -> [t, cp, np//cp, hn] - x = x.movedim(0, -3).movedim(0, seq_dim).contiguous() - # [b, 2, s//2, cp, np//cp, hn] -> [b*s, np, hn] - # or [2, s//2, b, cp, np//cp, hn] -> [s*b, np, hn] - # or [t, cp, np//cp, hn] -> [t, np, hn] + # [cp, 2, b, s//2, h//cp, d] -> [2, b, s//2, cp, h//cp, d] + # [cp, 2, s//2, b, h//cp, d] -> [2, s//2, b, cp, h//cp, d] + # [cp, 2, b, h//cp, s//2, d] -> [2, b, cp, h//cp, s//2, d] + # [cp, t, h//cp, d] -> [t, cp, h//cp, d] + tmp_list = list(qkv_format) + if "t" not in qkv_format: + tmp_list.insert(0, "2") + tmp_list.insert(0, "c") + tmp_format = "".join(tmp_list) + head_dim_ = tmp_format.index("h") - 1 + tmp_list.insert(head_dim_, tmp_list.pop(0)) + x = x.movedim(0, head_dim_) + # [2, b, s//2, cp, h//cp, d] -> [b, 2, s//2, cp, h//cp, d] + # [2, s//2, b, cp, h//cp, d] -> [2, s//2, b, cp, h//cp, d] + # [2, b, cp, h//cp, s//2, d] -> [b, cp, h//cp, 2, s//2, d] + # [t, cp, h//cp, d] -> [t, cp, h//cp, d] + if "t" not in qkv_format: + tmp_format = "".join(tmp_list) + seq_dim_ = tmp_format.index("s") - 1 + tmp_list.insert(seq_dim_, tmp_list.pop(0)) + x = x.movedim(0, seq_dim_) + else: + seq_dim_ = 0 + x = x.contiguous() + # [b, 2, s//2, cp, h//cp, d] -> [b*s, h, d] + # [2, s//2, b, cp, h//cp, d] -> [s*b, h, d] + # [b, cp, h//cp, 2, s//2, d] -> [b*h, s, d] + # [t, cp, h//cp, d] -> [t, h, d] a2a_outputs[i - 2] = x.view(-1, x.shape[-3] * x.shape[-2], x.shape[-1]) torch.cuda.current_stream().wait_stream(cp_stream) return a2a_outputs[0] if len(a2a_inputs) == 1 else a2a_outputs @@ -775,13 +843,16 @@ def cp_p2p_fwd_fused_attn( softmax_scale, dropout_p, qkv_layout, + o_format, attn_mask_type, attn_bias_type, fp8, + fp8_recipe, q_fp8, k_fp8, v_fp8, fwd_nominal_dtype, + QKV_quantizer, S_quantizer_per_step, O_quantizer_per_step, rank, @@ -867,11 +938,18 @@ def cp_p2p_fwd_fused_attn( cu_seqlens_kv_padded_ = cu_seqlens_kv_padded fp8_meta_kwargs = {} + new_qkv_layout = qkv_layout + qkv_scale_inv_format = None if fp8: - q_part, k_part, v_part = [ - Float8Tensor.make_like(x, data=y, dtype=fwd_nominal_dtype) - for x, y in zip([q_fp8, k_fp8, v_fp8], [q_part, k_part, v_part]) - ] + if not fp8_recipe.mxfp8(): + q_part, k_part, v_part = [ + Float8Tensor.make_like(x, data=y, dtype=fwd_nominal_dtype) + for x, y in zip([q_fp8, k_fp8, v_fp8], [q_part, k_part, v_part]) + ] + else: + q_part, k_part, v_part, new_qkv_layout, qkv_scale_inv_format = combine_and_quantize( + qkv_layout, q_part, k_part, v_part, QKV_quantizer + ) fp8_meta_kwargs["s_quantizer"] = S_quantizer_per_step fp8_meta_kwargs["o_quantizer"] = O_quantizer_per_step @@ -888,7 +966,8 @@ def cp_p2p_fwd_fused_attn( fused_attention_backend=fused_attn_backend, attn_scale=softmax_scale, dropout=dropout_p, - qkv_layout=qkv_layout, + qkv_layout=new_qkv_layout, + o_format=o_format, attn_mask_type=attn_mask_type_, attn_bias_type=attn_bias_type, attn_bias=attn_bias_inputs, @@ -897,10 +976,14 @@ def cp_p2p_fwd_fused_attn( **fp8_meta_kwargs, return_max_logit=return_max_logit, cuda_graph=is_graph_capturing(), + qkv_scale_inv_format=qkv_scale_inv_format, ) if fp8: - softmax_lse_per_step, _, rng_states = aux_ctx_tensors + if qkv_layout != "t3hd": + softmax_lse_per_step, rng_states = aux_ctx_tensors + else: + softmax_lse_per_step, _, rng_states = aux_ctx_tensors else: softmax_lse_per_step, rng_states, *rest = aux_ctx_tensors attn_bias = rest[0] if len(rest) > 0 else None @@ -1065,15 +1148,19 @@ def cp_p2p_bwd_fused_attn( softmax_scale, dropout_p, qkv_layout, + o_format, + do_format, + dqkv_layout, attn_mask_type, attn_bias_type, deterministic, fwd_nominal_dtype, bwd_nominal_dtype, - bwd_output_te_dtype, S_quantizer, dP_quantizer_per_step, dQKV_quantizer_per_step, + QKV_quantizer_per_step, + dO_quantizer_per_step, q_part, k_part, v_part, @@ -1083,11 +1170,14 @@ def cp_p2p_bwd_fused_attn( ): """Per-tile backward call of CP P2P with FusedAttention backend""" if fp8: - aux_tensors = [ - softmax_lse, - softmax_lse, - rng_states[cp_size - step - 1], - ] + if qkv_layout == "t3hd": + aux_tensors = [ + softmax_lse, + softmax_lse, + rng_states[cp_size - step - 1], + ] + else: + aux_tensors = [softmax_lse, rng_states[cp_size - step - 1]] else: aux_tensors = [softmax_lse, rng_states[cp_size - step - 1]] @@ -1106,11 +1196,14 @@ def cp_p2p_bwd_fused_attn( elif section == "upper-triangle": q_part, out_part, dout_part = [x.contiguous() for x in [q_part, out_part, dout_part]] if fp8: - aux_tensors = [ - softmax_lse_, - softmax_lse_, - rng_states[cp_size - step - 1], - ] + if qkv_layout == "t3hd": + aux_tensors = [ + softmax_lse_, + softmax_lse_, + rng_states[cp_size - step - 1], + ] + else: + aux_tensors = [softmax_lse_, rng_states[cp_size - step - 1]] else: aux_tensors = [softmax_lse_, rng_states[cp_size - step - 1]] @@ -1122,17 +1215,37 @@ def cp_p2p_bwd_fused_attn( aux_tensors += [attn_biases[cp_size - step - 1]] fp8_meta_kwargs = {} + qkv_scale_inv_format = None + do_scale_inv_format = None if fp8: - q_part, k_part, v_part = [ - Float8Tensor.make_like(x, data=y, dtype=fwd_nominal_dtype) - for x, y in zip( - [q_fp8, kv_fp8, kv_fp8], - [q_part, k_part, v_part], + if not fp8_recipe.mxfp8(): + q_part, k_part, v_part = [ + Float8Tensor.make_like(x, data=y, dtype=fwd_nominal_dtype) + for x, y in zip( + [q_fp8, kv_fp8, kv_fp8], + [q_part, k_part, v_part], + ) + ] + else: + q_part, k_part, v_part, qkv_layout, qkv_scale_inv_format = combine_and_quantize( + qkv_layout, + q_part, + k_part, + v_part, + QKV_quantizer_per_step, + used_in_forward=False, + used_in_backward=True, + ) + if not fp8_recipe.mxfp8(): + if not (fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16): + out_part = Float8Tensor.make_like(out_fp8, data=out_part, dtype=fwd_nominal_dtype) + dout_part = Float8Tensor.make_like(dout_fp8, data=dout_part, dtype=bwd_nominal_dtype) + else: + aux_tensors.append(dout_part) + (dout_part,), do_scale_inv_format = mxfp8_quantize_fast_path( + [(dout_part, dO_quantizer_per_step)], + do_format, ) - ] - if not (fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16): - out_part = Float8Tensor.make_like(out_fp8, data=out_part, dtype=fwd_nominal_dtype) - dout_part = Float8Tensor.make_like(dout_fp8, data=dout_part, dtype=bwd_nominal_dtype) fp8_meta_kwargs["s_quantizer"] = S_quantizer fp8_meta_kwargs["dp_quantizer"] = dP_quantizer_per_step fp8_meta_kwargs["dqkv_quantizer"] = dQKV_quantizer_per_step @@ -1148,7 +1261,6 @@ def cp_p2p_bwd_fused_attn( out_part, dout_part, bwd_nominal_dtype, - bwd_output_te_dtype, aux_tensors, fused_attn_backend, cu_seqlens_q_padded=cu_seqlens_q_padded_, @@ -1156,10 +1268,15 @@ def cp_p2p_bwd_fused_attn( attn_scale=softmax_scale, dropout=dropout_p, qkv_layout=qkv_layout, + o_format=o_format, + do_format=do_format, + dqkv_layout=dqkv_layout, attn_mask_type=attn_mask_type_, attn_bias_type=attn_bias_type, deterministic=deterministic, cuda_graph=is_graph_capturing(), + qkv_scale_inv_format=qkv_scale_inv_format, + do_scale_inv_format=do_scale_inv_format, **fp8_meta_kwargs, ) @@ -1313,16 +1430,15 @@ def forward( ) # set up attention args - enable_mla = k.shape[-1] != v.shape[-1] - causal = "causal" in attn_mask_type - if softmax_scale is None: softmax_scale = q.shape[-1] ** (-0.5) - + causal = "causal" in attn_mask_type + qkv_layout = qkv_format + "_" + qkv_format + "_" + qkv_format + orig_q_shape, orig_k_shape, orig_v_shape = q.shape, k.shape, v.shape + orig_o_shape = q.shape[:-1] + v.shape[-1:] batch_dim = None seq_dim = None cu_seqlens_q_half, cu_seqlens_kv_half = None, None - qkv_layout = qkv_format + "_" + qkv_format + "_" + qkv_format if qkv_format in ["bshd", "sbhd"]: seq_dim = qkv_format.index("s") cu_seqlens_q_padded, cu_seqlens_kv_padded = None, None @@ -1337,13 +1453,10 @@ def forward( else: cu_seqlens_q_padded = cu_seqlens_q_padded // cp_size cu_seqlens_kv_padded = cu_seqlens_kv_padded // cp_size - max_seqlen_q = max_seqlen_q // cp_size max_seqlen_kv = max_seqlen_kv // cp_size cu_seqlens_q_per_step = [None for _ in range(cp_size)] cu_seqlens_kv_per_step = [None for _ in range(cp_size)] - - fused_attn_backend = None amax_per_step = None S_quantizer_per_step = [None for _ in range(cp_size)] O_quantizer_per_step = [None for _ in range(cp_size)] @@ -1352,9 +1465,9 @@ def forward( assert isinstance(k, q.__class__) and isinstance( v, q.__class__ - ), "q, k, v must be of the same class, e.g. torch.Tensor or Float8Tensor." + ), "q, k, v must be of the same class, e.g. torch.Tensor or QuantizedTensorStorage." fwd_nominal_dtype = q.dtype - is_input_fp8 = isinstance(q, Float8Tensor) + is_input_fp8 = isinstance(q, QuantizedTensorStorage) is_output_fp8 = fp8_output is_bwd_fp8 = int(os.getenv("NVTE_FP8_DPA_BWD", "1")) # recipe passed in through autocast or set by NVTE_DPA_FP8_RECIPE; @@ -1362,7 +1475,6 @@ def forward( fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None: fp8_recipe = fp8_meta["local_recipes"][0] - ( QKV_quantizer, O_quantizer, @@ -1370,43 +1482,58 @@ def forward( dQKV_quantizer, dO_quantizer, dP_quantizer, - ) = dpa_utils.get_attention_quantizers(fp8, quantizers) + ) = dpa_utils.get_attention_quantizers(fp8, fp8_recipe, quantizers) - q_f16 = None + # q, k, v a2a: gather s and split h + # FP8DS/CS: Float8Tensor -> torch.uint8 -> Float8Tensor + # MXFP8/F16: fwd_nominal_dtype q_fp8, k_fp8, v_fp8 = (None, None, None) - # communicate for the 'a2a' part of 'a2a+p2p' if cp_size_a2a > 1: if fp8 and is_input_fp8: - QKV_quantizer = q._quantizer q_fp8, k_fp8, v_fp8 = q, k, v - q, k, v = (q._data, k._data, v._data) + if not fp8_recipe.mxfp8(): + q, k, v = [q_fp8._data, k_fp8._data, v_fp8._data] chunk_ids_for_a2a = get_seq_chunk_ids_for_reordering_before_attn(cp_size_a2a, q.device) q, k, v = flash_attn_a2a_communicate( - [q, k, v], chunk_ids_for_a2a, seq_dim, cp_size_a2a, cp_group_a2a, cp_stream, True + [q, k, v], + chunk_ids_for_a2a, + seq_dim, + cp_size_a2a, + cp_group_a2a, + cp_stream, + True, + qkv_format=qkv_format, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, + a2a_input_names=["q", "k", "v"], ) - if fp8 and is_input_fp8: + if fp8 and is_input_fp8 and not fp8_recipe.mxfp8(): q_fp8, k_fp8, v_fp8 = [ Float8Tensor.make_like(x, data=y, dtype=fwd_nominal_dtype) for x, y in zip([q_fp8, k_fp8, v_fp8], [q, k, v]) ] q, k, v = q_fp8, k_fp8, v_fp8 + post_a2a_o_shape = q.shape[:-1] + v.shape[-1:] # convert qkv to the right type + q_f16 = None + fused_attn_backend = None if fp8: assert use_fused_attention, "FP8 is only supported with Fused Attention!" fused_attn_backend = FusedAttnBackend["FP8"] - if is_input_fp8: # q_fp8, k_fp8, v_fp8: Float8Tensor, dtype=fwd_nominal_dtype # q, k, v: torch.Tensor, dtype=torch.uint8 q_fp8, k_fp8, v_fp8 = q, k, v - q, k, v = [q_fp8._data, k_fp8._data, v_fp8._data] - else: + elif not fp8_recipe.mxfp8(): # q_f16: torch.Tensor, dtype=fwd_nominal_dtype # q_fp8, k_fp8, v_fp8: Float8Tensor, dtype=fwd_nominal_dtype # q, k, v: torch.Tensor, dtype=torch.uint8 q_f16 = q - q_fp8, k_fp8, v_fp8 = combine_and_quantize(qkv_layout, q, k, v, QKV_quantizer) + q_fp8, k_fp8, v_fp8, qkv_layout, _ = combine_and_quantize( + qkv_layout, q, k, v, QKV_quantizer + ) + if not fp8_recipe.mxfp8(): q, k, v = [q_fp8._data, k_fp8._data, v_fp8._data] # print quantizers @@ -1427,10 +1554,11 @@ def forward( # per_step tensors are not reduced even if Float8CurrentScaling.with_amax_reduction=True; # only used to hold temporary scale/amax values (output only, no quantization op) for i in range(cp_size): - S_quantizer_per_step[i] = S_quantizer.copy() - S_quantizer_per_step[i].amax = amax_per_step[0][i].reshape((1,)) + S_quantizer_per_step[i] = S_quantizer.copy() if S_quantizer is not None else None O_quantizer_per_step[i] = O_quantizer.copy() - O_quantizer_per_step[i].amax = amax_per_step[1][i].reshape((1,)) + if not fp8_recipe.mxfp8(): + S_quantizer_per_step[i].amax = amax_per_step[0][i].reshape((1,)) + O_quantizer_per_step[i].amax = amax_per_step[1][i].reshape((1,)) else: # q_f16: torch.Tensor, dtype=fwd_nominal_dtype # q, k, v: torch.Tensor, dtype=fwd_nominal_dtype @@ -1482,7 +1610,6 @@ def forward( attn_bias_ = attn_bias.view( *attn_bias.shape[:-1], 2 * cp_size, attn_bias.shape[-1] // (2 * cp_size) ) - # [b, h, sq, sk] -> [b, h, sq, 2*cp, sk//(2*cp)] attn_bias = attn_bias.view( *attn_bias.shape[:-1], 2 * cp_size, attn_bias.shape[-1] // (2 * cp_size) @@ -1557,17 +1684,22 @@ def forward( # synchronize fwd results correction across steps fwd_results_correction_done = torch.cuda.Event() + # q, k, v, o: + # causal: [b, 2, s//2, h, d] or [2, s//2, b, h, d] + # non-causal: [b, s, h, d] or [s, b, h, d] p2p_comm_buffers = [None for _ in range(cp_size)] k_shape = k.shape k_numel = k.numel() v_shape = v.shape + o_shape = q.shape[:-1] + v.shape[-1:] p2p_comm_buffers[0] = torch.cat((k.view(-1), v.view(-1)), dim=-1) send_recv_reqs = [[], []] # P2P communication and compute: each rank has cp_size steps - # f16 attention: q, k, v: torch.Tensor, dtype=fwd_nominal_dtype - # fp8 attention: q, k, v: torch.Tensor, dtype=torch.uint8 + # MXFP8/F16 attention: q, k, v: torch.Tensor, dtype=fwd_nominal_dtype + # FP8DS/CS attention: q, k, v: torch.Tensor, dtype=torch.uint8 out = None + o_format = qkv_format for i in range(cp_size + 1): if i < cp_size: with torch.cuda.stream(flash_attn_streams[i % 2]): @@ -1621,13 +1753,16 @@ def forward( softmax_scale, dropout_p, qkv_layout, + o_format, attn_mask_type, attn_bias_type, fp8, + fp8_recipe, q_fp8, k_fp8, v_fp8, fwd_nominal_dtype, + QKV_quantizer, S_quantizer_per_step[i], O_quantizer_per_step[i], rank, @@ -1775,8 +1910,8 @@ def forward( with torch.cuda.stream(flash_attn_streams[(i - 1) % 2]): if use_fused_attention: - # [b, h, sq, 1] -> [b, h, sq] or - # [t, h, 1] -> [t, np] + # [b, h, sq, 1] -> [b, h, sq] + # [t, h, 1] -> [t, h] softmax_lse_per_step[i - 1].squeeze_(-1) if softmax_lse_in_packed_format: softmax_lse_per_step[i - 1] = ( @@ -1788,21 +1923,16 @@ def forward( out_per_step[i - 1] = out_per_step[i - 1].dequantize( dtype=torch.float32 ) - if fp8_recipe.float8_current_scaling(): + if fp8_recipe.float8_current_scaling() or fp8_recipe.mxfp8(): out_per_step[i - 1] = out_per_step[i - 1].to(dtype=torch.float32) if i == 1: softmax_lse = torch.clone(softmax_lse_per_step[0]) if qkv_format == "thd": - if enable_mla: - out = torch.zeros_like(v if not fp8 else out_per_step[0]).view( - v_shape - ) + if fp8: + out = torch.zeros_like(out_per_step[0]).view(o_shape) else: - # MHA or GQA - out = torch.zeros_like(q if not fp8 else out_per_step[0]).view( - q.shape - ) + out = torch.zeros(o_shape, dtype=q.dtype, device=q.device) elif (i - 1) <= rank or not causal: flash_attn_fwd_softmax_lse_correction( softmax_lse, softmax_lse_per_step[i - 1] @@ -1842,7 +1972,7 @@ def forward( # fwd output correction: out in torch.float32 for i in range(cp_size): if i <= rank or not causal: - if qkv_format in ["bshd", "sbhd"]: + if o_format in ["bshd", "sbhd"]: if i == 0: out = flash_attn_fwd_out_correction_init( out_per_step[0], @@ -1850,10 +1980,7 @@ def forward( softmax_lse_per_step[0], seq_dim, ) - if enable_mla: - out = out.view(v_shape) - else: - out = out.view(q.shape) + out = out.view(o_shape) else: flash_attn_fwd_out_correction( out.view(*out_per_step[i].shape), @@ -1862,7 +1989,7 @@ def forward( softmax_lse_per_step[i], seq_dim, ) - elif qkv_format == "thd": + elif o_format == "thd": tex.thd_out_correction( out, out_per_step[i], @@ -1873,7 +2000,7 @@ def forward( softmax_lse_in_packed_format, ) else: - if qkv_format in ["bshd", "sbhd"]: + if o_format in ["bshd", "sbhd"]: flash_attn_fwd_second_half_out_correction( out, out_per_step[i], @@ -1881,7 +2008,7 @@ def forward( softmax_lse_per_step[i], seq_dim, ) - elif qkv_format == "thd": + elif o_format == "thd": tex.thd_out_correction( out, out_per_step[i], @@ -1891,35 +2018,31 @@ def forward( True, softmax_lse_in_packed_format, ) - - if qkv_format == "bshd": - out = out.view(out.shape[0], -1, *out.shape[-2:]) - ctx.batch_size = out.shape[0] - elif qkv_format == "sbhd": - out = out.view(-1, *out.shape[-3:]) - ctx.batch_size = out.shape[1] + out = out.view(post_a2a_o_shape) + out_part = out.to(fwd_nominal_dtype) if cp_size_a2a > 1: chunk_ids_for_a2a = get_seq_chunk_ids_for_reordering_after_attn(cp_size_a2a, out.device) out = flash_attn_a2a_communicate( - out, chunk_ids_for_a2a, seq_dim, cp_size_a2a, cp_group_a2a, cp_stream, False + out, + chunk_ids_for_a2a, + seq_dim, + cp_size_a2a, + cp_group_a2a, + cp_stream, + False, + qkv_format=o_format, + cu_seqlens_q_padded=cu_seqlens_q_padded, + a2a_input_names=["out"], ) - if use_fused_attention: - if qkv_format == "bshd": - # [b*s, h, d] -> [b, s, h, d] - out = out.view(ctx.batch_size, -1, *out.shape[-2:]) - elif qkv_format == "sbhd": - # [s*b, h, d] -> [s, b, h, d] - out = out.view(-1, ctx.batch_size, *out.shape[-2:]) + out = out.view(orig_o_shape) if return_max_logit: max_logit = flash_attn_a2a_communicate_softmax_offset( max_logit, 0, cp_size_a2a, cp_group_a2a, cp_stream, False ) - elif not use_fused_attention: - out = out.view(-1, *out.shape[-2:]) # update FP8 quantizers: amax across cp_size steps - if fp8 and use_fused_attention: + if fp8 and use_fused_attention and not fp8_recipe.mxfp8(): amax_cp_fwd = amax_per_step.amax(dim=1) S_quantizer.amax.copy_(amax_cp_fwd[0]) O_quantizer.amax.copy_(amax_cp_fwd[1]) @@ -1942,7 +2065,11 @@ def forward( out_f16 = out.to(fwd_nominal_dtype) if fp8 and ( is_output_fp8 - or (is_bwd_fp8 and not (fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16)) + or ( + is_bwd_fp8 + and not (fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16) + and not fp8_recipe.mxfp8() + ) ): out_fp8 = O_quantizer(out_f16) out_ret = out_fp8 if (fp8 and is_output_fp8) else out_f16 @@ -1953,7 +2080,7 @@ def forward( kv_fp8 = None kv = p2p_comm_buffers[-1] - if fp8: + if fp8 and not fp8_recipe.mxfp8(): q_fp8, kv_fp8 = [ Float8Tensor.make_like(x, data=y, dtype=fwd_nominal_dtype) for x, y in zip([q_fp8, k_fp8], [q, kv]) @@ -1961,17 +2088,28 @@ def forward( # q, kv, out fp8_tensors = (None, None, None) f16_tensors = (None, None, None) + out_f16 = out_part if ctx.fp8: # fwd: fp8, bwd: fp8, save all fp8 fp8_tensors = (q_fp8, kv_fp8, out_fp8) if fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16: f16_tensors = (None, None, out_f16) - elif fp8 and is_input_fp8: + elif fp8_recipe.mxfp8(): + f16_tensors = (q, kv, out_f16) + elif fp8 and is_input_fp8 and not fp8_recipe.mxfp8(): # fwd: fp8, bwd: f16, save all f16 # dequantize fp8 inputs q_f16 = q_fp8.dequantize() kv_f16 = kv_fp8.dequantize() f16_tensors = (q_f16, kv_f16, out_f16) + elif fp8 and is_input_fp8 and fp8_recipe.mxfp8(): + # fwd: fp8, bwd: f16, save all f16 + # there is already an F16 version of the inputs + q_f16, k_f16, v_f16 = combine_and_dequantize(qkv_layout, q, k, v) + kv_f16 = torch.cat((k_f16.view(-1), v_f16.view(-1)), dim=-1) + f16_tensors = (q_f16, kv_f16, out_f16) + elif fp8 and not is_input_fp8 and fp8_recipe.mxfp8(): + f16_tensors = (q, kv, out_f16) elif fp8: # fwd: fp8, bwd: f16, save all f16 # inputs are already in f16 @@ -2009,7 +2147,6 @@ def forward( ctx.max_seqlen_q = max_seqlen_q ctx.max_seqlen_kv = max_seqlen_kv ctx.softmax_scale = softmax_scale - ctx.qkv_format = qkv_format ctx.attn_mask_type = attn_mask_type ctx.attn_bias_type = attn_bias_type ctx.attn_bias_shape = None if attn_bias is None else attn_bias.shape @@ -2022,12 +2159,19 @@ def forward( ctx.is_output_fp8 = is_output_fp8 ctx.use_flash_attn_3 = use_flash_attn_3 - ctx.enable_mla = enable_mla + ctx.orig_q_shape = orig_q_shape + ctx.orig_k_shape = orig_k_shape + ctx.orig_v_shape = orig_v_shape + ctx.orig_o_shape = orig_o_shape + ctx.post_a2a_o_shape = post_a2a_o_shape ctx.k_numel = k_numel ctx.k_shape = k_shape ctx.v_shape = v_shape - + ctx.o_shape = o_shape + ctx.qkv_format = qkv_format + ctx.qkv_layout = qkv_layout ctx.fwd_nominal_dtype = fwd_nominal_dtype + ctx.dQKV_quantizer = dQKV_quantizer ctx.dO_quantizer = dO_quantizer ctx.dP_quantizer = dP_quantizer @@ -2036,14 +2180,14 @@ def forward( ctx.S_quantizer = S_quantizer if ctx.fp8: ctx.QKV_quantizer = QKV_quantizer.copy() - ctx.QKV_quantizer.scale = QKV_quantizer.scale.clone() ctx.O_quantizer = O_quantizer.copy() - ctx.O_quantizer.scale = O_quantizer.scale.clone() - ctx.S_quantizer = S_quantizer.copy() - ctx.S_quantizer.scale = S_quantizer.scale.clone() + ctx.S_quantizer = S_quantizer.copy() if S_quantizer is not None else None + if not ctx.fp8_recipe.mxfp8(): + ctx.QKV_quantizer.scale = QKV_quantizer.scale.clone() + ctx.O_quantizer.scale = O_quantizer.scale.clone() + ctx.S_quantizer.scale = S_quantizer.scale.clone() nvtx_range_pop(f"{nvtx_label}") - if return_max_logit: return out_ret, max_logit return out_ret @@ -2057,8 +2201,13 @@ def backward(ctx, dout, *_args): nvtx_range_push(f"{nvtx_label}") # dout is expected to be in FP8 if is_output_fp8=True, - # but in the case it's not, convert it to FP8 before any operation - if ctx.fp8 and ctx.is_output_fp8 and not isinstance(dout, QuantizedTensorStorage): + # but in the case it's not, convert it to FP8 (except for MXFP8) before any operation + if ( + ctx.fp8 + and ctx.is_output_fp8 + and not isinstance(dout, QuantizedTensorStorage) + and not ctx.fp8_recipe.mxfp8() + ): dout = ctx.dO_quantizer(dout) if ctx.use_fused_attention: dout._data = dout._data.contiguous() @@ -2098,7 +2247,6 @@ def backward(ctx, dout, *_args): # set up attention args causal = "causal" in ctx.attn_mask_type seq_dim = None - qkv_layout = ctx.qkv_format + "_" + ctx.qkv_format + "_" + ctx.qkv_format if ctx.qkv_format in ["bshd", "sbhd"]: seq_dim = ctx.qkv_format.index("s") @@ -2137,13 +2285,13 @@ def backward(ctx, dout, *_args): if ctx.softmax_lse_in_packed_format: softmax_lse_ = softmax_lse_.transpose(0, 1).contiguous() # [b, h, sq//2] -> [b, h, sq//2, 1] or - # [t//2, np] -> [t//2, h, 1] + # [t//2, h] -> [t//2, h, 1] softmax_lse_.unsqueeze_(-1) if ctx.use_fused_attention: if ctx.softmax_lse_in_packed_format: softmax_lse = softmax_lse.transpose(0, 1).contiguous() # [b, h, sq] -> [b, h, sq, 1] or - # [t, np] -> [t, h, 1] + # [t, h] -> [t, h, 1] softmax_lse.unsqueeze_(-1) # assume fwd and bwd always use the same high precision, i.e. torch.float16 or torch.bfloat16 @@ -2158,28 +2306,29 @@ def backward(ctx, dout, *_args): buffer_dtype = torch.uint8 dq_buffer = None dout_fp8 = None - bwd_output_te_dtype = None dkv_buffer = None if ctx.fp8: - assert ctx.use_fused_attention, "FP8 is only supported with Fused Attention!" + assert ctx.use_fused_attention, "FP8 is only supported with FusedAttention backend!" fused_attn_backend = FusedAttnBackend["FP8"] - q, kv, out = ( - q_fp8._data, - kv_fp8._data, - ( - out - if ctx.fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16 - else out_fp8._data - ), - ) + if not ctx.fp8_recipe.mxfp8(): + q, kv, out = ( + q_fp8._data, + kv_fp8._data, + ( + out + if ctx.fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16 + else out_fp8._data + ), + ) # dout_fp8: Float8Tensor, dtype=bwd_nominal_dtype # dout: torch.Tensor, dtype=torch.uint8 - if ctx.is_output_fp8: + if isinstance(dout, QuantizedTensorStorage): dout_fp8 = dout - else: + elif not ctx.fp8_recipe.mxfp8(): dout_fp8 = ctx.dO_quantizer(dout) - dout = dout_fp8._data + if not ctx.fp8_recipe.mxfp8(): + dout = dout_fp8._data # print quantizers print_quantizers( @@ -2193,9 +2342,6 @@ def backward(ctx, dout, *_args): ctx.dP_quantizer, ) - # dout_fp8._fp8_dtype - bwd_output_te_dtype = ctx.dO_quantizer.dtype - # create buffers for reduction in float32 if ctx.fp8_recipe.delayed(): dq_buffer = torch.empty( @@ -2203,7 +2349,7 @@ def backward(ctx, dout, *_args): dtype=buffer_dtype, device=q.device, ) - if ctx.fp8_recipe.float8_current_scaling(): + if ctx.fp8_recipe.float8_current_scaling() or ctx.fp8_recipe.mxfp8(): dq_buffer = torch.empty( q.shape, dtype=torch.float32, @@ -2217,7 +2363,7 @@ def backward(ctx, dout, *_args): ) dkv_recv_buffer = torch.empty_like(dkv_send_buffer) p2p_comm_buffers = [[kv, dkv_send_buffer], [kv_recv_buffer, dkv_recv_buffer]] - if ctx.fp8_recipe.float8_current_scaling(): + if ctx.fp8_recipe.float8_current_scaling() or ctx.fp8_recipe.mxfp8(): dkv_buffer = torch.zeros( kv.shape, dtype=torch.float32, @@ -2230,10 +2376,13 @@ def backward(ctx, dout, *_args): # per_step tensors are not reduced even if Float8CurrentScaling.with_amax_reduction=True; # only used to hold temporary scale/amax values (output only, no quantization op) for i in range(cp_size): - dP_quantizer_per_step[i] = ctx.dP_quantizer.copy() - dP_quantizer_per_step[i].amax = amax_per_step[0][i].reshape((1,)) + dP_quantizer_per_step[i] = ( + ctx.dP_quantizer.copy() if ctx.dP_quantizer is not None else None + ) dQKV_quantizer_per_step[i] = ctx.dQKV_quantizer.copy() - dQKV_quantizer_per_step[i].amax = amax_per_step[1][i].reshape((1,)) + if not ctx.fp8_recipe.mxfp8(): + dP_quantizer_per_step[i].amax = amax_per_step[0][i].reshape((1,)) + dQKV_quantizer_per_step[i].amax = amax_per_step[1][i].reshape((1,)) else: if isinstance(dout, QuantizedTensorStorage): dout = dout.dequantize(dtype=bwd_nominal_dtype) @@ -2244,34 +2393,28 @@ def backward(ctx, dout, *_args): ] p2p_comm_buffers[0][0].copy_(kv) if ctx.use_fused_attention: - bwd_output_te_dtype = TE_DType[bwd_nominal_dtype] fused_attn_backend = FusedAttnBackend["F16_arbitrary_seqlen"] # communicate for the 'a2a' part of 'a2a+p2p' + dout = dout.view(*ctx.orig_o_shape) if cp_size_a2a > 1: - if not ctx.use_fused_attention: - out = out.view(ctx.batch_size, -1, *out.shape[-2:]) - dout = dout.view(*out.shape) chunk_ids_for_a2a = get_seq_chunk_ids_for_reordering_before_attn( cp_size_a2a, out.device ) - out, dout = flash_attn_a2a_communicate( - [out, dout], + dout = flash_attn_a2a_communicate( + dout, chunk_ids_for_a2a, seq_dim, cp_size_a2a, ctx.cp_group_a2a, ctx.cp_stream, True, + qkv_format=ctx.qkv_format, + cu_seqlens_q_padded=cu_seqlens_q_padded, + a2a_input_names=["dout"], ) - - if ctx.enable_mla: - out = out.view(*ctx.v_shape) - dout = dout.view(*ctx.v_shape) - else: - # MHA or GQA - out = out.view(*q.shape) - dout = dout.view(*q.shape) + out = out.view(*ctx.o_shape) + dout = dout.view(*ctx.o_shape) flash_attn_bwd = None if not ctx.use_fused_attention: @@ -2368,10 +2511,11 @@ def backward(ctx, dout, *_args): kv_fp8, ( out - if ctx.fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16 + if (ctx.fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16) + or ctx.fp8_recipe.mxfp8() else out_fp8 ), - dout_fp8, + dout_fp8 if not ctx.fp8_recipe.mxfp8() else dout, softmax_lse, softmax_lse_, rng_states, @@ -2388,16 +2532,20 @@ def backward(ctx, dout, *_args): fused_attn_backend, ctx.softmax_scale, ctx.dropout_p, - qkv_layout, + ctx.qkv_layout, + ctx.qkv_format, + ctx.qkv_format, + ctx.qkv_layout, ctx.attn_mask_type, ctx.attn_bias_type, ctx.deterministic, ctx.fwd_nominal_dtype, bwd_nominal_dtype, - bwd_output_te_dtype, ctx.S_quantizer, dP_quantizer_per_step[i], dQKV_quantizer_per_step[i], + ctx.QKV_quantizer, + ctx.dO_quantizer, ] else: flash_attn_inputs = [ @@ -2471,7 +2619,7 @@ def backward(ctx, dout, *_args): if ctx.fp8 and ctx.use_fused_attention: if ctx.fp8_recipe.delayed(): dq_, dk_, dv_ = [x._data for x in [dq_, dk_, dv_]] - if ctx.fp8_recipe.float8_current_scaling(): + if ctx.fp8_recipe.float8_current_scaling() or ctx.fp8_recipe.mxfp8(): dq_, dk_, dv_ = [x.to(torch.float32) for x in [dq_, dk_, dv_]] # copy dq_ into the right buffer position @@ -2555,7 +2703,7 @@ def backward(ctx, dout, *_args): # dkv correction if ctx.fp8 and ctx.fp8_recipe.delayed(): dkv = dkv_recv_buffer[(rank + i + 1) % cp_size] - elif ctx.fp8 and ctx.fp8_recipe.float8_current_scaling(): + elif ctx.fp8 and (ctx.fp8_recipe.float8_current_scaling() or ctx.fp8_recipe.mxfp8()): dkv = dkv_buffer else: dkv = p2p_comm_buffers[(i + 1) % 2][1] @@ -2645,9 +2793,10 @@ def backward(ctx, dout, *_args): # sum up all cp_size for dq, dk, dv if ctx.fp8 and ctx.use_fused_attention: - amax_cp_bwd = amax_per_step.amax(dim=1) - ctx.dP_quantizer.amax.copy_(amax_cp_bwd[0]) - ctx.dQKV_quantizer.amax.copy_(amax_cp_bwd[1]) + if not ctx.fp8_recipe.mxfp8(): + amax_cp_bwd = amax_per_step.amax(dim=1) + ctx.dP_quantizer.amax.copy_(amax_cp_bwd[0]) + ctx.dQKV_quantizer.amax.copy_(amax_cp_bwd[1]) dq = dq_buffer if ctx.fp8_recipe.delayed(): @@ -2661,7 +2810,7 @@ def backward(ctx, dout, *_args): for x in [dq, dk, dv] ] dq, dk, dv = combine_and_dequantize( - qkv_layout, + ctx.qkv_layout, dq, dk, dv, @@ -2670,7 +2819,7 @@ def backward(ctx, dout, *_args): ) dq, dk, dv = [x.sum(dim=0).to(bwd_nominal_dtype) for x in [dq, dk, dv]] - if ctx.fp8_recipe.float8_current_scaling(): + if ctx.fp8_recipe.float8_current_scaling() or ctx.fp8_recipe.mxfp8(): dk = dkv[: ctx.k_numel].view(ctx.k_shape) dv = dkv[ctx.k_numel :].view(ctx.v_shape) @@ -2686,7 +2835,7 @@ def backward(ctx, dout, *_args): dv[cu_seqlens_kv_padded[-1] :].fill_(0) if ctx.fp8 and ctx.is_input_fp8: - dq, dk, dv = combine_and_quantize(qkv_layout, dq, dk, dv, ctx.dQKV_quantizer) + dq, dk, dv, _, _ = combine_and_quantize(ctx.qkv_layout, dq, dk, dv, ctx.dQKV_quantizer) if ctx.fp8: # print quantizers @@ -2704,7 +2853,8 @@ def backward(ctx, dout, *_args): if cp_size_a2a > 1: if ctx.fp8 and ctx.is_input_fp8: dq_fp8, dk_fp8, dv_fp8 = dq, dk, dv - dq, dk, dv = (dq_fp8._data, dk_fp8._data, dv_fp8._data) + if not ctx.fp8_recipe.mxfp8(): + dq, dk, dv = (dq_fp8._data, dk_fp8._data, dv_fp8._data) chunk_ids_for_a2a = get_seq_chunk_ids_for_reordering_after_attn(cp_size_a2a, q.device) dq, dk, dv = flash_attn_a2a_communicate( [dq, dk, dv], @@ -2714,16 +2864,22 @@ def backward(ctx, dout, *_args): ctx.cp_group_a2a, ctx.cp_stream, False, + qkv_format=ctx.qkv_format, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, + a2a_input_names=["dq", "dk", "dv"], ) - if ctx.fp8 and ctx.is_input_fp8: + if ctx.fp8 and ctx.is_input_fp8 and not ctx.fp8_recipe.mxfp8(): dq, dk, dv = [ Float8Tensor.make_like(x, data=y, dtype=bwd_nominal_dtype) for x, y in zip([dq_fp8, dk_fp8, dv_fp8], [dq, dk, dv]) ] - if ctx.qkv_format == "bshd": - dq, dk, dv = [x.view(ctx.batch_size, -1, *x.shape[-2:]) for x in [dq, dk, dv]] - elif ctx.qkv_format == "sbhd": - dq, dk, dv = [x.view(-1, ctx.batch_size, *x.shape[-2:]) for x in [dq, dk, dv]] + dq, dk, dv = [ + x.view(y) + for x, y in zip( + [dq, dk, dv], [ctx.orig_q_shape, ctx.orig_k_shape, ctx.orig_v_shape] + ) + ] if attn_dbias is not None: # [b, h, sq, 2*cp, sk//(2*cp)] -> [b, h, sq, sk] @@ -2821,27 +2977,42 @@ def forward( cp_group, cp_stream, use_flash_attn_3, + fp8, + fp8_meta, + quantizers, + fp8_output, ): # pylint: disable=missing-function-docstring nvtx_range_push("transformer_engine.AttnFuncWithCPAndKVAllGather.forward") - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) cp_size = get_distributed_world_size(cp_group) rank = get_distributed_rank(cp_group) + qkv_layout = qkv_format + "_" + qkv_format + "_" + qkv_format + o_format = qkv_format + _, seq_dim_qkv, _ = get_bsh_dims(qkv_format) + if softmax_scale is None: + softmax_scale = q.shape[-1] ** (-0.5) - qkv_dtype = q.dtype - - causal = "causal" in attn_mask_type - padding = "padding" in attn_mask_type - assert not padding, f"{attn_mask_type} mask type is not supported!" - if use_fused_attention and causal and "bottom_right" not in attn_mask_type: - attn_mask_type = attn_mask_type + "_bottom_right" - assert attn_bias_type == "no_bias", f"{attn_bias_type} bias type is not supported!" - assert q.shape[-1] % 8 == 0, "Hidden size per attention head should be multiple of 8!" + assert qkv_format != "thd", f"No support for cp_comm_type='all_gather' and {qkv_format=}." + assert ( + "padding" not in attn_mask_type + ), f"No support for cp_comm_type='all_gather' and {attn_mask_type=}." + assert ( + attn_bias_type == "no_bias" + ), f"No support for cp_comm_type='all_gather' and {attn_bias_type=}." assert ( - use_fused_attention or fa_utils.v2_3_plus - ), "Sliding window attention only can work with FusedAttention or FlashAttention >= 2.3!" + window_size == (-1, 0) + or window_size == (-1, -1) + or use_fused_attention + or fa_utils.v2_3_plus + ), ( + "cp_comm_type='all_gather' only supports SWA through FusedAttention or FlashAttention" + f" >= 2.3. Found {use_fused_attention=} and {fa_utils.v2_3_plus=}." + ) + assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, ( + "cp_comm_type='all_gather' requires seq_len % 2 == 0 for Q, K, V. Found seq_len_q =" + f" {q.shape[seq_dim_qkv]}, seq_len_kv = {k.shape[seq_dim_qkv]}." + ) flash_attn_fwd = None if not use_fused_attention: @@ -2874,14 +3045,6 @@ def forward( if fa_utils.v2_6_0_plus: fa_forward_kwargs["softcap"] = 0.0 - assert qkv_format != "thd", f"{qkv_format} format is not supported!" - qkv_layout = qkv_format + "_" + qkv_format + "_" + qkv_format - - seq_dim = qkv_format.index("s") - assert ( - q.shape[seq_dim] % 2 == 0 and k.shape[seq_dim] % 2 == 0 - ), "Sequence length per GPU needs to be divisible by 2!" - max_seqlen_q = max_seqlen_q // (2 * cp_size) max_seqlen_kv = max_seqlen_kv // (2 * cp_size) if use_fused_attention or qkv_format == "thd": @@ -2890,30 +3053,90 @@ def forward( cu_seqlens_q_padded = cu_seqlens_q_padded // (2 * cp_size) else: cu_seqlens_q_padded = None + if use_fused_attention and attn_mask_type == "causal": + attn_mask_type = attn_mask_type + "_bottom_right" + causal = "causal" in attn_mask_type - # [b, s, h, d] -> [b, 2, s//2, h, d] or [s, b, h, d] -> [2, s//2, b, h, d] - q = q.view(*q.shape[:seq_dim], 2, q.shape[seq_dim] // 2, *q.shape[(seq_dim + 1) :]) - # [b, s, h, d] or [s, b, h, d] -> [s, b, h, d] - k, v = [x.movedim(seq_dim, 0).contiguous() for x in [k, v]] + # FP8 setup + assert isinstance(k, q.__class__) and isinstance( + v, q.__class__ + ), "q, k, v must be of the same class, e.g. torch.Tensor or QuantizedTensorStorage." + is_input_fp8 = isinstance(q, QuantizedTensorStorage) + is_output_fp8 = fp8_output + is_bwd_fp8 = int(os.getenv("NVTE_FP8_DPA_BWD", "1")) + fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() + if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None: + fp8_recipe = fp8_meta["local_recipes"][0] + ( + QKV_quantizer, + O_quantizer, + S_quantizer, + dQKV_quantizer, + dO_quantizer, + dP_quantizer, + ) = dpa_utils.get_attention_quantizers(fp8, fp8_recipe, quantizers) + fwd_nominal_dtype = q.dtype + q_fp8, k_fp8, v_fp8 = (q, k, v) if is_input_fp8 else (None, None, None) + q_f16, k_f16, v_f16 = (None, None, None) if is_input_fp8 else (q, k, v) + fused_attn_backend = None + fp8_meta_kwargs = {} + if fp8: + assert use_fused_attention, "FP8 is only supported with FusedAttention backend!" + fused_attn_backend = tex.NVTE_Fused_Attn_Backend.NVTE_FP8 + if not is_input_fp8 and not fp8_recipe.mxfp8(): + q_fp8, k_fp8, v_fp8, qkv_layout, _ = combine_and_quantize( + qkv_layout, q, k, v, QKV_quantizer + ) + if not fp8_recipe.mxfp8(): + q, k, v = [q_fp8._data, k_fp8._data, v_fp8._data] + fp8_meta_kwargs["s_quantizer"] = S_quantizer + fp8_meta_kwargs["o_quantizer"] = O_quantizer + elif use_fused_attention: + fused_attn_backend = tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen + orig_q_shape, _, orig_v_shape = q.shape, k.shape, v.shape + orig_o_shape = orig_q_shape[:-1] + orig_v_shape[-1:] + + # q, k, v: + # FP8DS/CS: torch.uint8 + # MXFP8/F16: torch.float16 or torch.bfloat16 + # reshape: split s + # [b, s, h, d] -> [b, 2, s//2, h, d] + # [s, b, h, d] -> [2, s//2, b, h, d] + q = q.view( + *q.shape[:seq_dim_qkv], 2, q.shape[seq_dim_qkv] // 2, *q.shape[(seq_dim_qkv + 1) :] + ) + # s dim first for all-gather + # [b, s, h, d]/[s, b, h, d] -> [s, b, h, d] + k, v = [x.movedim(seq_dim_qkv, 0).contiguous() for x in [k, v]] - # [s, b, h, d] -> [cp, s, b, h, d] + # gather along s: [s, b, h, d] -> [cp, s, b, h, d] k_ag, _ = gather_along_first_dim(k, cp_group) v_ag, _ = gather_along_first_dim(v, cp_group) - - # [cp, s, b, h, d] -> [cp*2, s//2, b, h, d] + # split s:[cp, s, b, h, d] -> [cp*2, s//2, b, h, d] k_ag = k_ag.view(2 * cp_size, k.shape[0] // 2, *k.shape[1:]) v_ag = v_ag.view(2 * cp_size, v.shape[0] // 2, *v.shape[1:]) + # pick out specific chunks for each rank chunk_ids_for_kv_ag = get_seq_chunk_ids_for_reordering_before_attn(cp_size, k.device) k_ag = torch.index_select(k_ag, dim=0, index=chunk_ids_for_kv_ag) v_ag = torch.index_select(v_ag, dim=0, index=chunk_ids_for_kv_ag) - # [cp*2, s//2, b, h, d] -> [cp*s, b, h, d] + # reshape/flatten: [cp*2, s//2, b, h, d] -> [cp*s, b, h, d] k_ag = k_ag.view(-1, *k.shape[1:]) v_ag = v_ag.view(-1, *v.shape[1:]) cp_stream.wait_stream(torch.cuda.current_stream()) + # q: [b, 2, s//2, h, d] or [2, s//2, b, h, d] + # k: [s, b, h, d] + # v: [s, b, h, d] + # k_ag: [cp*s, b, h, d] + # v_ag: [cp*s, b, h, d] + # out_f16: [b, 2, s//2, h, d] or [2, s//2, b, h, d] + q_shape, k_shape, v_shape = q.shape, k.shape, v.shape + o_shape = q.shape[:-1] + v.shape[-1:] + out_f16 = torch.empty(o_shape, dtype=fwd_nominal_dtype, device=q.device) + # create two streams to resolve wave quantization issue of Flash Attn in each step flash_attn_streams = [torch.cuda.current_stream(), cp_stream] - + # prepare per-step tensors local_seq_chunk_ids = [rank, 2 * cp_size - rank - 1] kv_seq_range_per_step = [None, None] window_size_per_step = [None, None] @@ -2921,16 +3144,15 @@ def forward( out_per_step = [None, None] softmax_lse_per_step = [None, None] rng_states = [None, None] - out = torch.empty_like(q) max_logit_per_step = [None, None] max_logit = None for i in range(len(local_seq_chunk_ids) + 1): if i < len(local_seq_chunk_ids): with torch.cuda.stream(flash_attn_streams[i]): - # [b, 2, sq//2, h, d] -> [b, sq//2, h, d] - # or [2, sq//2, b, h, d] -> [sq//2, b, h, d] - q_ = q.select(seq_dim, i).contiguous() + # [b, 2, s//2, h, d] -> [b, s//2, h, d] + # [2, s//2, b, h, d] -> [s//2, b, h, d] + q_part = q.select(seq_dim_qkv, i).contiguous() kv_seq_range_per_step[i], window_size_per_step[i] = ( get_kv_seq_info_after_all_gather( local_seq_chunk_ids[i], @@ -2950,13 +3172,30 @@ def forward( cu_seqlens_kv_per_step[i] = dpa_utils.get_full_cu_seqlens( k.shape[1], max_seqlen_kv_, k.device ) - k_, v_ = [x[seq_start_idx:seq_end_idx] for x in [k_ag, v_ag]] - # [s_range, b, h, d] -> [b, s_range, h, d] or [s_range, b, h, d] - k_, v_ = [x.movedim(0, seq_dim).contiguous() for x in [k_, v_]] + # select range: [s_range, b, h, d] + k_part, v_part = [x[seq_start_idx:seq_end_idx] for x in [k_ag, v_ag]] + # reshape to original format: [b, s_range, h, d] or [s_range, b, h, d] + k_part, v_part = [ + x.movedim(0, seq_dim_qkv).contiguous() for x in [k_part, v_part] + ] if use_fused_attention: + new_qkv_layout = qkv_layout + qkv_scale_inv_format = None + if fp8: + if not fp8_recipe.mxfp8(): + q_part, k_part, v_part = [ + Float8Tensor.make_like(x, data=y, dtype=fwd_nominal_dtype) + for x, y in zip([q_fp8, k_fp8, v_fp8], [q_part, k_part, v_part]) + ] + else: + q_part, k_part, v_part, new_qkv_layout, qkv_scale_inv_format = ( + combine_and_quantize( + qkv_layout, q_part, k_part, v_part, QKV_quantizer + ) + ) ( out_per_step[i], - [softmax_lse_per_step[i], rng_states[i]], + aux_ctx_tensors, *max_logit_, ) = fused_attn_fwd( is_training, @@ -2964,14 +3203,15 @@ def forward( max_seqlen_kv_, cu_seqlens_q, cu_seqlens_kv_per_step[i], - q_, - k_, - v_, - qkv_dtype, - tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, + q_part, + k_part, + v_part, + fwd_nominal_dtype, + fused_attn_backend, attn_scale=softmax_scale, dropout=dropout_p, - qkv_layout=qkv_layout, + qkv_layout=new_qkv_layout, + o_format=o_format, attn_mask_type=attn_mask_type, attn_bias_type=attn_bias_type, attn_bias=attn_bias, @@ -2980,9 +3220,20 @@ def forward( window_size=window_size_per_step[i], return_max_logit=return_max_logit, cuda_graph=is_graph_capturing(), + qkv_scale_inv_format=qkv_scale_inv_format, + **fp8_meta_kwargs, ) + if fp8: + if qkv_layout != "t3hd": + softmax_lse_per_step[i], rng_states[i] = aux_ctx_tensors + else: + softmax_lse_per_step[i], _, rng_states[i] = aux_ctx_tensors + else: + softmax_lse_per_step[i], rng_states[i], *_ = aux_ctx_tensors if return_max_logit: max_logit_per_step[i] = max_logit_[0] + if fp8 and isinstance(out_per_step[i], QuantizedTensorStorage): + out_per_step[i] = out_per_step[i].dequantize(dtype=fwd_nominal_dtype) else: fa_forward_args_thd = get_fa_args( True, @@ -2999,9 +3250,9 @@ def forward( fa_forward_kwargs["window_size_left"] = window_size_per_step[i][0] fa_forward_kwargs["window_size_right"] = window_size_per_step[i][1] fa_outputs = flash_attn_fwd( - q_, - k_, - v_, + q_part, + k_part, + v_part, *fa_forward_args_thd, causal=causal, **fa_forward_kwargs, @@ -3017,61 +3268,152 @@ def forward( if not use_flash_attn_3: rng_states[i] = fa_outputs[3] + # out_per_step[i]: fwd_nominal_dtype, [b, s//2, h, d] or [s//2, b, h, d] + # out_f16: fwd_nominal_dtype, [b, 2, s//2, h, d] or [2, s//2, b, h, d] + # max_logit_per_step[i]: torch.float32, [h] + # max_logit: torch.float32, [h] if return_max_logit and i == 0: max_logit = torch.clone(max_logit_per_step[0]) if i > 0: with torch.cuda.stream(flash_attn_streams[i - 1]): - if qkv_format == "bshd": - out[:, i - 1].copy_(out_per_step[i - 1]) - elif qkv_format == "sbhd": - out[i - 1].copy_(out_per_step[i - 1]) + if o_format == "bshd": + out_f16[:, i - 1].copy_(out_per_step[i - 1]) + elif o_format == "sbhd": + out_f16[i - 1].copy_(out_per_step[i - 1]) if return_max_logit: max_logit = torch.maximum(max_logit, max_logit_per_step[i - 1]) torch.cuda.current_stream().wait_stream(cp_stream) + + # all reduce max_logit across ranks if return_max_logit: torch.distributed.all_reduce( max_logit, op=torch.distributed.ReduceOp.MAX, group=cp_group ) - if use_fused_attention: - if qkv_format == "bshd": - out = out.view(out.shape[0], -1, *out.shape[-2:]) - elif qkv_format == "sbhd": - out = out.view(-1, *out.shape[-3:]) - else: - out = out.view(-1, *out.shape[-2:]) + # out_f16: fwd_nominal_dtype + # [b, 2, s//2, h, d] -> [b, s, h, d] + # [2, s//2, b, h, d] -> [s, b, h, d] + out_f16 = out_f16.view(orig_o_shape) - ctx.save_for_backward( - q, - k, - v, + # prepare for forward output and backward saves of out + out_fp8 = None + bwd_requires_o_fp8 = ( + is_training + and is_bwd_fp8 + and ( + fp8_recipe.delayed() + or (fp8_recipe.float8_current_scaling() and not _dpa_fp8_cs_o_in_f16) + ) + ) + if fp8 and (is_output_fp8 or bwd_requires_o_fp8): + out_fp8 = O_quantizer(out_f16) + out_ret = out_fp8 if is_output_fp8 else out_f16 + + # save tensors for backward + ctx.fp8 = fp8 and is_bwd_fp8 + ctx.fp8_recipe = fp8_recipe + fp8_tensors = (None, None, None, None) + f16_tensors = (None, None, None, None) + # True: q split along s; k/v with s first, i.e. [s, b, h, d] + # False: original [b, s, h, d] or [s, b, h, d] + ctx.qkv_reshaped = True + # no load-balance related token shuffling; original token order in q/k/v/out_f16 + # q: [b, 2, s//2, h, d] or [2, s//2, b, h, d] + # k: [s, b, h, d] + # v: [s, b, h, d] + # out_f16/out_fp8: [b, s, h, d] or [s, b, h, d] + if ctx.fp8: + # q_fp8_save: [b, 2, s//2, h, d] or [2, s//2, b, h, d] + # k_fp8_save: [s, b, h, d] + # v_fp8_save: [s, b, h, d] + q_fp8_save, k_fp8_save, v_fp8_save = None, None, None + if fp8_recipe.delayed() or fp8_recipe.float8_current_scaling(): + q_fp8_save = Float8Tensor.make_like(q_fp8, data=q, dtype=fwd_nominal_dtype) + k_fp8_save = Float8Tensor.make_like(k_fp8, data=k, dtype=fwd_nominal_dtype) + v_fp8_save = Float8Tensor.make_like(v_fp8, data=v, dtype=fwd_nominal_dtype) + # FP8DS or (FP8CS+not _dpa_fp8_cs_o_in_f16): q/k/v/o all in FP8 + # FP8CS+_dpa_fp8_cs_o_in_f16: q/k/v in FP8, o in f16 + # MXFP8: q/k/v/o all in f16 + if fp8_recipe.delayed() or ( + fp8_recipe.float8_current_scaling() and not _dpa_fp8_cs_o_in_f16 + ): + fp8_tensors = (q_fp8_save, k_fp8_save, v_fp8_save, out_fp8) + elif fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16: + fp8_tensors = (q_fp8_save, k_fp8_save, v_fp8_save, None) + f16_tensors = (None, None, None, out_f16) + elif fp8_recipe.mxfp8(): + f16_tensors = (q, k, v, out_f16) + elif fp8: + # convert q/k/v to F16 if necessary, and save q/k/v/o all in F16 and original format + if is_input_fp8: + q_f16, k_f16, v_f16 = combine_and_dequantize(qkv_layout, q_fp8, k_fp8, v_fp8) + f16_tensors = (q_f16, k_f16, v_f16, out_f16) + ctx.qkv_reshaped = False + else: + # save all in F16 + # q: [b, 2, s//2, h, d] or [2, s//2, b, h, d] + # k: [s, b, h, d] + # v: [s, b, h, d] + # out_f16: [b, s, h, d] or [s, b, h, d] + f16_tensors = (q, k, v, out_f16) + tensors_to_save, tensor_objects = prepare_for_saving( + *fp8_tensors, + *f16_tensors, cu_seqlens_q, cu_seqlens_q_padded, *cu_seqlens_kv_per_step, - *out_per_step, *softmax_lse_per_step, *rng_states, ) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects - ctx.qkv_dtype = qkv_dtype + ctx.qkv_format = qkv_format + ctx.qkv_layout = qkv_layout + ctx.o_format = o_format + ctx.dqkv_format = qkv_format + ctx.dqkv_layout = qkv_layout + ctx.fwd_nominal_dtype = fwd_nominal_dtype + ctx.q_shape = q_shape + ctx.k_shape = k_shape + ctx.v_shape = v_shape + ctx.o_shape = o_shape ctx.kv_seq_range_per_step = kv_seq_range_per_step ctx.window_size_per_step = window_size_per_step + ctx.cp_group = cp_group ctx.cp_stream = cp_stream ctx.dropout_p = dropout_p ctx.max_seqlen_q = max_seqlen_q ctx.softmax_scale = softmax_scale - ctx.qkv_format = qkv_format ctx.attn_bias_type = attn_bias_type ctx.attn_mask_type = attn_mask_type ctx.deterministic = deterministic ctx.use_fused_attention = use_fused_attention ctx.use_flash_attn_3 = use_flash_attn_3 + ctx.fp8_meta = fp8_meta + ctx.is_input_fp8 = is_input_fp8 + + ctx.dQKV_quantizer = dQKV_quantizer + ctx.dO_quantizer = dO_quantizer + ctx.dP_quantizer = dP_quantizer + ctx.QKV_quantizer = QKV_quantizer + ctx.O_quantizer = O_quantizer + ctx.S_quantizer = S_quantizer + if ctx.fp8: + ctx.QKV_quantizer = QKV_quantizer.copy() + ctx.O_quantizer = O_quantizer.copy() + ctx.S_quantizer = S_quantizer.copy() if S_quantizer is not None else None + if not ctx.fp8_recipe.mxfp8(): + ctx.QKV_quantizer.scale = QKV_quantizer.scale.clone() + ctx.O_quantizer.scale = O_quantizer.scale.clone() + ctx.S_quantizer.scale = S_quantizer.scale.clone() + nvtx_range_pop("transformer_engine.AttnFuncWithCPAndKVAllGather.forward") if return_max_logit: - return out, max_logit - return out + return out_ret, max_logit + return out_ret @staticmethod def backward(ctx, dout, *_args): @@ -3080,22 +3422,94 @@ def backward(ctx, dout, *_args): cp_size = get_distributed_world_size(ctx.cp_group) rank = get_distributed_rank(ctx.cp_group) - (*saved_tensors,) = ctx.saved_tensors - (q, k, v, cu_seqlens_q, cu_seqlens_q_padded) = saved_tensors[:5] - cu_seqlens_kv_per_step = saved_tensors[5:7] - out_per_step = saved_tensors[7:9] - softmax_lse_per_step = saved_tensors[9:11] - rng_states = saved_tensors[11:13] + cu_seqlens_kv_per_step = [None, None] + softmax_lse_per_step = [None, None] + rng_states = [None, None] + ( + q_fp8, + k_fp8, + v_fp8, + out_fp8, + q, + k, + v, + out, + cu_seqlens_q, + cu_seqlens_q_padded, + cu_seqlens_kv_per_step[0], + cu_seqlens_kv_per_step[1], + softmax_lse_per_step[0], + softmax_lse_per_step[1], + rng_states[0], + rng_states[1], + ) = restore_from_func_ctx(ctx) kv_seq_range_per_step = ctx.kv_seq_range_per_step window_size_per_step = ctx.window_size_per_step - seq_dim = ctx.qkv_format.index("s") - qkv_layout = ctx.qkv_format + "_" + ctx.qkv_format + "_" + ctx.qkv_format + _, seq_dim_qkv, _ = get_bsh_dims(ctx.qkv_format) + _, seq_dim_dqkv, _ = get_bsh_dims(ctx.dqkv_format) + _, seq_dim_o, _ = get_bsh_dims(ctx.o_format) + causal = "causal" in ctx.attn_mask_type - dout = dout.view(q.shape) - dq = torch.empty_like(q) - dk = torch.zeros((k.shape[0] * cp_size, *k.shape[1:]), dtype=k.dtype, device=k.device) - dv = torch.zeros_like(dk) + # set up dout: + # FP8DS/CS: torch.uint8, [b, s, h, d] or [s, b, h, d] + # MXFP8/F16: torch.float16 or torch.bfloat16, [b, s, h, d] or [s, b, h, d] + dout_fp8 = None + if ctx.fp8: + assert ctx.use_fused_attention, "FP8 is only supported with FusedAttention backend!" + if isinstance(dout, QuantizedTensorStorage): + dout_fp8 = dout + elif not ctx.fp8_recipe.mxfp8(): + dout = ctx.dO_quantizer(dout) + dout_fp8 = dout + if not ctx.fp8_recipe.mxfp8(): + dout = dout_fp8._data + # [b, s, h, d] -> [b, 2, s//2, h, d] + # [s, b, h, d] -> [2, s//2, b, h, d] + dout = dout.view(ctx.o_shape) + + # set up q, k, v: + # FP8DS/CS: torch.uint8 + # MXFP8/F16: torch.float16 or torch.bfloat16 + # q: [b, 2, s//2, h, d] or [2, s//2, b, h, d] + # k: [s, b, h, d] + # v: [s, b, h, d] + if ctx.fp8 and not ctx.fp8_recipe.mxfp8(): + q, k, v = [x._data for x in [q_fp8, k_fp8, v_fp8]] + if not ctx.qkv_reshaped: + q = q.view( + *q.shape[:seq_dim_qkv], 2, q.shape[seq_dim_qkv] // 2, *q.shape[(seq_dim_qkv + 1) :] + ) + k, v = [x.movedim(seq_dim_qkv, 0).contiguous() for x in [k, v]] + + # set up out: + # FP8DS or (FP8CS+not _dpa_fp8_cs_o_in_f16): torch.uint8 + # FP8CS+_dpa_fp8_cs_o_in_f16: torch.float16 or torch.bfloat16 + # MXFP8/F16: torch.float16 or torch.bfloat16 + # [b, s, h, d] -> [b, 2, s//2, h, d] + # [s, b, h, d] -> [2, s//2, b, h, d] + if ctx.fp8 and ( + ctx.fp8_recipe.delayed() + or (ctx.fp8_recipe.float8_current_scaling() and not _dpa_fp8_cs_o_in_f16) + ): + out = out_fp8._data + out = out.view(ctx.o_shape) + + # set up dq, dk, dv: + # dq: fwd_nominal_dtype, [b, 2, s//2, h, d] or [2, s//2, b, h, d] + # dk: fwd_nominal_dtype, [cp*s, b, h, d] + # dv: fwd_nominal_dtype, [cp*s, b, h, d] + dq = torch.empty(ctx.q_shape, dtype=ctx.fwd_nominal_dtype, device=q.device) + dk = torch.zeros( + (ctx.k_shape[0] * cp_size, *ctx.k_shape[1:]), + dtype=ctx.fwd_nominal_dtype, + device=k.device, + ) + dv = torch.zeros( + (ctx.v_shape[0] * cp_size, *ctx.v_shape[1:]), + dtype=ctx.fwd_nominal_dtype, + device=v.device, + ) dq_per_step = [None, None] dk_per_step = [None, None] dv_per_step = [None, None] @@ -3105,23 +3519,22 @@ def backward(ctx, dout, *_args): # synchronize dkv update across steps dkv_update_done = torch.cuda.Event() - # [s, b, h, d] -> [cp, s, b, h, d] + # gather k and v along s: [s, b, h, d] -> [cp, s, b, h, d] k_ag, _ = gather_along_first_dim(k, ctx.cp_group) v_ag, _ = gather_along_first_dim(v, ctx.cp_group) - - # [cp, s, b, h, d] -> [cp*2, s//2, b, h, d] + # split s: [cp, s, b, h, d] -> [cp*2, s//2, b, h, d] k_ag = k_ag.view(2 * cp_size, k.shape[0] // 2, *k.shape[1:]) v_ag = v_ag.view(2 * cp_size, v.shape[0] // 2, *v.shape[1:]) + # select appropriate chunks for each rank chunk_ids_for_kv_ag = get_seq_chunk_ids_for_reordering_before_attn(cp_size, k.device) k_ag = torch.index_select(k_ag, dim=0, index=chunk_ids_for_kv_ag) v_ag = torch.index_select(v_ag, dim=0, index=chunk_ids_for_kv_ag) - # [cp*2, s//2, b, h, d] -> [cp*s, b, h, d] + # flatten: [cp*2, s//2, b, h, d] -> [cp*s, b, h, d] k_ag = k_ag.view(-1, *k.shape[1:]) v_ag = v_ag.view(-1, *v.shape[1:]) ctx.cp_stream.wait_stream(torch.cuda.current_stream()) - local_seq_chunk_ids = [rank, 2 * cp_size - rank - 1] - + # set up flash_attn_bwd flash_attn_bwd = None if not ctx.use_fused_attention: fa_backward_kwargs = {"softmax_scale": ctx.softmax_scale} @@ -3153,57 +3566,132 @@ def backward(ctx, dout, *_args): if fa_utils.v2_6_0_plus: fa_backward_kwargs["softcap"] = 0.0 + local_seq_chunk_ids = [rank, 2 * cp_size - rank - 1] for i in range(len(local_seq_chunk_ids) + 1): if i < len(local_seq_chunk_ids): with torch.cuda.stream(flash_attn_streams[i]): - # [b, 2, sq//2, h, d] -> [b, sq//2, h, d] - # or [2, sq//2, b, h, d] -> [sq//2, b, h, d] - q_ = q.select(seq_dim, i).contiguous() + # [b, 2, s//2, h, d] -> [b, s//2, h, d] + # [2, s//2, b, h, d] -> [s//2, b, h, d] + q_part = q.select(seq_dim_qkv, i).contiguous() seq_start_idx, seq_end_idx = ( kv_seq_range_per_step[i][0], kv_seq_range_per_step[i][1], ) max_seqlen_kv = seq_end_idx - seq_start_idx - k_, v_ = [x[seq_start_idx:seq_end_idx] for x in [k_ag, v_ag]] - # [cp*s, b, h, d] -> [b, s_range, h, d] or [s_range, b, h, d] - k_, v_ = [x.movedim(0, seq_dim).contiguous() for x in [k_, v_]] - out_ = out_per_step[i] - dout_ = dout.select(seq_dim, i).contiguous().view(out_.shape) + # select range: [s_range, b, h, d] + k_part, v_part = [x[seq_start_idx:seq_end_idx] for x in [k_ag, v_ag]] + # reshape to original format: [b, s_range, h, d] or [s_range, b, h, d] + k_part, v_part = [ + x.movedim(0, seq_dim_qkv).contiguous() for x in [k_part, v_part] + ] + # [b, 2, s//2, h, d] -> [b, s//2, h, d] + # [2, s//2, b, h, d] -> [s//2, b, h, d] + out_part = out.select(seq_dim_o, i).contiguous() + dout_part = dout.select(seq_dim_o, i).contiguous() if ctx.use_fused_attention: - aux_ctx_tensors = [softmax_lse_per_step[i], rng_states[i]] + if ctx.fp8 and ctx.qkv_layout == "t3hd": + aux_ctx_tensors = [ + softmax_lse_per_step[i], + softmax_lse_per_step[i], + rng_states[i], + ] + else: + aux_ctx_tensors = [ + softmax_lse_per_step[i], + rng_states[i], + ] + fused_attn_backend = tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen + fp8_meta_kwargs = {} + new_qkv_layout = ctx.qkv_layout + do_format = ctx.o_format + qkv_scale_inv_format = None + do_scale_inv_format = None + if ctx.fp8: + fused_attn_backend = tex.NVTE_Fused_Attn_Backend.NVTE_FP8 + fp8_meta_kwargs["s_quantizer"] = ctx.S_quantizer + fp8_meta_kwargs["dp_quantizer"] = ctx.dP_quantizer + fp8_meta_kwargs["dqkv_quantizer"] = ctx.dQKV_quantizer + # FP8DS or (FP8CS+not _dpa_fp8_cs_o_in_f16): q/k/v/o/do all in FP8 + # FP8CS+_dpa_fp8_cs_o_in_f16: q/k/v/do in FP8, o in f16 + # MXFP8: q/k/v/do all in MXFP8, o/do_f16 in F16 + if not ctx.fp8_recipe.mxfp8(): + q_part, k_part, v_part = [ + Float8Tensor.make_like(x, data=y, dtype=ctx.fwd_nominal_dtype) + for x, y in zip([q_fp8, k_fp8, v_fp8], [q_part, k_part, v_part]) + ] + if ctx.fp8_recipe.delayed() or ( + ctx.fp8_recipe.float8_current_scaling() + and not _dpa_fp8_cs_o_in_f16 + ): + out_part = Float8Tensor.make_like( + out_fp8, data=out_part, dtype=ctx.fwd_nominal_dtype + ) + dout_part = Float8Tensor.make_like( + dout_fp8, data=dout_part, dtype=ctx.fwd_nominal_dtype + ) + else: + q_part, k_part, v_part, new_qkv_layout, qkv_scale_inv_format = ( + combine_and_quantize( + ctx.qkv_layout, + q_part, + k_part, + v_part, + ctx.QKV_quantizer, + used_in_forward=False, + used_in_backward=True, + ) + ) + aux_ctx_tensors.append(dout_part) + (dout_part,), do_scale_inv_format = mxfp8_quantize_fast_path( + [(dout_part, ctx.dO_quantizer)], + do_format, + ) dq_per_step[i], dk_per_step[i], dv_per_step[i], *_ = fused_attn_bwd( ctx.max_seqlen_q, max_seqlen_kv, cu_seqlens_q, cu_seqlens_kv_per_step[i], - q_, - k_, - v_, - out_, - dout_, - ctx.qkv_dtype, - TE_DType[dout.dtype], + q_part, + k_part, + v_part, + out_part, + dout_part, + ctx.fwd_nominal_dtype, aux_ctx_tensors, - tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, + fused_attn_backend, cu_seqlens_q_padded=cu_seqlens_q_padded, cu_seqlens_kv_padded=cu_seqlens_kv_per_step[i], attn_scale=ctx.softmax_scale, dropout=ctx.dropout_p, - qkv_layout=qkv_layout, + qkv_layout=new_qkv_layout, + o_format=ctx.o_format, + do_format=do_format, + dqkv_layout=ctx.dqkv_layout, attn_mask_type=ctx.attn_mask_type, attn_bias_type=ctx.attn_bias_type, window_size=window_size_per_step[i], deterministic=ctx.deterministic, cuda_graph=is_graph_capturing(), + qkv_scale_inv_format=qkv_scale_inv_format, + do_scale_inv_format=do_scale_inv_format, + **fp8_meta_kwargs, ) + if ctx.fp8 and all( + isinstance(x, QuantizedTensorStorage) + for x in [dq_per_step[i], dk_per_step[i], dv_per_step[i]] + ): + dq_per_step[i], dk_per_step[i], dv_per_step[i] = [ + x.dequantize(dtype=ctx.fwd_nominal_dtype) + for x in [dq_per_step[i], dk_per_step[i], dv_per_step[i]] + ] else: dq_per_step[i], dk_per_step[i], dv_per_step[i] = [ - torch.empty_like(x) for x in [q_, k_, v_] + torch.empty_like(x) for x in [q_part, k_part, v_part] ] fa_backward_args_thd = get_fa_args( False, ctx.use_flash_attn_3, - ctx.qkv_format, + ctx.dqkv_format, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv_per_step[i], max_seqlen_q=ctx.max_seqlen_q, @@ -3220,29 +3708,34 @@ def backward(ctx, dout, *_args): fa_backward_kwargs["window_size_left"] = window_size_per_step[i][0] fa_backward_kwargs["window_size_right"] = window_size_per_step[i][1] if ctx.use_flash_attn_3: - fa_backward_kwargs["is_causal"] = "causal" in ctx.attn_mask_type + fa_backward_kwargs["is_causal"] = causal else: - fa_backward_kwargs["causal"] = "causal" in ctx.attn_mask_type + fa_backward_kwargs["causal"] = causal flash_attn_bwd( - dout_, - q_, - k_, - v_, - out_, + dout_part, + q_part, + k_part, + v_part, + out_part, softmax_lse_per_step[i], *fa_backward_args_thd, **fa_backward_kwargs, ) if i > 0: + # dq/dk/dv, dq_per_step/dk_per_step/dv_per_step: ctx.fwd_nominal_dtype with torch.cuda.stream(flash_attn_streams[i - 1]): - if ctx.qkv_format == "bshd": + # dq: [b, 2, s//2, h, d] or [2, s//2, b, h, d] + # dq_per_step[i]: [b, s//2, h, d] or [s//2, b, h, d] + if ctx.dqkv_format == "bshd": dq[:, i - 1].copy_(dq_per_step[i - 1]) - elif ctx.qkv_format == "sbhd": + elif ctx.dqkv_format == "sbhd": dq[i - 1].copy_(dq_per_step[i - 1]) - # [b, s_range, h, d] or [s_range, b, h, d] -> [s_range, b, h, d] + # dk/dv: [cp*s, b, h, d] + # dk_per_step[i - 1]/dv_per_step[i - 1]: [s_range, b, h, d] or [b, s_range, h, d] + # move s to first dim: [s_range, b, h, d] dk_per_step[i - 1], dv_per_step[i - 1] = [ - x.movedim(seq_dim, 0).contiguous() + x.movedim(seq_dim_dqkv, 0).contiguous() for x in [dk_per_step[i - 1], dv_per_step[i - 1]] ] # wait until dkv update of last step is done @@ -3252,6 +3745,7 @@ def backward(ctx, dout, *_args): kv_seq_range_per_step[i - 1][0], kv_seq_range_per_step[i - 1][1], ) + # add to dk/dv: [cp*s, b, h, d] dk[seq_start_idx:seq_end_idx].add_(dk_per_step[i - 1]) dv[seq_start_idx:seq_end_idx].add_(dv_per_step[i - 1]) if i < len(local_seq_chunk_ids): @@ -3259,23 +3753,33 @@ def backward(ctx, dout, *_args): torch.cuda.current_stream().wait_stream(ctx.cp_stream) - # [cp*s, b, h, d] -> [cp*2, s//2, b, h, d] + # split s:[cp*s, b, h, d] -> [cp*2, s//2, b, h, d] dk = dk.view(2 * cp_size, -1, *dk.shape[-3:]) dv = dv.view(2 * cp_size, -1, *dv.shape[-3:]) + # put back together the right chunks for each rank chunk_ids_for_kv_ag = get_seq_chunk_ids_for_reordering_after_attn(cp_size, dk.device) dk = torch.index_select(dk, dim=0, index=chunk_ids_for_kv_ag) dv = torch.index_select(dv, dim=0, index=chunk_ids_for_kv_ag) - # [cp*2, s//2, b, h, d] -> [cp*s, b, h, d] + # flatten: [cp*2, s//2, b, h, d] -> [cp*s, b, h, d] dk = dk.view(-1, *dk.shape[-3:]) dv = dv.view(-1, *dv.shape[-3:]) + # reduce scatter: [cp*s, b, h, d] -> [s, b, h, d] dk, _ = reduce_scatter_along_first_dim(dk, ctx.cp_group) dv, _ = reduce_scatter_along_first_dim(dv, ctx.cp_group) - dq = dq.view(*dq.shape[:seq_dim], -1, *dq.shape[(seq_dim + 2) :]) - dk = dk.movedim(0, seq_dim).contiguous() - dv = dv.movedim(0, seq_dim).contiguous() - nvtx_range_pop("transformer_engine.AttnFuncWithCPAndKVAllGather.backward") + # reshape to original format: + # dq: [b, 2, s//2, h, d] or [2, s//2, b, h, d] -> [b, s, h, d] or [s, b, h, d] + # dk: [s, b, h, d] -> [b, s, h, d] or [s, b, h, d] + # dv: [s, b, h, d] -> [b, s, h, d] or [s, b, h, d] + dq = dq.view(*dq.shape[:seq_dim_dqkv], -1, *dq.shape[(seq_dim_dqkv + 2) :]) + dk = dk.movedim(0, seq_dim_dqkv).contiguous() + dv = dv.movedim(0, seq_dim_dqkv).contiguous() + # quantize if necessary + if ctx.fp8 and ctx.is_input_fp8: + dq, dk, dv, _, _ = combine_and_quantize(ctx.dqkv_layout, dq, dk, dv, ctx.dQKV_quantizer) + + nvtx_range_pop("transformer_engine.AttnFuncWithCPAndKVAllGather.backward") return ( None, dq, @@ -3298,6 +3802,10 @@ def backward(ctx, dout, *_args): None, None, None, + None, + None, + None, + None, ) @@ -3342,24 +3850,43 @@ def forward( ): # pylint: disable=missing-function-docstring nvtx_range_push("transformer_engine.AttnFuncWithCPAndQKVOA2A.forward") - if softmax_scale is None: - softmax_scale = q.shape[-1] ** (-0.5) cp_size = get_distributed_world_size(cp_group) - + qkv_layout = qkv_format + "_" + qkv_format + "_" + qkv_format + original_qkv_layout = qkv_layout + orig_q_shape, orig_k_shape, orig_v_shape = q.shape, k.shape, v.shape + orig_o_shape = orig_q_shape[:-1] + orig_v_shape[-1:] + o_format = qkv_format + _, seq_dim_qkv, _ = get_bsh_dims(qkv_format) + _, seq_dim_o, _ = get_bsh_dims(o_format) + if softmax_scale is None: + softmax_scale = q.shape[-1] ** (-0.5) causal = "causal" in attn_mask_type - padding = "padding" in attn_mask_type + + if qkv_format in ["bshd", "sbhd"]: + assert ( + "padding" not in attn_mask_type + ), f"No support for cp_comm_type='a2a', {attn_mask_type=} and {qkv_format=}." assert ( - not padding or qkv_format == "thd" - ), f"{attn_mask_type} mask type is not supported for BSHD and SBHD!" - assert attn_bias_type == "no_bias", f"{attn_bias_type} bias type is not supported!" - assert q.shape[-1] % 8 == 0, "Hidden size per attention head should be multiple of 8!" + attn_bias_type == "no_bias" + ), f"No support for cp_comm_type='a2a' and {attn_bias_type=}." assert ( window_size == (-1, 0) or window_size == (-1, -1) or use_fused_attention or fa_utils.v2_3_plus - ), "Sliding window attention only can work with FusedAttention or FlashAttention >= 2.3!" + ), ( + "cp_comm_type='a2a' only supports SWA through FusedAttention or FlashAttention >= 2.3." + f" Found {use_fused_attention=} and {fa_utils.v2_3_plus=}." + ) + assert q.shape[seq_dim_qkv] % 2 == 0 and k.shape[seq_dim_qkv] % 2 == 0, ( + "cp_comm_type='a2a' requires seq_len % 2 == 0 for Q, K, V. Found seq_len_q =" + f" {q.shape[seq_dim_qkv]}, seq_len_kv = {k.shape[seq_dim_qkv]}, cp_size = {cp_size}." + ) + assert q.shape[-2] % cp_size == 0 and k.shape[-2] % cp_size == 0, ( + "cp_comm_type='a2a' requires num_heads % cp_size == 0 for Q, K, V. Found num_heads_q =" + f" {q.shape[-2]}, num_heads_kv = {k.shape[-2]}, cp_size = {cp_size}." + ) flash_attn_fwd = None if not use_fused_attention: @@ -3399,26 +3926,10 @@ def forward( if fa_utils.v2_6_0_plus: fa_forward_kwargs["softcap"] = 0.0 - assert ( - q.shape[-2] % cp_size == 0 and k.shape[-2] % cp_size == 0 - ), "The number of attention heads needs to be divisible by CP size!" - - qkv_layout = qkv_format + "_" + qkv_format + "_" + qkv_format - - if qkv_format in ["bshd", "sbhd"]: - batch_dim = qkv_format.index("b") - seq_dim = qkv_format.index("s") - else: # qkv_format == "thd" - batch_dim = seq_dim = qkv_format.index("t") - - assert ( - q.shape[seq_dim] % 2 == 0 and k.shape[seq_dim] % 2 == 0 - ), "Sequence length per GPU needs to be divisible by 2!" - assert isinstance(k, q.__class__) and isinstance( v, q.__class__ - ), "q, k, v must be of the same class, e.g. torch.Tensor or Float8Tensor." - is_input_fp8 = isinstance(q, Float8Tensor) + ), "q, k, v must be of the same class, e.g. torch.Tensor or QuantizedTensorStorage." + is_input_fp8 = isinstance(q, QuantizedTensorStorage) is_output_fp8 = fp8_output is_bwd_fp8 = int(os.getenv("NVTE_FP8_DPA_BWD", "1")) # recipe passed in through autocast or set by NVTE_DPA_FP8_RECIPE; @@ -3426,62 +3937,104 @@ def forward( fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None: fp8_recipe = fp8_meta["local_recipes"][0] + fwd_nominal_dtype = q.dtype fused_attn_backend = None max_logit = None QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer = ( - dpa_utils.get_attention_quantizers(fp8, quantizers) + dpa_utils.get_attention_quantizers(fp8, fp8_recipe, quantizers) ) q_fp8, k_fp8, v_fp8 = (None, None, None) + fp8_meta_kwargs = {} if fp8: - if use_fused_attention: - fused_attn_backend = FusedAttnBackend["FP8"] - if is_input_fp8: - q_fp8, k_fp8, v_fp8 = q, k, v - q, k, v = q_fp8._data, k_fp8._data, v_fp8._data - else: - q_fp8, k_fp8, v_fp8 = combine_and_quantize(qkv_layout, q, k, v, QKV_quantizer) - q, k, v = [q_fp8._data, k_fp8._data, v_fp8._data] - fp8_meta_kwargs = {} - fp8_meta_kwargs["s_quantizer"] = S_quantizer - fp8_meta_kwargs["o_quantizer"] = O_quantizer - else: - assert False, "FP8 is only supported with Fused Attention!" + assert use_fused_attention, "FP8 is only supported with FusedAttention backend!" + fused_attn_backend = FusedAttnBackend["FP8"] + if is_input_fp8: + q_fp8, k_fp8, v_fp8 = q, k, v + elif not fp8_recipe.mxfp8(): + q_fp8, k_fp8, v_fp8, qkv_layout, _ = combine_and_quantize( + qkv_layout, q, k, v, QKV_quantizer + ) + if not fp8_recipe.mxfp8(): + q, k, v = [q_fp8._data, k_fp8._data, v_fp8._data] + fp8_meta_kwargs["s_quantizer"] = S_quantizer + fp8_meta_kwargs["o_quantizer"] = O_quantizer else: if use_fused_attention: - fp8_meta_kwargs = {} fused_attn_backend = FusedAttnBackend["F16_arbitrary_seqlen"] + # q, k, v: + # FP8DS/FP8CS: torch.uint8 + # MXFP8: torch.float16 or torch.bfloat16 + # F16: torch.float16 or torch.bfloat16 + # a2a: gather s and split h + # [b, s//cp, h, d] -> [b, s, h//cp, d] + # [s//cp, b, h, d] -> [s, b, h//cp, d] + # [t//cp, h, d] -> [t, h//cp, d] chunk_ids_for_a2a = get_seq_chunk_ids_for_reordering_before_attn(cp_size, q.device) q, k, v = flash_attn_a2a_communicate( [q, k, v], chunk_ids_for_a2a, - seq_dim, + seq_dim_qkv, cp_size, cp_group, cp_stream, before_attn=True, qkv_format=qkv_format, - cu_seqlens_padded=cu_seqlens_q_padded, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, + a2a_input_names=["q", "k", "v"], ) + + # softmax_offset: split h + # [1, h, 1, 1] -> [1, h//cp, 1, 1] if softmax_type != "vanilla": softmax_offset = flash_attn_a2a_communicate_softmax_offset( softmax_offset, 1, cp_size, cp_group, cp_stream, True ) - out_fp8 = None - out_f16 = None - batch_size = q.shape[batch_dim] + # _part: inputs to attention kernel and saved for backward + # note: they have post a2a shapes q_part, k_part, v_part = q, k, v - out_part = None + out_part, out_fp8, out_f16 = None, None, None + bwd_requires_o_f16 = is_training and ( + not is_bwd_fp8 + or ( + is_bwd_fp8 + and ( + (fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16) + or fp8_recipe.mxfp8() + ) + ) + ) + bwd_requires_o_fp8 = ( + is_training + and is_bwd_fp8 + and ( + fp8_recipe.delayed() + or (fp8_recipe.float8_current_scaling() and not _dpa_fp8_cs_o_in_f16) + ) + ) + qkv_scale_inv_format = None if use_fused_attention: if fp8: - q_part, k_part, v_part = [ - Float8Tensor.make_like(x, data=y, dtype=fwd_nominal_dtype) - for x, y in zip([q_fp8, k_fp8, v_fp8], [q_part, k_part, v_part]) - ] + if fp8_recipe.mxfp8(): + q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format = combine_and_quantize( + qkv_layout, + q_part, + k_part, + v_part, + QKV_quantizer, + used_in_backward=is_training, + ) + q_part, k_part, v_part = [q_fp8, k_fp8, v_fp8] + else: + q_part, k_part, v_part = [ + Float8Tensor.make_like(x, data=y, dtype=fwd_nominal_dtype) + for x, y in zip([q_fp8, k_fp8, v_fp8], [q_part, k_part, v_part]) + ] out_, aux_ctx_tensors, *max_logit = fused_attn_fwd( is_training, max_seqlen_q, @@ -3496,6 +4049,7 @@ def forward( attn_scale=softmax_scale, dropout=dropout_p, qkv_layout=qkv_layout, + o_format=o_format, attn_mask_type=attn_mask_type, attn_bias_type=attn_bias_type, attn_bias=attn_bias, @@ -3507,25 +4061,20 @@ def forward( softmax_offset=softmax_offset, return_max_logit=return_max_logit, cuda_graph=is_graph_capturing(), + qkv_scale_inv_format=qkv_scale_inv_format, ) - if isinstance(out_, Float8Tensor): - out_fp8 = out_ - out_ = out_._data - if is_bwd_fp8 and not ( - fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16 - ): - out_part = out_fp8 - else: - out_part = out_fp8.dequantize(dtype=fwd_nominal_dtype) - else: - out_f16 = out_ - out_part = out_ - if ( - fp8 - and is_bwd_fp8 - and not (fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16) - ): - out_part = O_quantizer(out_) + # construct out_part for backward + # out_fp8 and out_f16 store the FP8 or F16 tensor for backward saves + out_fp8 = out_ + out_f16 = out_ + if bwd_requires_o_fp8: + if not isinstance(out_, QuantizedTensorStorage): + out_fp8 = O_quantizer(out_) + out_part = out_fp8 + if bwd_requires_o_f16: + if isinstance(out_, QuantizedTensorStorage): + out_f16 = out_.dequantize(dtype=fwd_nominal_dtype) + out_part = out_f16 else: fa_forward_args_thd = get_fa_args( True, @@ -3553,60 +4102,95 @@ def forward( aux_ctx_tensors = [softmax_lse, rng_state] out_part = out_ + # a2a: split s and gather h + # [b, s, h//cp, d] -> [b*s//cp, h, d] + # [s, b, h//cp, d] -> [s//cp*b, h, d] + # [t, h//cp, d] -> [t//cp, h, d] + if isinstance(out_, Float8TensorStorage): + out_ = out_._data chunk_ids_for_a2a = get_seq_chunk_ids_for_reordering_after_attn(cp_size, out_.device) out_ = flash_attn_a2a_communicate( out_, chunk_ids_for_a2a, - seq_dim, + seq_dim_o, cp_size, cp_group, cp_stream, before_attn=False, - qkv_format=qkv_format, - cu_seqlens_padded=cu_seqlens_q_padded, + qkv_format=o_format, + cu_seqlens_q_padded=cu_seqlens_q_padded, + a2a_input_names=["out"], ) - if return_max_logit: - max_logit = flash_attn_a2a_communicate_softmax_offset( - *max_logit, 0, cp_size, cp_group, cp_stream, False - ) - - if use_fused_attention: - if qkv_format == "bshd": - # [b*s, h, d] -> [b, s, h, d] - out_ = out_.view(batch_size, -1, *out_.shape[-2:]) - elif qkv_format == "sbhd": - # [s*b, h, d] -> [s, b, h, d] - out_ = out_.view(-1, batch_size, *out_.shape[-2:]) + # [b*s//cp, h, d] -> [b, s//cp, h, d] + # [s//cp*b, h, d] -> [s//cp, b, h, d] + # [t//cp, h, d] -> [t//cp, h, d] + out_ = out_.view(orig_o_shape) - if fp8 and use_fused_attention: - if fp8_recipe.float8_current_scaling(): - out_f16 = out_ - if is_output_fp8: - out_fp8 = O_quantizer(out_) + # out_ret: output tensor for forward pass + # out_fp8 and out_f16 are reused here to store the FP8 or F16 tensor for forward returns + if fp8: if fp8_recipe.delayed(): out_fp8 = Float8Tensor.make_like(out_fp8, data=out_, dtype=fwd_nominal_dtype) - if not is_output_fp8: + if is_output_fp8: + if fp8_recipe.float8_current_scaling() or fp8_recipe.mxfp8(): + out_fp8 = O_quantizer(out_) + out_f16 = out_ + else: + if fp8_recipe.delayed(): out_f16 = out_fp8.dequantize(dtype=fwd_nominal_dtype) + else: + out_f16 = out_ else: out_f16 = out_ - out_ret = out_fp8 if is_output_fp8 else out_f16 + # all gather max logit + if return_max_logit: + max_logit = flash_attn_a2a_communicate_softmax_offset( + *max_logit, 0, cp_size, cp_group, cp_stream, False + ) + + ctx.qkv_layout = qkv_layout + ctx.o_format = o_format + ctx.qkv_scale_inv_format = qkv_scale_inv_format + ctx.dqkv_layout = original_qkv_layout + ctx.dqkv_format = qkv_format + ctx.orig_q_shape = orig_q_shape + ctx.orig_k_shape = orig_k_shape + ctx.orig_v_shape = orig_v_shape + ctx.orig_o_shape = orig_o_shape + + # save tensors for backward ctx.fp8 = fp8 and is_bwd_fp8 fp8_tensors = (None, None, None, None) f16_tensors = (None, None, None, None) - if ctx.fp8: - if fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16: - fp8_tensors = (q_part, k_part, v_part, None) - f16_tensors = (None, None, None, out_part) + if is_training: + if ctx.fp8: + # FP8DS or (FP8CS+not _dpa_fp8_cs_o_in_f16): q/k/v/o all in FP8 + # (FP8CS+_dpa_fp8_cs_o_in_f16) or MXFP8: q/k/v in FP8, o in F16 + if ( + fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16 + ) or fp8_recipe.mxfp8(): + fp8_tensors = (q_part, k_part, v_part, None) + f16_tensors = (None, None, None, out_part) + elif fp8_recipe.delayed() or ( + fp8_recipe.float8_current_scaling() and not _dpa_fp8_cs_o_in_f16 + ): + fp8_tensors = (q_part, k_part, v_part, out_part) + elif fp8: + # FP8DS/CS: convert post-a2a FP8 q/k/v to F16; out_part already in F16 + # MXFP8: save post-a2a pre-quantization F16 q/k/v; out_part already in F16 + if fp8_recipe.mxfp8(): + f16_tensors = (q, k, v, out_part) + ctx.qkv_layout = original_qkv_layout + else: + q_part, k_part, v_part = combine_and_dequantize( + qkv_layout, q_part, k_part, v_part + ) + f16_tensors = (q_part, k_part, v_part, out_part) else: - fp8_tensors = (q_part, k_part, v_part, out_part) - elif fp8: - q_part, k_part, v_part = combine_and_dequantize(qkv_layout, q_part, k_part, v_part) - f16_tensors = (q_part, k_part, v_part, out_part) - else: - f16_tensors = (q_part, k_part, v_part, out_part) - + # all tensors are in F16 + f16_tensors = (q_part, k_part, v_part, out_part) tensors_to_save, tensor_objects = prepare_for_saving( *fp8_tensors, *f16_tensors, @@ -3618,16 +4202,13 @@ def forward( ) ctx.save_for_backward(*tensors_to_save) ctx.tensor_objects = tensor_objects - ctx.out_shape = out_ret.shape - ctx.batch_size = batch_size ctx.cp_group = cp_group ctx.cp_stream = cp_stream ctx.dropout_p = dropout_p ctx.max_seqlen_q = max_seqlen_q ctx.max_seqlen_kv = max_seqlen_kv ctx.softmax_scale = softmax_scale - ctx.qkv_format = qkv_format ctx.attn_mask_type = attn_mask_type ctx.attn_bias_type = attn_bias_type ctx.deterministic = deterministic @@ -3649,11 +4230,13 @@ def forward( ctx.S_quantizer = S_quantizer if ctx.fp8: ctx.QKV_quantizer = QKV_quantizer.copy() - ctx.QKV_quantizer.scale = QKV_quantizer.scale.clone() ctx.O_quantizer = O_quantizer.copy() - ctx.O_quantizer.scale = O_quantizer.scale.clone() - ctx.S_quantizer = S_quantizer.copy() - ctx.S_quantizer.scale = S_quantizer.scale.clone() + ctx.S_quantizer = S_quantizer.copy() if S_quantizer is not None else None + if not ctx.fp8_recipe.mxfp8(): + ctx.QKV_quantizer.scale = QKV_quantizer.scale.clone() + ctx.O_quantizer.scale = O_quantizer.scale.clone() + ctx.S_quantizer.scale = S_quantizer.scale.clone() + nvtx_range_pop("transformer_engine.AttnFuncWithCPAndQKVOA2A.forward") if return_max_logit: return out_ret, max_logit @@ -3681,60 +4264,53 @@ def backward(ctx, dout, *_args): *aux_ctx_tensors, ) = restore_from_func_ctx(ctx) - qkv_format = ctx.qkv_format - qkv_layout = qkv_format + "_" + qkv_format + "_" + qkv_format - causal = "causal" in ctx.attn_mask_type - - if qkv_format in ["bshd", "sbhd"]: - seq_dim = qkv_format.index("s") - else: # qkv_format == "thd" - seq_dim = qkv_format.index("t") - + _, seq_dim_dqkv, _ = get_bsh_dims(ctx.dqkv_format) + _, seq_dim_do, _ = get_bsh_dims(ctx.o_format) bwd_nominal_dtype = ctx.fwd_nominal_dtype - dqkv_te_dtype = None fused_attn_backend = None - dout_fp8 = dout + causal = "causal" in ctx.attn_mask_type + + dout_fp8 = None + fp8_meta_kwargs = {} if ctx.fp8: - if ctx.use_fused_attention: - fused_attn_backend = FusedAttnBackend["FP8"] - if not isinstance(dout, QuantizedTensorStorage): - dout = ctx.dO_quantizer(dout) - dout_fp8 = dout - dqkv_te_dtype = dout._fp8_dtype + assert ctx.use_fused_attention, "FP8 is only supported with FusedAttention backend!" + fused_attn_backend = FusedAttnBackend["FP8"] + if isinstance(dout, QuantizedTensorStorage): + dout_fp8 = dout + elif not ctx.fp8_recipe.mxfp8(): + dout = ctx.dO_quantizer(dout) + dout_fp8 = dout + if not ctx.fp8_recipe.mxfp8(): dout = dout._data - fp8_meta_kwargs = {} - fp8_meta_kwargs["s_quantizer"] = ctx.S_quantizer - fp8_meta_kwargs["dp_quantizer"] = ctx.dP_quantizer - fp8_meta_kwargs["dqkv_quantizer"] = ctx.dQKV_quantizer - - else: - assert False, "FP8 is only supported with Fused Attention!" + fp8_meta_kwargs["s_quantizer"] = ctx.S_quantizer + fp8_meta_kwargs["dp_quantizer"] = ctx.dP_quantizer + fp8_meta_kwargs["dqkv_quantizer"] = ctx.dQKV_quantizer else: if isinstance(dout, QuantizedTensorStorage): dout = dout.dequantize(dtype=bwd_nominal_dtype) if ctx.use_fused_attention: - fp8_meta_kwargs = {} - dqkv_te_dtype = TE_DType[dout.dtype] fused_attn_backend = FusedAttnBackend["F16_arbitrary_seqlen"] - - if not ctx.use_fused_attention: - if qkv_format in ["bshd", "sbhd"]: - out = out.view(ctx.batch_size, -1, *out.shape[-2:]) - dout = dout.view(ctx.batch_size, -1, *dout.shape[-2:]) - else: - dout = dout.view(*ctx.out_shape) - + dout = dout.view(*ctx.orig_o_shape) + + # dout: + # FP8DS/CS: torch.uint8 + # MXFP8/F16: torch.float16 or torch.bfloat16 + # a2a: gather s and split h + # [b, s//cp, h, d] -> [b, s, h//cp, d] + # [s//cp, b, h, d] -> [s, b, h//cp, d] + # [t//cp, h, d] -> [t, h//cp, d] chunk_ids_for_a2a = get_seq_chunk_ids_for_reordering_before_attn(cp_size, dout.device) dout = flash_attn_a2a_communicate( dout, chunk_ids_for_a2a, - seq_dim, + seq_dim_do, cp_size, ctx.cp_group, ctx.cp_stream, before_attn=True, - qkv_format=qkv_format, - cu_seqlens_padded=cu_seqlens_q_padded, + qkv_format=ctx.o_format, + cu_seqlens_q_padded=cu_seqlens_q_padded, + a2a_input_names=["dout"], ) flash_attn_bwd = None @@ -3752,7 +4328,7 @@ def backward(ctx, dout, *_args): fa_backward_kwargs["window_size_right"] = ctx.window_size[1] fa_backward_kwargs["deterministic"] = ctx.deterministic else: - if qkv_format == "thd": + if ctx.o_format == "thd": from transformer_engine.pytorch.attention.dot_product_attention.backends import ( _flash_attn_varlen_bwd, ) @@ -3779,12 +4355,23 @@ def backward(ctx, dout, *_args): dq_fp8, dk_fp8, dv_fp8 = None, None, None if ctx.use_fused_attention: + do_format = ctx.o_format + do_scale_inv_format = None q_part, k_part, v_part, out_part, dout_part = q, k, v, out, dout if ctx.fp8: q_part, k_part, v_part, out_part = q_fp8, k_fp8, v_fp8, out_fp8 - if ctx.fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16: + if ( + ctx.fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16 + ) or ctx.fp8_recipe.mxfp8(): out_part = out - dout_part = Float8Tensor.make_like(dout_fp8, data=dout, dtype=bwd_nominal_dtype) + if not ctx.fp8_recipe.mxfp8(): + dout_part = Float8Tensor.make_like(dout_fp8, data=dout, dtype=bwd_nominal_dtype) + else: + aux_ctx_tensors.append(dout) + (dout_part,), do_scale_inv_format = mxfp8_quantize_fast_path( + [(dout, ctx.dO_quantizer)], + do_format, + ) dq, dk, dv, *rest = fused_attn_bwd( ctx.max_seqlen_q, ctx.max_seqlen_kv, @@ -3796,23 +4383,27 @@ def backward(ctx, dout, *_args): out_part, dout_part, bwd_nominal_dtype, - dqkv_te_dtype, aux_ctx_tensors, fused_attn_backend, cu_seqlens_q_padded=cu_seqlens_q_padded, cu_seqlens_kv_padded=cu_seqlens_kv_padded, attn_scale=ctx.softmax_scale, dropout=ctx.dropout_p, - qkv_layout=qkv_layout, + qkv_layout=ctx.qkv_layout, + o_format=ctx.o_format, + do_format=do_format, + dqkv_layout=ctx.dqkv_layout, attn_mask_type=ctx.attn_mask_type, attn_bias_type=ctx.attn_bias_type, window_size=ctx.window_size, deterministic=ctx.deterministic, cuda_graph=is_graph_capturing(), + qkv_scale_inv_format=ctx.qkv_scale_inv_format, + do_scale_inv_format=do_scale_inv_format, **fp8_meta_kwargs, softmax_type=ctx.softmax_type, ) - if isinstance(dq, Float8Tensor): + if all(isinstance(x, Float8TensorStorage) for x in [dq, dk, dv]): dq_fp8, dk_fp8, dv_fp8 = dq, dk, dv dq, dk, dv = [x._data for x in [dq, dk, dv]] else: @@ -3821,7 +4412,7 @@ def backward(ctx, dout, *_args): fa_backward_args_thd = get_fa_args( False, ctx.use_flash_attn_3, - qkv_format, + ctx.dqkv_format, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, max_seqlen_q=ctx.max_seqlen_q, @@ -3847,24 +4438,33 @@ def backward(ctx, dout, *_args): **fa_backward_kwargs, ) + # dq, dk, dv: + # FP8DS: torch.uint8 + # FP8CS/MXFP8/F16: torch.float16 or torch.bfloat16 + # a2a: gather s and split h + # [b, s//cp, h, d] -> [b, s, h//cp, d] + # [s//cp, b, h, d] -> [s, b, h//cp, d] + # [t//cp, h, d] -> [t, h//cp, d] chunk_ids_for_a2a = get_seq_chunk_ids_for_reordering_after_attn(cp_size, dq.device) dq, dk, dv = flash_attn_a2a_communicate( [dq, dk, dv], chunk_ids_for_a2a, - seq_dim, + seq_dim_dqkv, cp_size, ctx.cp_group, ctx.cp_stream, before_attn=False, - qkv_format=qkv_format, - cu_seqlens_padded=cu_seqlens_q_padded, + qkv_format=ctx.dqkv_format, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, + a2a_input_names=["dq", "dk", "dv"], ) + dq, dk, dv = [ + x.view(y) + for x, y in zip([dq, dk, dv], [ctx.orig_q_shape, ctx.orig_k_shape, ctx.orig_v_shape]) + ] - if qkv_format == "bshd": - dq, dk, dv = [x.view(ctx.batch_size, -1, *x.shape[-2:]) for x in [dq, dk, dv]] - elif qkv_format == "sbhd": - dq, dk, dv = [x.view(-1, ctx.batch_size, *x.shape[-2:]) for x in [dq, dk, dv]] - + # d_bias, d_softmax_offset d_bias = None d_softmax_offset = None if ctx.use_fused_attention: @@ -3876,9 +4476,14 @@ def backward(ctx, dout, *_args): d_softmax_offset, 1, cp_size, ctx.cp_group, ctx.cp_stream, False ) + # convert dq, dk, dv to appropriate types if ctx.fp8: - if ctx.fp8_recipe.float8_current_scaling() and ctx.is_input_fp8: - dq, dk, dv = combine_and_quantize(qkv_layout, dq, dk, dv, ctx.dQKV_quantizer) + if ( + ctx.fp8_recipe.float8_current_scaling() or ctx.fp8_recipe.mxfp8() + ) and ctx.is_input_fp8: + dq, dk, dv, _, _ = combine_and_quantize( + ctx.dqkv_layout, dq, dk, dv, ctx.dQKV_quantizer + ) if ctx.fp8_recipe.delayed(): dq, dk, dv = [ Float8Tensor.make_like(x, data=y, dtype=bwd_nominal_dtype) @@ -3886,7 +4491,7 @@ def backward(ctx, dout, *_args): ] if not ctx.is_input_fp8: dq, dk, dv = combine_and_dequantize( - qkv_layout, + ctx.dqkv_layout, dq, dk, dv, @@ -3894,7 +4499,6 @@ def backward(ctx, dout, *_args): ) nvtx_range_pop("transformer_engine.AttnFuncWithCPAndQKVOA2A.backward") - return ( None, dq, @@ -4069,17 +4673,6 @@ def attn_forward_func_with_cp( "all_gather", ], f"Context parallelism does not support sliding window attention with {cp_comm_type=}!" - enable_mla = k.shape[-1] != v.shape[-1] - assert not enable_mla or cp_comm_type in [ - "p2p", - "a2a+p2p", - ], f"Context parallelism does not support MLA with {cp_comm_type=}!" - - if fp8 and fp8_meta is not None: - if fp8_meta["recipe"].fp8_dpa: - assert ( - softmax_type == "vanilla" - ), f"Context parallelism does not support {softmax_type=} with FP8 attention!" assert ( softmax_type == "vanilla" or use_fused_attention ), f"Context parallelism only supports {softmax_type=} with FusedAttention backend!" @@ -4131,7 +4724,16 @@ def attn_forward_func_with_cp( elif cp_comm_type == "all_gather": args.pop(5) args.pop(8) - args += [window_size, cp_group, cp_stream, use_flash_attn_3] + args += [ + window_size, + cp_group, + cp_stream, + use_flash_attn_3, + fp8, + fp8_meta, + quantizers, + fp8_output, + ] out = AttnFuncWithCPAndKVAllGather.apply(*args) elif cp_comm_type == "a2a": args += [ diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 588c708e1..17e9a337a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -19,6 +19,7 @@ Recipe, DelayedScaling, Float8CurrentScaling, + MXFP8BlockScaling, ) from transformer_engine.pytorch.utils import get_cudnn_version from transformer_engine.pytorch.quantization import ( @@ -30,7 +31,7 @@ Float8CurrentScalingRecipeState, Float8BlockScalingRecipeState, ) -from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor +from transformer_engine.pytorch.tensor.storage.float8_tensor_storage import Float8TensorStorage from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.export import is_in_onnx_export_mode from transformer_engine.pytorch.constants import ( @@ -98,19 +99,26 @@ +-------------------+-----------+-----------------------------------------------------------------------------------+ | Linear | Attention | Configuration | +===================+===========+===================================================================================+ -| FP8DS/FP8CS/NVFP4 | FP16/BF16 | Pass FP8DS, FP8CS or NVFP4 to autocast(); | -| | | export NVTE_DPA_FP8_RECIPE="F16" | +| FP8DS/FP8CS/NVFP4 | FP16/BF16 | Pass FP8DS, FP8CS, NVFP4 or MXFP8 to autocast(); | +| /MXFP8 | | export NVTE_DPA_FP8_RECIPE="F16" | +-------------------+-----------+-----------------------------------------------------------------------------------+ -| FP8DS | FP8DS | Pass FP8DS to autocast(); | +| FP8DS | FP8DS | Pass FP8DS to autocast(); | +-------------------+-----------+-----------------------------------------------------------------------------------+ -| FP8CS | FP8DS | Pass FP8CS to autocast(); | +| FP8CS | FP8DS | Pass FP8CS to autocast(); | | | | Attention FP8DS reuses the fp8_format, fp8_dpa, fp8_mha values from linear FP8CS; | | | | export NVTE_DPA_FP8_RECIPE="DelayedScaling" # switch to DS | | | | export NVTE_DPA_FP8DS_AMAX_ALGO="most_recent" # or "max" | | | | export NVTE_DPA_FP8DS_AMAX_HISTLEN=1 # or any other integer | | | | export NVTE_DPA_FP8DS_REDUCE_AMAX=1 # or 0 | +-------------------+-----------+-----------------------------------------------------------------------------------+ -| NVFP4 | FP8DS | Pass NVFP4 to autocast(); | +| MXFP8 | FP8DS | Pass MXFP8 to autocast(); | +| | | Attention FP8DS reuses the fp8_format, fp8_dpa, fp8_mha values from linear MXFP8; | +| | | export NVTE_DPA_FP8_RECIPE="DelayedScaling" # switch to DS | +| | | export NVTE_DPA_FP8DS_AMAX_ALGO="most_recent" # or "max" | +| | | export NVTE_DPA_FP8DS_AMAX_HISTLEN=1 # or any other integer | +| | | export NVTE_DPA_FP8DS_REDUCE_AMAX=1 # or 0 | ++-------------------+-----------+-----------------------------------------------------------------------------------+ +| NVFP4 | FP8DS | Pass NVFP4 to autocast(); | | | | Attention FP8DS reuses the fp8_dpa, fp8_mha values from linear NVFP4; | | | | export NVTE_DPA_FP8_RECIPE="DelayedScaling" # switch to DS | | | | export NVTE_DPA_FP8_FORMAT="HYBRID" # or "E4M3", "E5M2" | @@ -118,19 +126,27 @@ | | | export NVTE_DPA_FP8DS_AMAX_HISTLEN=1 # or any other integer | | | | export NVTE_DPA_FP8DS_REDUCE_AMAX=1 # or 0 | +-------------------+-----------+-----------------------------------------------------------------------------------+ -| FP8DS | FP8CS | Pass FP8DS to autocast(); | +| FP8DS | FP8CS | Pass FP8DS to autocast(); | | | | Attention uses FP8DS for S, dP tensors, and creates a new FP8CS recipe for QKV, O,| | | | dO, dQKV tensors based on fp8_format, fp8_dpa, fp8_mha from linear FP8DS; | | | | export NVTE_DPA_FP8_RECIPE="Float8CurrentScaling" # switch to CS | +-------------------+-----------+-----------------------------------------------------------------------------------+ -| FP8CS | FP8CS | Pass FP8CS to autocast(); | +| FP8CS | FP8CS | Pass FP8CS to autocast(); | | | | Attention uses FP8CS for QKV, O, dO, dQKV tensors, and creates a new FP8DS recipe | | | | for S, dP tensors based on fp8_format, fp8_dpa, fp8_mha from linear FP8CS and: | | | | export NVTE_DPA_FP8DS_AMAX_ALGO="most_recent" # or "max" | | | | export NVTE_DPA_FP8DS_AMAX_HISTLEN=1 # or any other integer | | | | export NVTE_DPA_FP8DS_REDUCE_AMAX=1 # or 0 | +-------------------+-----------+-----------------------------------------------------------------------------------+ -| NVFP4 | FP8CS | Pass NVFP4 to autocast(); | +| MXFP8 | FP8CS | Pass MXFP8 to autocast(); | +| | | Attention creates a new FP8CS recipe based on fp8_format, fp8_dpa, fp8_mha from | +| | | linear MXFP8, and: | +| | | export NVTE_DPA_FP8_RECIPE="Float8CurrentScaling" # switch to CS | +| | | export NVTE_DPA_FP8DS_AMAX_ALGO="most_recent" # or "max" | +| | | export NVTE_DPA_FP8DS_AMAX_HISTLEN=1 # or any other integer | +| | | export NVTE_DPA_FP8DS_REDUCE_AMAX=1 # or 0 | ++-------------------+-----------+-----------------------------------------------------------------------------------+ +| NVFP4 | FP8CS | Pass NVFP4 to autocast(); | | | | Attention creates a new FP8CS recipe for QKV, O, dO, dQKV, and a new FP8DS recipe | | | | for S, dP, based on the fp8_dpa, fp8_mha values from linear NVFP4 and: | | | | export NVTE_DPA_FP8_RECIPE="Float8CurrentScaling" # switch to CS | @@ -139,6 +155,18 @@ | | | export NVTE_DPA_FP8DS_AMAX_HISTLEN=1 # or any other integer | | | | export NVTE_DPA_FP8DS_REDUCE_AMAX=1 # or 0 | +-------------------+-----------+-----------------------------------------------------------------------------------+ +| FP8DS/FP8CS | MXFP8 | Pass FP8DS/FP8CS to autocast(); | +| | | Attention creates a new MXFP8 recipe based on fp8_format, fp8_dpa, fp8_mha from | +| | | linear FP8DS/FP8CS | +| | | export NVTE_DPA_FP8_RECIPE="MXFP8BlockScaling" # switch to MXFP8BS | ++-------------------+-----------+-----------------------------------------------------------------------------------+ +| MXFP8 | MXFP8 | Pass MXFP8 to autocast(); | ++-------------------+-----------+-----------------------------------------------------------------------------------+ +| NVFP4 | MXFP8 | Pass NVFP4 to autocast(); | +| | | Attention MXFP8 reuses the fp8_dpa, fp8_mha values from linear NVFP4; | +| | | export NVTE_DPA_FP8_RECIPE="MXFP8BlockScaling" # switch to MXFP8BS | +| | | export NVTE_DPA_FP8_FORMAT="HYBRID" # or "E4M3", "E5M2" | ++-------------------+-----------+-----------------------------------------------------------------------------------+ """ _dpa_fp8_recipe = os.getenv("NVTE_DPA_FP8_RECIPE", "") formats = {"HYBRID": Format.HYBRID, "E4M3": Format.E4M3, "E5M2": Format.E5M2} @@ -600,7 +628,9 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: # ignore the recipe from autocast, set fp8_dpa = False, fp8_mha = False fp8_recipe.fp8_dpa = False fp8_recipe.fp8_mha = False - elif fp8_recipe.float8_current_scaling() and _dpa_fp8_recipe == "DelayedScaling": + elif ( + fp8_recipe.float8_current_scaling() or fp8_recipe.mxfp8() + ) and _dpa_fp8_recipe == "DelayedScaling": # reuse fp8_format, fp8_dpa, fp8_mha from fp8_recipe, and construct a DS recipe fake_recipe = DelayedScaling( fp8_format=fp8_recipe.fp8_format, @@ -653,6 +683,25 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: ) fp8_recipe_dpa = fake_recipe fp8_recipes = [fp8_recipe, fp8_recipe_dpa] + elif fp8_recipe.mxfp8() and _dpa_fp8_recipe == "Float8CurrentScaling": + # reuse fp8_format, fp8_dpa, fp8_mha from fp8_recipe, and construct a CS+DS recipe + fake_recipes = [ + Float8CurrentScaling( + fp8_format=fp8_recipe.fp8_format, + fp8_dpa=fp8_recipe.fp8_dpa, + fp8_mha=fp8_recipe.fp8_mha, + ), + DelayedScaling( + fp8_format=fp8_recipe.fp8_format, + amax_history_len=_dpa_fp8ds_amax_histlen, + amax_compute_algo=_dpa_fp8ds_amax_algo, + fp8_dpa=fp8_recipe.fp8_dpa, + fp8_mha=fp8_recipe.fp8_mha, + reduce_amax=_dpa_fp8ds_reduce_amax, + ), + ] + fp8_recipe_dpa = fake_recipes[1] + fp8_recipes = fake_recipes elif fp8_recipe.nvfp4() and _dpa_fp8_recipe == "Float8CurrentScaling": # reuse fp8_dpa, fp8_mha from fp8_recipe but not fp8_format # construct a CS recipe for QKV, O, dO, dQKV and a DS recipe for S, dP @@ -673,11 +722,26 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: ] fp8_recipe_dpa = fake_recipes[1] fp8_recipes = fake_recipes - # DPA only support DS and CS; other recipes should have fp8_dpa=False, fp8_mha=False - if not fp8_recipe_dpa.float8_per_tensor_scaling(): - assert not ( - fp8_recipe_dpa.fp8_dpa or fp8_recipe_dpa.fp8_mha - ), f"DotProductAttention does not support {fp8_recipe_dpa.__class__.__name__} recipe" + elif ( + fp8_recipe.delayed() or fp8_recipe.float8_current_scaling() + ) and _dpa_fp8_recipe == "MXFP8BlockScaling": + # reuse fp8_format, fp8_dpa, fp8_mha from fp8_recipe, and construct a MXFP8 recipe + fake_recipe = MXFP8BlockScaling( + fp8_format=fp8_recipe.fp8_format, + fp8_dpa=fp8_recipe.fp8_dpa, + fp8_mha=fp8_recipe.fp8_mha, + ) + fp8_recipe_dpa = fake_recipe + fp8_recipes = fp8_recipe_dpa + elif fp8_recipe.nvfp4() and _dpa_fp8_recipe == "MXFP8BlockScaling": + # reuse fp8_dpa, fp8_mha from fp8_recipe but not fp8_format; construct a MXFP8 recipe + fake_recipe = MXFP8BlockScaling( + fp8_format=_dpa_fp8_format, + fp8_dpa=fp8_recipe.fp8_dpa, + fp8_mha=fp8_recipe.fp8_mha, + ) + fp8_recipe_dpa = fake_recipe + fp8_recipes = fp8_recipe_dpa # reduce over TP+CP groups; expect fp8_group to be set up so # assume attention uses the same fp8_group as GEMMs @@ -1203,7 +1267,9 @@ def forward( cu_seqlens_kv_padded = None # get qkv's memory layout - if all(isinstance(x, Float8Tensor) for x in [query_layer, key_layer, value_layer]): + if all( + isinstance(x, Float8TensorStorage) for x in [query_layer, key_layer, value_layer] + ): ( qkv_layout, query_layer._data, @@ -1365,6 +1431,7 @@ def forward( attention_dropout=self.attention_dropout, context_parallel=context_parallel, cp_comm_type=self.cp_comm_type, + cp_size=cp_size, deterministic=self.deterministic, is_training=self.training, fp8=self.fp8, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 20228ddb8..c416e49da 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -35,13 +35,18 @@ META_DP, ) from transformer_engine.pytorch.attention.inference import InferenceParams +from transformer_engine.pytorch.quantized_tensor import QuantizedTensorStorage from transformer_engine.pytorch.tensor.float8_tensor import ( Float8Tensor, Float8Quantizer, Float8CurrentScalingQuantizer, ) +from transformer_engine.pytorch.tensor.float8_tensor import Float8TensorStorage +from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer, MXFP8Tensor +from transformer_engine.pytorch.tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage + from transformer_engine.pytorch.quantization import get_fp8_te_dtype -from transformer_engine.pytorch.constants import TE_DType +from transformer_engine.pytorch.constants import TE_DType, MXFP8_BLOCK_SCALING_SIZE from transformer_engine.pytorch.utils import ( @@ -231,6 +236,8 @@ class AttentionParams: Whether context parallelism is used or not. cp_comm_type : str, default = "p2p" The communication type of context parallelism. + cp_size : int, default = 1 + The group size of context parallelism. deterministic : bool, default = False Whether to run `DotProductAttention` with determinism or not. is_training : bool, default = True @@ -272,6 +279,7 @@ class AttentionParams: attention_dropout: float = 0.0 context_parallel: bool = False cp_comm_type: str = "p2p" + cp_size: int = 1 deterministic: bool = False is_training: bool = True fp8: bool = False @@ -349,6 +357,7 @@ def get_attention_backend( attention_dropout = attention_params.attention_dropout context_parallel = attention_params.context_parallel cp_comm_type = attention_params.cp_comm_type + cp_size = attention_params.cp_size # pylint: disable=unused-variable deterministic = attention_params.deterministic is_training = attention_params.is_training fp8 = attention_params.fp8 @@ -368,9 +377,9 @@ def get_attention_backend( cudnn_version = get_cudnn_version() run_config = { "transformer_engine_version": te.__version__, - "compute_capability": ( - "sm" + str(10 * device_compute_capability[0] + device_compute_capability[1]) - ), + "compute_capability": "sm" + + str(10 * device_compute_capability[0] + device_compute_capability[1]), + "cuda_version": torch.version.cuda, "flash_attn_version": ( str(FlashAttentionUtils.version) if FlashAttentionUtils.is_installed @@ -488,21 +497,30 @@ def get_attention_backend( if qkv_dtype not in [torch.bfloat16, torch.float16, torch.float8_e4m3fn] or qkv_type not in [ torch.Tensor, Float8Tensor, + Float8TensorStorage, ]: if use_flash_attention_3 and FlashAttentionUtils.v3_is_installed: logger.debug( - "Disabling FlashAttention 3 for unsupported qkv_dtype = %s, qkv_type = %s. " - "Supported: qkv_dtype = {torch.bfloat16, torch.float16, torch.float8_e4m3fn}, " - "qkv_type = {torch.Tensor, Float8Tensor}. ", + "Disabling FlashAttention 3 for unsupported qkv_dtype = %s, qkv_type = %s." + " Supported: qkv_dtype = {torch.bfloat16, torch.float16, torch.float8_e4m3fn}," + " qkv_type = {torch.Tensor, Float8Tensor, Float8TensorStorage}. ", qkv_dtype, qkv_type, ) use_flash_attention_3 = False + if qkv_dtype not in [torch.bfloat16, torch.float16, torch.float8_e4m3fn] or qkv_type not in ( + torch.Tensor, + Float8Tensor, + Float8TensorStorage, + MXFP8Tensor, + MXFP8TensorStorage, + ): if use_fused_attention: logger.debug( - "Disabling FusedAttention for unsupported qkv_dtype = %s, qkv_type = %s. " - "Supported: qkv_dtype = {torch.bfloat16, torch.float16, torch.float8_e4m3fn}, " - "qkv_type = {torch.Tensor, Float8Tensor}. ", + "Disabling FusedAttention for unsupported qkv_dtype = %s, qkv_type = %s. Supported:" + " qkv_dtype = {torch.bfloat16, torch.float16, torch.float8_e4m3fn}, qkv_type =" + " {torch.Tensor, Float8Tensor, Float8TensorStorage, MXFP8Tensor," + " MXFP8TensorStorage}. ", qkv_dtype, qkv_type, ) @@ -510,6 +528,9 @@ def get_attention_backend( # Filter: Execution type if fp8 and fp8_meta["recipe"].fp8_dpa: + fp8_recipe = fp8_meta["recipe"] + if fp8_meta.get("local_recipes", None) is not None: + fp8_recipe = fp8_meta["local_recipes"][0] if use_flash_attention_2 and FlashAttentionUtils.is_installed: logger.debug("Disabling FlashAttention 2 for FP8 attention") use_flash_attention_2 = False @@ -520,6 +541,12 @@ def get_attention_backend( if FlashAttentionUtils.v3_is_installed: logger.debug("Disabling FlashAttention 3 for FP8 training") use_flash_attention_3 = False + if use_flash_attention_3 and not ( + fp8_recipe.delayed() or fp8_recipe.float8_current_scaling() + ): + if FlashAttentionUtils.v3_is_installed: + logger.debug("Disabling FlashAttention 3 for %s", fp8_recipe.__class__.__name__) + use_flash_attention_3 = False if use_unfused_attention: allow_emulation = ( os.getenv("NVTE_UnfusedDPA_Emulate_FP8", "0") == "1" or is_in_onnx_export_mode() @@ -527,15 +554,21 @@ def get_attention_backend( if not allow_emulation: logger.debug("Disabling UnfusedDotProductAttention for FP8 attention") use_unfused_attention = False - fp8_recipe = fp8_meta["recipe"] - if fp8_meta.get("local_recipes", None) is not None: - fp8_recipe = fp8_meta["local_recipes"][0] + if use_fused_attention and fp8_recipe.delayed(): + if ( + device_compute_capability >= (10, 0) + and deterministic + and cudnn_version < (9, 18, 0) + ): + logger.debug( + "Disabling FusedAttention for FP8 delayed scaling on arch >= sm100 with" + " determinism for cuDNN < 9.18.0" + ) + use_fused_attention = False if use_fused_attention and fp8_recipe.float8_current_scaling(): if device_compute_capability < (10, 0): logger.debug("Disabling FusedAttention for FP8 current scaling on arch < sm100") use_fused_attention = False - # TODO(cyanguwa): Modify the min cuDNN version supporting FP8 current scaling - # determinism for Blackwell else: if cudnn_version < (9, 14, 0): logger.debug( @@ -545,10 +578,27 @@ def get_attention_backend( else: if deterministic and cudnn_version < (9, 18, 0): logger.debug( - "Disabling FusedAttention for FP8 current scaling requiring determinism" - " with cuDNN < 9.18.0" + "Disabling FusedAttention for FP8 current scaling with determinism" + " for cuDNN < 9.18.0" ) use_fused_attention = False + if use_fused_attention and fp8_recipe.mxfp8(): + if device_compute_capability < (10, 0): + logger.debug("Disabling FusedAttention for MXFP8 on arch < sm100") + use_fused_attention = False + elif fp8_recipe.fp8_mha: + logger.debug("Disabling FusedAttention for MXFP8 with fp8_mha=True") + use_fused_attention = False + else: + if cudnn_version < (9, 21, 0): + logger.debug("Disabling FusedAttention for MXFP8 with cuDNN < 9.21.0") + use_fused_attention = False + elif qkv_format == "thd": + logger.debug("Disabling FusedAttention for MXFP8 with qkv_format = thd") + use_fused_attention = False + if use_fused_attention and (fp8_recipe.float8_block_scaling() or fp8_recipe.nvfp4()): + logger.debug("Disabling FusedAttention for %s", fp8_recipe.__class__.__name__) + use_fused_attention = False if device_compute_capability == (12, 0): if use_flash_attention: @@ -837,29 +887,36 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt logger.debug("Disabling FlashAttention for softmax_type = %s", softmax_type) use_flash_attention = False if fp8 and fp8_meta["recipe"].fp8_dpa: - logger.debug("Disabling FusedAttention for softmax_type = %s in FP8", softmax_type) - use_fused_attention = False - logger.debug( - "Disabling UnfusedDotProductAttention for softmax_type = %s in FP8", softmax_type - ) - use_unfused_attention = False - if qkv_format == "thd": - if cudnn_version < (9, 18, 0): + if use_fused_attention and ( + device_compute_capability < (10, 0) or cudnn_version < (9, 21, 0) + ): logger.debug( - "Disabling FusedAttention for softmax_type = %s, qkv_format = thd and cuDNN" - " version < 9.18", + "Disabling FusedAttention for softmax_type = %s in FP8 on sm < 100 with cuDNN" + " version < 9.21", softmax_type, ) use_fused_attention = False - if context_parallel: - if cp_comm_type != "a2a": + if use_unfused_attention: logger.debug( - "Disabling FusedAttention for context parallelism with softmax_type = %s and" - " cp_comm_type = %s", + "Disabling UnfusedDotProductAttention for softmax_type = %s in FP8", softmax_type, - cp_comm_type, ) - use_fused_attention = False + use_unfused_attention = False + if qkv_format == "thd" and cudnn_version < (9, 18, 0): + logger.debug( + "Disabling FusedAttention for softmax_type = %s, qkv_format = thd and cuDNN" + " version < 9.18", + softmax_type, + ) + use_fused_attention = False + if context_parallel and cp_comm_type != "a2a": + logger.debug( + "Disabling FusedAttention for context parallelism with softmax_type = %s and" + " cp_comm_type = %s", + softmax_type, + cp_comm_type, + ) + use_fused_attention = False # Filter: Context parallelism # qkv_format | attn_mask_type | attn_bias_type | supported backends @@ -946,10 +1003,50 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt " bias for THD format" ) use_fused_attention = False - elif fp8 and fp8_meta["recipe"].fp8_dpa and head_dim_qk != head_dim_v: + elif fp8 and fp8_meta["recipe"].fp8_dpa and qkv_format == "thd": logger.debug( "Disabling FusedAttention as it does not support context parallelism with FP8" - " MLA attention" + " attention and THD format" + ) + use_fused_attention = False + elif fp8 and fp8_meta["recipe"].fp8_dpa and core_attention_bias_type != "no_bias": + logger.debug( + "Disabling FusedAttention as it does not support context parallelism with FP8" + " attention and bias" + ) + use_fused_attention = False + elif core_attention_bias_type != "no_bias" and cp_comm_type != "p2p": + logger.debug( + "Disabling FusedAttention as it does not support context parallelism with bias" + " and cp_comm_type = %s", + cp_comm_type, + ) + use_fused_attention = False + elif qkv_format == "thd" and cp_comm_type in ["all_gather", "a2a+p2p"]: + logger.debug( + "Disabling FusedAttention as it does not support context parallelism with THD" + " format and cp_comm_type = %s", + cp_comm_type, + ) + use_fused_attention = False + elif ( + window_size is not None + and (window_size[0] != -1 or window_size[1] not in [-1, 0]) + and cp_comm_type in ["p2p", "a2a+p2p"] + ): + logger.debug( + "Disabling FusedAttention as it does not support context parallelism with sliding" + " window attention and cp_comm_type = %s", + cp_comm_type, + ) + use_fused_attention = False + elif cp_comm_type in ["a2a", "a2a+p2p"] and (num_heads % 2 != 0 or num_gqa_groups % 2 != 0): + logger.debug( + "Disabling FusedAttention as cp_comm_type = %s requires num_heads and" + " num_gqa_groups divisible by 2 (got num_heads = %s, num_gqa_groups = %s)", + cp_comm_type, + num_heads, + num_gqa_groups, ) use_fused_attention = False @@ -1004,9 +1101,14 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if window_size is None: window_size = check_set_window_size(attn_mask_type, window_size) if use_fused_attention and (window_size[0] != -1 or window_size[1] not in [-1, 0]): - if fp8 and (fp8_meta["recipe"].fp8_dpa or fp8_meta["recipe"].fp8_mha): + if ( + fp8 + and (fp8_meta["recipe"].fp8_dpa or fp8_meta["recipe"].fp8_mha) + and (device_compute_capability < (10, 0) or cudnn_version < (9, 21, 0)) + ): logger.debug( "Disabling FusedAttention as it does not support sliding window attention for FP8" + " on sm < 100 with cuDNN version < 9.21" ) use_fused_attention = False elif attention_dropout != 0.0: @@ -1150,8 +1252,8 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if ( use_fused_attention and window_size is not None - and window_size[0] != -1 - and fused_attention_backend != FusedAttnBackend["F16_arbitrary_seqlen"] + and (window_size[0] != -1 or window_size[1] not in [-1, 0]) + and fused_attention_backend == FusedAttnBackend["F16_max512_seqlen"] ): logger.debug( "Disabling FusedAttention as only sub-backend %s does not support " @@ -2256,28 +2358,45 @@ def check_set_window_size( return window_size -def get_attention_quantizers(fp8, quantizers): +def get_attention_quantizers(fp8, fp8_recipe, quantizers): """Get the list of quantizers used in attention from the quantizers list.""" if not fp8: return [None] * 6 + QKV_quantizer = quantizers["scaling_fwd"][META_QKV] - QKV_quantizer.internal = True + QKV_quantizer.internal = False QKV_quantizer.set_usage(rowwise=True, columnwise=False) - O_quantizer = quantizers["scaling_fwd"][META_O] - O_quantizer.set_usage(rowwise=True, columnwise=False) + S_quantizer = quantizers["scaling_fwd"][META_S] S_quantizer.internal = True S_quantizer.set_usage(rowwise=True, columnwise=False) - dQKV_quantizer = quantizers["scaling_bwd"][META_DQKV] - dQKV_quantizer.interal = True - dQKV_quantizer.set_usage(rowwise=True, columnwise=False) + O_quantizer = quantizers["scaling_fwd"][META_O] + O_quantizer.internal = False + O_quantizer.set_usage(rowwise=True, columnwise=False) + dO_quantizer = quantizers["scaling_bwd"][META_DO] + dO_quantizer.internal = False dO_quantizer.set_usage(rowwise=True, columnwise=False) - dO_quantizer.internal = True + dP_quantizer = quantizers["scaling_bwd"][META_DP] + dP_quantizer.internal = True dP_quantizer.set_usage(rowwise=True, columnwise=False) - dP_quantizer.interal = True + + dQKV_quantizer = quantizers["scaling_bwd"][META_DQKV] + dQKV_quantizer.internal = False + dQKV_quantizer.set_usage(rowwise=True, columnwise=False) + + if fp8_recipe.mxfp8(): + QKV_quantizer.columnwise_usage = True + QKV_quantizer.optimize_for_gemm = True + S_quantizer = None + O_quantizer.columnwise_usage = True + + dO_quantizer.columnwise_usage = True + dO_quantizer.optimize_for_gemm = True + dP_quantizer = None + dQKV_quantizer.columnwise_usage = True return QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer @@ -2331,18 +2450,289 @@ def print_quantizers( type_str = "DS" elif isinstance(q, Float8CurrentScalingQuantizer): type_str = "CS" - print( - f"{label} >> {names[i]:14s}: {type_str}, {q.scale.item():.4e} x" - f" {q.amax.item():.4e} = {q.scale.item()*q.amax.item():.4e}" + elif isinstance(q, MXFP8Quantizer): + type_str = "MXFP8" + if type_str in ["DS", "CS"]: + print( + f"{label} >> {names[i]:14s}: {type_str}, {q.scale.item():.4e} x" + f" {q.amax.item():.4e} = {q.scale.item()*q.amax.item():.4e}" + ) + else: + print(f"{label} >> {names[i]:14s}: {type_str}") + + +def transpose_to_bhsd_htd_pytorch(tensor, src_format): + """Permute to BHSD or HTD format using native PyTorch operations.""" + if src_format in ("bhsd", "htd"): + return tensor + dim_s = src_format.find("s") if "s" in src_format else src_format.find("t") + dim_others = [i for i in range(tensor.ndim) if i != dim_s] + new_dims = [*dim_others[:-1], dim_s, dim_others[-1]] + return tensor.permute(*new_dims).contiguous() + + +def mxfp8_quantize_fast_path(tensor_quantizer_pairs, src_format): + """MXFP8 attention requires quantization along S and D dimensions. This fast path + quantizes tensors without swizzle, and pads, permutes and swizzles the scale_invs + to achieve faster speed due to the smaller sizes of scale_invs compare to the data. + The output tensors have _rowwise_data and _columnwise_data in src_format, and + _rowwise_scale_inv and _columnwise_scale_inv in BHSD format. + + Parameters + ---------- + tensor_quantizer_pairs : list of (torch.Tensor, MXFP8Quantizer) + Each pair is a tensor and its quantizer (with the desired + rowwise_usage / columnwise_usage already set). + src_format : str + Layout of input tensors: ``"bshd"`` or ``"sbhd"``. + All tensors in the list must have the same src_format. + Returns + ------- + fp8_tensors : list of MXFP8Tensors + Data in ``src_format``, scale_inv in BHSD format. + scale_inv_format : str + Always ``"bhsd"``. + """ + if not tensor_quantizer_pairs: + return [], src_format + assert src_format in ( + "bshd", + "sbhd", + ), f"mxfp8_quantize_fast_path only supports bshd/sbhd, got {src_format!r}." + _s_dim = {"bshd": 1, "sbhd": 0} + _d_dim = {"bshd": 3, "sbhd": 3} + + fp8_tensors = [] + for tensor, quantizer in tensor_quantizer_pairs: + original_shape = tensor.shape + rs_shape = list(original_shape) + rs_shape[_d_dim[src_format]] //= MXFP8_BLOCK_SCALING_SIZE + cs_shape = list(original_shape) + cs_shape[_s_dim[src_format]] //= MXFP8_BLOCK_SCALING_SIZE + + # view tensor as 2D for quantization + # BSHD -> (B*S, H*D) + # SBHD -> (S, B*H*D) + if src_format == "bshd": + tensor = tensor.view(*tensor.shape[:2], -1) + else: + tensor = tensor.view(tensor.shape[0], -1) + + # quantize + orig_optimize = quantizer.optimize_for_gemm + quantizer.optimize_for_gemm = False + fp8_tensor = quantizer(tensor) + quantizer.optimize_for_gemm = orig_optimize + + # reshape rowwise/columnwise data to original shape + fp8_tensor._rowwise_data = ( + fp8_tensor._rowwise_data.view(original_shape) + if fp8_tensor._rowwise_data is not None + else None + ) + fp8_tensor._columnwise_data = ( + fp8_tensor._columnwise_data.view(original_shape) + if fp8_tensor._columnwise_data is not None + else None + ) + fp8_tensor._rowwise_scale_inv = ( + fp8_tensor._rowwise_scale_inv.view(rs_shape) + if fp8_tensor._rowwise_scale_inv is not None + else None + ) + fp8_tensor._columnwise_scale_inv = ( + fp8_tensor._columnwise_scale_inv.view(cs_shape) + if fp8_tensor._columnwise_scale_inv is not None + else None + ) + fp8_tensors.append(fp8_tensor) + + # ---- Pad + permute + swizzle scale_inv to BHSD ---- + rs_list = [t._rowwise_scale_inv for t in fp8_tensors] + cs_list = [t._columnwise_scale_inv for t in fp8_tensors] + + def _align_up(x, a): + return ((x + a - 1) // a) * a + + def _bhsd_shape(src_4d, d_pad): + if src_format == "sbhd": + S, B, H, _ = src_4d.shape + else: + B, S, H, _ = src_4d.shape + return (B, H, S, d_pad) + + def _build_outputs(scale_list, alignment): + entries = [] + total = 0 + for s in scale_list: + if s is None: + entries.append(None) + continue + d_pad = _align_up(s.shape[-1], alignment) + shape = _bhsd_shape(s, d_pad) + numel = 1 + for dim in shape: + numel *= dim + entries.append((total, numel, shape)) + total += numel + if total == 0: + return [None] * len(scale_list) + device = next(s for s in scale_list if s is not None).device + buf = torch.empty(total, dtype=torch.uint8, device=device) + return [buf[e[0] : e[0] + e[1]].view(e[2]) if e is not None else None for e in entries] + + # allocate buffers with padding in mind + rs_outs = _build_outputs(rs_list, 4) + cs_outs = _build_outputs(cs_list, 128) + + # permute scale_invs to BHSD; batched + rs_permuted = tex.multi_tensor_transpose_to_bhsd( + rs_list, + original_format=src_format, + outputs=rs_outs, + ) + cs_permuted = tex.multi_tensor_transpose_to_bhsd( + cs_list, + original_format=src_format, + outputs=cs_outs, + ) + + # build output tensors + result = [] + for t, rp, cp in zip(fp8_tensors, rs_permuted, cs_permuted): + rp = rp.view(-1, rp.shape[-1]) if rp is not None else None + cp = cp.view(-1, cp.shape[-1]) if cp is not None else None + result.append( + MXFP8Tensor( + shape=t.shape, + dtype=t.dtype, + rowwise_data=t._rowwise_data, + rowwise_scale_inv=rp, + columnwise_data=t._columnwise_data, + columnwise_scale_inv=cp, + quantizer=t._quantizer, + requires_grad=False, + fp8_dtype=t._fp8_dtype, + with_gemm_swizzled_scales=t._with_gemm_swizzled_scales, ) + ) + # swizzle in place; batched + tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(result, True, False) + tex.multi_tensor_swizzle_scales_for_gemm_unchecked_(result, False, True) + for t in result: + t._with_gemm_swizzled_scales = True + + return result, "bhsd" + + +def combine_and_quantize( + qkv_layout, + q, + k, + v, + qkv_quantizer, + used_in_forward=True, + used_in_backward=False, + keep_same_data_and_scale_inv_format=False, +): + """Combine Q, K, V tensors based on qkv_layout and quantize them together.""" + if isinstance(qkv_quantizer, MXFP8Quantizer): + qkv_format, q_format, kv_format = get_qkv_format(qkv_layout) + assert qkv_format in ("bshd", "sbhd"), ( + "combine_and_quantize only supports bshd/sbhd for MXFP8 quantization, got" + f" {qkv_format!r}." + ) + + _s_dim = {"sbhd": 0, "bshd": 1} + _d_dim = {"sbhd": 3, "bshd": 3} + d_qk = q.shape[_d_dim[qkv_format]] + d_v = v.shape[_d_dim[qkv_format]] + s_q = q.shape[_s_dim[q_format]] + s_kv = v.shape[_s_dim[kv_format]] + assert s_q % 128 == 0 and s_kv % 128 == 0 and d_qk % 32 == 0 and d_v % 32 == 0, ( + "MXFP8 quantization requires s_q % 128 == 0, s_kv % 128 == 0, d_qk % 32 == 0, d_v % 32" + f" == 0. Found {s_q=}, {s_kv=}, {d_qk=}, {d_v=}." + ) + + if qkv_layout not in ("bshd_bshd_bshd", "sbhd_sbhd_sbhd"): + keep_same_data_and_scale_inv_format = True + + # ---- Fast path: quantize in original layout, permute scale_inv to BHSD, then swizzle ---- + if not keep_same_data_and_scale_inv_format: + q_quantizer, k_quantizer, v_quantizer = [qkv_quantizer.copy() for _ in range(3)] + if used_in_forward and not used_in_backward: + q_quantizer.rowwise_usage = True + q_quantizer.columnwise_usage = False + k_quantizer.rowwise_usage = True + k_quantizer.columnwise_usage = False + v_quantizer.rowwise_usage = False + v_quantizer.columnwise_usage = True + elif (not used_in_forward) and used_in_backward: + q_quantizer.rowwise_usage = True + q_quantizer.columnwise_usage = True + k_quantizer.rowwise_usage = True + k_quantizer.columnwise_usage = True + v_quantizer.rowwise_usage = True + v_quantizer.columnwise_usage = False + (q_fp8, k_fp8, v_fp8), qkv_scale_inv_format = mxfp8_quantize_fast_path( + [(q, q_quantizer), (k, k_quantizer), (v, v_quantizer)], qkv_format + ) + return q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format + + # ---- Slow path: permute data to BHSD, then quantize with swizzle ---- + if qkv_layout in ("bshd_bshd_bshd", "sbhd_sbhd_sbhd"): + q, k, v = tex.multi_tensor_transpose_to_bhsd( + [q, k, v], + original_format=qkv_format, + ) + else: + q = transpose_to_bhsd_htd_pytorch(q, q_format) + k = transpose_to_bhsd_htd_pytorch(k, kv_format) + v = transpose_to_bhsd_htd_pytorch(v, kv_format) + qkv_layout = "bhsd_bhsd_bhsd" + qkv_scale_inv_format = "bhsd" + + original_shapes = [x.shape for x in [q, k, v]] + q, k, v = [x.view(-1, x.shape[-1]) for x in [q, k, v]] + + q_quantizer, k_quantizer, v_quantizer = [qkv_quantizer.copy() for _ in range(3)] + if used_in_forward and not used_in_backward: + q_quantizer.rowwise_usage = True + q_quantizer.columnwise_usage = False + k_quantizer.rowwise_usage = True + k_quantizer.columnwise_usage = False + v_quantizer.rowwise_usage = False + v_quantizer.columnwise_usage = True + elif (not used_in_forward) and used_in_backward: + q_quantizer.rowwise_usage = True + q_quantizer.columnwise_usage = True + k_quantizer.rowwise_usage = True + k_quantizer.columnwise_usage = True + v_quantizer.rowwise_usage = True + v_quantizer.columnwise_usage = False + q_fp8, k_fp8, v_fp8 = [ + quant(x) for quant, x in zip([q_quantizer, k_quantizer, v_quantizer], [q, k, v]) + ] + + for fp8_tensor, shape in zip([q_fp8, k_fp8, v_fp8], original_shapes): + fp8_tensor._rowwise_data = ( + fp8_tensor._rowwise_data.view(shape) + if fp8_tensor._rowwise_data is not None + else None + ) + fp8_tensor._columnwise_data = ( + fp8_tensor._columnwise_data.view(shape) + if fp8_tensor._columnwise_data is not None + else None + ) + + return q_fp8, k_fp8, v_fp8, qkv_layout, qkv_scale_inv_format -def combine_and_quantize(qkv_layout, q, k, v, qkv_quantizer): - """Combine q,k,v based on qkv_layout and quantize them together""" - # 1: qkv packed, 2: kv packed, 3: qkv separate qkv_layout = qkv_layout.replace("paged_kv_", "") qkv_group = len(qkv_layout.split("_")) src_nominal_dtype = q.dtype + # 1: qkv packed, 2: kv packed, 3: qkv separate match qkv_group: case 1: dim = qkv_layout.find("3") @@ -2382,24 +2772,28 @@ def combine_and_quantize(qkv_layout, q, k, v, qkv_quantizer): for x in [q_data, k_data, v_data] ] - return q_fp8, k_fp8, v_fp8 + return q_fp8, k_fp8, v_fp8, qkv_layout, None def combine_and_dequantize( qkv_layout, q_fp8, k_fp8, v_fp8, src_nominal_dtype=None, des_nominal_dtype=None ): """Combine q,k,v based on qkv_layout and dequantize them together""" - # 1: qkv packed, 2: kv packed, 3: qkv separate - qkv_layout = qkv_layout.replace("paged_kv_", "") - qkv_group = len(qkv_layout.split("_")) - if all(isinstance(x, Float8Tensor) for x in [q_fp8, k_fp8, v_fp8]): + if all(isinstance(x, QuantizedTensorStorage) for x in [q_fp8, k_fp8, v_fp8]): src_nominal_dtype = q_fp8.dtype else: assert src_nominal_dtype is not None, "The nominal dtype of input tensors is required!" if des_nominal_dtype is None: des_nominal_dtype = src_nominal_dtype + if all(isinstance(x, (MXFP8Tensor, MXFP8TensorStorage)) for x in [q_fp8, k_fp8, v_fp8]): + q, k, v = [x.dequantize(dtype=des_nominal_dtype) for x in [q_fp8, k_fp8, v_fp8]] + return q, k, v + + qkv_layout = qkv_layout.replace("paged_kv_", "") + qkv_group = len(qkv_layout.split("_")) q_data, k_data, v_data = [x._data for x in [q_fp8, k_fp8, v_fp8]] + # 1: qkv packed, 2: kv packed, 3: qkv separate match qkv_group: case 1: dim = qkv_layout.find("3") diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index d95d327c7..afc4622b2 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -795,15 +795,31 @@ def forward( fp8_dpa = fp8_recipe.fp8_dpa fp8_mha = fp8_recipe.fp8_mha float8_current_scaling = fp8_recipe.float8_current_scaling() + mxfp8_scaling = fp8_recipe.mxfp8() else: fp8_dpa = _dpa_fp8_recipe_dpa fp8_mha = _dpa_fp8_recipe_mha float8_current_scaling = _dpa_fp8_recipe == "Float8CurrentScaling" - # QKV Gemm: do not produce FP8 output when in Float8CurrentScaling recipe - qkv_fp8_output = fp8 and fp8_mha and rotary_pos_emb is None and not float8_current_scaling - # DPA: always produce FP8 output when fp8=True to take advantage of the O amax - dpa_fp8_output = fp8 and (fp8_dpa or fp8_mha) - # Proj Gemm: match DPA output except for Float8CurrentScaling + mxfp8_scaling = _dpa_fp8_recipe == "MXFP8BlockScaling" + + # QKV Gemm: do not produce FP8 output when fp8_mha = True if + # 1. RoPE is on: RoPE is only implemented in F16 currently + # 2. FP8CS recipe: due to cuBLAS limitation, FP8CS Gemms can not produce FP8 output + # 3. MXFP8 recipe: QKV Gemm produces QKV in bs(hd), sb(hd), t(hd) shapes, quantization of which would be along + # s/b/t and (hd) dimensions, whereas MXFP8 attention requires quantization along s and d, e.g. bhsd, sbhd, thd + qkv_fp8_output = ( + fp8 + and fp8_mha + and rotary_pos_emb is None + and not float8_current_scaling + and not mxfp8_scaling + ) + # DPA: produce FP8 output to take advantage of O amax from DPA; Projection Gemm can take FP8 or F16 inputs + # 1. FP8DS/FP8CS recipe: produce FP8 output + # 2. MXFP8 recipe: produce F16 output; again, due to quantization dimensions mismatch + dpa_fp8_output = fp8 and (fp8_dpa or fp8_mha) and not mxfp8_scaling + # Projection Gemm: match DPA output except + # 1. FP8CS recipe: produce F16 grads; again, due to cuBLAS limitation proj_fp8_grad = dpa_fp8_output and not float8_current_scaling layernorm_output = None diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 06bfb6ef3..01e139da4 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -35,6 +35,7 @@ } QKVFormat = { + None: NVTE_QKV_Format.NVTE_QKV_Format_NOT_SET, "bshd": NVTE_QKV_Format.NVTE_BSHD, "sbhd": NVTE_QKV_Format.NVTE_SBHD, "thd": NVTE_QKV_Format.NVTE_THD, @@ -42,6 +43,7 @@ "bshd_2sbhd": NVTE_QKV_Format.NVTE_BSHD_2SBHD, "thd_2bshd": NVTE_QKV_Format.NVTE_THD_2BSHD, "thd_2sbhd": NVTE_QKV_Format.NVTE_THD_2SBHD, + "bhsd": NVTE_QKV_Format.NVTE_BHSD, } QKVLayout = { @@ -70,6 +72,7 @@ "paged_kv_sbhd_sbhd_sbhd": NVTE_QKV_Layout.NVTE_Paged_KV_SBHD_SBHD_SBHD, "paged_kv_thd_bshd_bshd": NVTE_QKV_Layout.NVTE_Paged_KV_THD_BSHD_BSHD, "paged_kv_thd_sbhd_sbhd": NVTE_QKV_Layout.NVTE_Paged_KV_THD_SBHD_SBHD, + "bhsd_bhsd_bhsd": NVTE_QKV_Layout.NVTE_BHSD_BHSD_BHSD, } AttnBiasType = { @@ -134,6 +137,8 @@ def fused_attn_fwd( dropout: float = 0.0, fast_zero_fill: bool = True, qkv_layout: str = "sbh3d", + o_format: str = "sbhd", + qkv_scale_inv_format: str = None, attn_bias_type: str = "no_bias", attn_mask_type: str = "padding", softmax_type: str = "vanilla", @@ -203,6 +208,11 @@ def fused_attn_fwd( {"sb3hd", "sbh3d", "sbhd_sb2hd", "sbhd_sbh2d", "sbhd_sbhd_sbhd", "bs3hd", "bsh3d", "bshd_bs2hd", "bshd_bsh2d", "bshd_bshd_bshd", "t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"} + o_format : str, default = "sbhd" + format of O; {"sbhd", "bshd", "thd"} + qkv_scale_inv_format : str, default = None + format of the scale-inverse tensors for QKV; {"sbhd", "bshd", "thd", "bhsd"}; + if None, defaults to the format inferred from qkv_layout. attn_bias_type : str, default = "no_bias" type of the bias; {"no_bias", "pre_scale_bias", "post_scale_bias", "alibi"} attn_mask_type : str, default = "padding" @@ -251,7 +261,7 @@ def fused_attn_fwd( M: torch.Tensor max(Q*K.T) shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32 - ZInv: torch.Tensor + ZInv: torch.Tensor, only allocated for T3HD path 1/sum(e^(x - max(x))), where x=Q*K.T shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32 rng_state: torch.Tensor, optional, if backend is not F16_max512_seqlen @@ -302,17 +312,6 @@ def fused_attn_fwd( rng_elts_per_thread = ( max_seqlen_q * max_seqlen_q + BACKEND_F16m512_FP8_THREADS_PER_CTA - 1 ) // BACKEND_F16m512_FP8_THREADS_PER_CTA - - if s_quantizer is None: - raise ValueError( - "s_quantizer is required for FP8 fused attention forward" - f" (backend={fused_attention_backend}, qkv_layout={qkv_layout!r})." - ) - if o_quantizer is None: - raise ValueError( - "o_quantizer is required for FP8 fused attention forward" - f" (backend={fused_attention_backend}, qkv_layout={qkv_layout!r})." - ) else: raise ValueError(f"Unsupported backend {fused_attention_backend}") @@ -326,6 +325,8 @@ def fused_attn_fwd( dropout, fast_zero_fill, QKVLayout[qkv_layout], + QKVFormat[o_format], + QKVFormat[qkv_scale_inv_format], AttnBiasType[attn_bias_type], AttnMaskType[attn_mask_type], SoftmaxType[softmax_type], @@ -415,7 +416,6 @@ def fused_attn_bwd( o: torch.Tensor, d_o: torch.Tensor, fake_dtype: torch.dtype, - dqkv_dtype: tex.DType, aux_ctx_tensors: List[torch.Tensor], fused_attention_backend: tex.NVTE_Fused_Attn_Backend, cu_seqlens_q_padded: torch.Tensor = None, @@ -427,6 +427,11 @@ def fused_attn_bwd( dropout: float = 0.0, fast_zero_fill: bool = True, qkv_layout: str = "sbh3d", + o_format: str = "sbhd", + do_format: str = "sbhd", + dqkv_layout: str = "sbh3d", + qkv_scale_inv_format: str = None, + do_scale_inv_format: str = None, attn_bias_type: str = "no_bias", attn_mask_type: str = "padding", softmax_type: str = "vanilla", @@ -465,8 +470,6 @@ def fused_attn_bwd( fake_dtype : tex.DType data type of Q, K and V - in case of high precision, fake dtype in case of FP8; in torch.dtype - dqkv_dtype : tex.DType - data type of dQ, dK and dV; in tex.DType, not torch.dtype aux_ctx_tensors : List[torch.Tensor] auxiliary output tensors of the forward pass when its is_training is True, e.g. aux_ctx_tensors = [M, ZInv, rng_state] @@ -482,6 +485,9 @@ def fused_attn_bwd( Quantizer object for the intermediate value dP. dqkv_quantizer : Quantizer, default = None Quantizer object for the output values of the fused_attn_bwd. + attn_scale : float, default = None + if not None, use attn_scale as the attention scale for Q*K.T BMM; + if None, use 1.0/sqrt(head_dim_qk) as the default dropout : float, default = 0.0 dropout probability, 0.0 means no dropout, 1.0 means no output; dropout must be 0.0 if is_training is False @@ -493,6 +499,21 @@ def fused_attn_bwd( {"sb3hd", "sbh3d", "sbhd_sb2hd", "sbhd_sbh2d", "sbhd_sbhd_sbhd", "bs3hd", "bsh3d", "bshd_bs2hd", "bshd_bsh2d", "bshd_bshd_bshd", "t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"} + o_format : str, default = "sbhd" + format of O; {"sbhd", "bshd", "thd"} + do_format : str, default = "sbhd" + format of dO; {"sbhd", "bshd", "thd"} + dqkv_layout : str, default = "sbh3d" + layout of dQ, dK and dV; + {"sb3hd", "sbh3d", "sbhd_sb2hd", "sbhd_sbh2d", "sbhd_sbhd_sbhd", + "bs3hd", "bsh3d", "bshd_bs2hd", "bshd_bsh2d", "bshd_bshd_bshd", + "t3hd", "th3d", "thd_t2hd", "thd_th2d", "thd_thd_thd"} + qkv_scale_inv_format : str, default = None + format of the scale-inverse tensors for QKV; {"sbhd", "bshd", "thd", "bhsd"}; + if None, defaults to the format inferred from qkv_layout. + do_scale_inv_format : str, default = None + format of the scale-inverse tensors for dO; {"sbhd", "bshd", "thd", "bhsd"}; + if None, defaults to the format inferred from the output layout. attn_bias_type : str, default = "no_bias" type of the bias; {"no_bias", "pre_scale_bias", "post_scale_bias", "alibi"} attn_mask_type : str, default = "padding" @@ -553,29 +574,6 @@ def fused_attn_bwd( f" for backend={fused_attention_backend}." ) - if fused_attention_backend == FusedAttnBackend["FP8"]: - if s_quantizer is None: - raise ValueError( - "s_quantizer is required for FP8 fused attention backward" - f" (backend={fused_attention_backend}, qkv_layout={qkv_layout!r})." - ) - if dp_quantizer is None: - raise ValueError( - "dp_quantizer is required for FP8 fused attention backward" - f" (backend={fused_attention_backend}, qkv_layout={qkv_layout!r})." - ) - if dqkv_dtype is None: - raise ValueError( - "dqkv_dtype is required for FP8 fused attention backward" - f" (backend={fused_attention_backend}, qkv_layout={qkv_layout!r})." - ) - if len(aux_ctx_tensors) != 3: - raise ValueError( - "aux_ctx_tensors must be [M, ZInv, rng_state] for FP8 fused attention," - f" but got len(aux_ctx_tensors)={len(aux_ctx_tensors)}" - f" (backend={fused_attention_backend})." - ) - output_tensors = tex.fused_attn_bwd( max_seqlen_q, max_seqlen_kv, @@ -583,6 +581,11 @@ def fused_attn_bwd( dropout, fast_zero_fill, QKVLayout[qkv_layout], + QKVFormat[o_format], + QKVFormat[do_format], + QKVLayout[dqkv_layout], + QKVFormat[qkv_scale_inv_format], + QKVFormat[do_scale_inv_format], AttnBiasType[attn_bias_type], AttnMaskType[attn_mask_type], SoftmaxType[softmax_type], @@ -597,7 +600,6 @@ def fused_attn_bwd( o, d_o, fake_dtype, - dqkv_dtype, aux_ctx_tensors, cu_seqlens_q_padded, cu_seqlens_kv_padded, diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index fb5783dfc..929be8906 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -84,11 +84,11 @@ NVTE_Fused_Attn_Backend get_fused_attn_backend( std::vector fused_attn_fwd( size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, float attn_scale, float p_dropout, - bool set_zero, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - const std::vector window_size, bool bottom_right_diagonal, - const at::Tensor cu_seqlens_q, const at::Tensor cu_seqlens_kv, const py::handle Q, - const py::handle K, const py::handle V, const at::ScalarType fake_dtype, + bool set_zero, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, + NVTE_Softmax_Type softmax_type, const std::vector window_size, + bool bottom_right_diagonal, const at::Tensor cu_seqlens_q, const at::Tensor cu_seqlens_kv, + const py::handle Q, const py::handle K, const py::handle V, const at::ScalarType fake_dtype, const std::optional cu_seqlens_q_padded, const std::optional cu_seqlens_kv_padded, const std::optional page_table_k, const std::optional page_table_v, @@ -98,11 +98,13 @@ std::vector fused_attn_fwd( std::vector fused_attn_bwd( size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float p_dropout, bool set_zero, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, + NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format, + NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, const std::vector window_size, bool bottom_right_diagonal, bool deterministic, const at::Tensor cu_seqlens_q, const at::Tensor cu_seqlens_kv, const py::handle Q, const py::handle K, const py::handle V, - const py::handle O, const py::handle dO, const at::ScalarType fake_dtype, const DType dqkv_type, + const py::handle O, const py::handle dO, const at::ScalarType fake_dtype, const std::vector Aux_CTX_Tensors, const std::optional cu_seqlens_q_padded, const std::optional cu_seqlens_kv_padded, py::handle s_quantizer, @@ -111,6 +113,13 @@ std::vector fused_attn_bwd( at::Tensor fa_prepare_fwd(at::Tensor qkvi); at::Tensor fa_prepare_bwd(at::Tensor q, at::Tensor k, at::Tensor v); +std::vector> multi_tensor_transpose_to_bhsd( + std::vector> inputs, const std::string &original_format, + std::vector> outputs = {}); + +std::vector multi_tensor_pad_last_dim(std::vector inputs, + int64_t alignment); + at::Tensor convert_thd_to_bshd(at::Tensor tensor, at::Tensor cu_seqlens, int b, int max_seq_len); at::Tensor convert_bshd_to_thd(at::Tensor tensor, at::Tensor cu_seqlens, int t); void copy_to_kv_cache(at::Tensor new_k, at::Tensor new_v, at::Tensor k_cache, at::Tensor v_cache, @@ -572,6 +581,13 @@ void fused_multi_row_unpadding(at::Tensor input, at::Tensor output, void inplace_swizzle_scale_for_gemm(py::handle &tensor); +void inplace_multi_tensor_swizzle_scales_for_gemm(std::vector &tensors, + bool rowwise_usage, bool columnwise_usage); + +void inplace_multi_tensor_swizzle_scales_for_gemm_unchecked(std::vector &tensors, + bool rowwise_usage, + bool columnwise_usage); + void grouped_swizzle_for_gemm(py::handle &tensor, bool rowwise, bool columnwise); /*************************************************************************************************** diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index ff60bb87b..8a2e54a73 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -57,7 +57,7 @@ NVTE_Fused_Attn_Backend get_fused_attn_backend( // helper function for S and dP quantizers std::pair quantizer_helper(py::handle quantizer, const std::vector &shape, DType dtype, - bool create_hp_tensor_for_cs, + bool create_hp_tensor, std::optional data) { std::unique_ptr T_quantizer = convert_quantizer(quantizer); TensorWrapper te_T; @@ -78,7 +78,7 @@ std::pair quantizer_helper(py::handle quantizer, } else if (detail::IsFloat8CurrentScalingQuantizers(quantizer.ptr())) { // current scaling auto *T_quantizer_fp8 = dynamic_cast(T_quantizer.get()); - if (create_hp_tensor_for_cs) { + if (create_hp_tensor) { if (data.has_value()) { std::tie(te_T, py_T) = T_quantizer_fp8->create_unquantized_tensor_with_amax(shape, dtype, data.value()); @@ -91,6 +91,20 @@ std::pair quantizer_helper(py::handle quantizer, !data.has_value(), "Float8CurrentScalingQuantizer::create_tensor() does not take data tensor as input!"); } + } else if (detail::IsMXFP8Quantizers(quantizer.ptr())) { + // MXFP8 + if (create_hp_tensor) { + if (data.has_value()) { + std::tie(te_T, py_T) = NoneQuantizer(py::none()).create_tensor(shape, dtype, data.value()); + } else { + std::tie(te_T, py_T) = NoneQuantizer(py::none()).create_tensor(shape, dtype); + } + } else { + auto *T_quantizer_fp8 = dynamic_cast(T_quantizer.get()); + std::tie(te_T, py_T) = T_quantizer_fp8->create_tensor(shape, dtype); + NVTE_CHECK(!data.has_value(), + "MXFP8Quantizer::create_tensor() does not take data tensor as input!"); + } } return {std::move(te_T), std::move(py_T)}; } @@ -98,11 +112,11 @@ std::pair quantizer_helper(py::handle quantizer, // fused attention FWD with separate Q, K and V tensors std::vector fused_attn_fwd( size_t max_seqlen_q, size_t max_seqlen_kv, bool is_training, float attn_scale, float p_dropout, - bool set_zero, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, - const std::vector window_size, bool bottom_right_diagonal, - const at::Tensor cu_seqlens_q, const at::Tensor cu_seqlens_kv, const py::handle Q, - const py::handle K, const py::handle V, const at::ScalarType fake_dtype, + bool set_zero, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, + NVTE_QKV_Format qkv_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, + NVTE_Softmax_Type softmax_type, const std::vector window_size, + bool bottom_right_diagonal, const at::Tensor cu_seqlens_q, const at::Tensor cu_seqlens_kv, + const py::handle Q, const py::handle K, const py::handle V, const at::ScalarType fake_dtype, const std::optional cu_seqlens_q_padded, const std::optional cu_seqlens_kv_padded, const std::optional page_table_k, const std::optional page_table_v, @@ -134,8 +148,13 @@ std::vector fused_attn_fwd( std::unique_ptr O_quantizer = convert_quantizer(o_quantizer); std::vector q_shape = convertShape(te_Q.shape()); std::vector v_shape = convertShape(te_V.shape()); - auto o_shape = std::vector{q_shape.begin(), q_shape.end()}; - o_shape[o_shape.size() - 1] = v_shape[v_shape.size() - 1]; + auto o_shape_tmp = std::vector{q_shape.begin(), q_shape.end()}; + o_shape_tmp[o_shape_tmp.size() - 1] = v_shape[v_shape.size() - 1]; + auto o_shape = std::vector{o_shape_tmp.begin(), o_shape_tmp.end()}; + NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + AttentionShape o_parsed(q_format, o_shape_tmp.data()); + size_t h = o_parsed.h(), d = o_parsed.d(); + o_parsed.to_format(o_format, o_shape.data()); const DType fake_dtype_te = GetTransformerEngineDType(fake_dtype); std::tie(te_O, py_O) = quantizer_helper(o_quantizer, o_shape, fake_dtype_te, true, std::nullopt); @@ -146,9 +165,7 @@ std::vector fused_attn_fwd( TensorWrapper te_page_table_k, te_page_table_v; if (qkv_type == DType::kFloat8E4M3 || qkv_type == DType::kFloat8E5M2) { // FP8 - auto h = q_shape[q_shape.size() - 2]; - auto d = q_shape[q_shape.size() - 1]; - if (set_zero && (nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD)) { + if (set_zero && (o_format == NVTE_QKV_Format::NVTE_THD)) { if ((h * d) % block_size == 0) { mha_fill(te_O, cu_seqlens_q.index({torch::indexing::Slice(-1, torch::indexing::None)})); } else { @@ -156,7 +173,7 @@ std::vector fused_attn_fwd( } } } else if (qkv_type == DType::kBFloat16 || qkv_type == DType::kFloat16) { - if (nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD) { + if (o_format == NVTE_QKV_Format::NVTE_THD) { te_O.zero_(at::cuda::getCurrentCUDAStream()); } } else { @@ -235,9 +252,9 @@ std::vector fused_attn_fwd( te_O.data(), &nvte_aux_tensor_pack, te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), te_page_table_k.data(), te_page_table_v.data(), te_rng_state.data(), max_seqlen_q, max_seqlen_kv, is_training, - return_max_logit, cuda_graph, attn_scale, p_dropout, qkv_layout, bias_type, attn_mask_type, - softmax_type, window_size[0], window_size[1], bottom_right_diagonal, workspace.data(), - at::cuda::getCurrentCUDAStream()); + return_max_logit, cuda_graph, attn_scale, p_dropout, qkv_layout, o_format, + qkv_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size[0], + window_size[1], bottom_right_diagonal, workspace.data(), at::cuda::getCurrentCUDAStream()); }); // allocate memory for workspace and auxiliary output tensors @@ -260,7 +277,7 @@ std::vector fused_attn_fwd( // f16_arbitrary: // return_max_logit=false: S [b, h, sq, 1], rng_state [2], (optional) Bias [1, h, sq, skv], (optional) SoftmaxOffset [1, h, 1, 1] // return_max_logit=true: S [b, h, sq, 1], Max [b, h, sq, 1], rng_state [2], (optional) Bias [1, h, sq, skv], (optional) SoftmaxOffset [1, h, 1, 1] - // fp8 : M [b, h, sq, 1], ZInv [b, h, sq, 1], rng_state [2] + // fp8 : M [b, h, sq, 1], optional ZInv [b, h, sq, 1] (T3HD path), rng_state [2] size_t i = 0; at::Tensor output_tensor; // intermediate softmax tensor, S or M (for fp8) @@ -268,8 +285,10 @@ std::vector fused_attn_fwd( allocateSpace(nvte_shape_to_vector(nvte_tensor_shape(nvte_aux_tensor_pack.tensors[i])), static_cast(nvte_tensor_type(nvte_aux_tensor_pack.tensors[i])), false); set_tensor_param(i++, output_tensor); - // fp8 has an additional softmax stats tensor, ZInv; return_max_logit=true has an additional Max tensor - if (return_max_logit || qkv_type == DType::kFloat8E4M3 || qkv_type == DType::kFloat8E5M2) { + // fp8 T3HD has an additional softmax stats tensor, ZInv; return_max_logit=true has an additional Max tensor + if (((qkv_type == DType::kFloat8E4M3 || qkv_type == DType::kFloat8E5M2) && + qkv_layout == NVTE_QKV_Layout::NVTE_T3HD) || + return_max_logit) { output_tensor = allocateSpace(nvte_shape_to_vector(nvte_tensor_shape(nvte_aux_tensor_pack.tensors[i])), static_cast(nvte_tensor_type(nvte_aux_tensor_pack.tensors[i])), false); @@ -295,9 +314,9 @@ std::vector fused_attn_fwd( te_O.data(), &nvte_aux_tensor_pack, te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), te_page_table_k.data(), te_page_table_v.data(), te_rng_state.data(), max_seqlen_q, max_seqlen_kv, is_training, - return_max_logit, cuda_graph, attn_scale, p_dropout, qkv_layout, bias_type, attn_mask_type, - softmax_type, window_size[0], window_size[1], bottom_right_diagonal, workspace.data(), - at::cuda::getCurrentCUDAStream()); + return_max_logit, cuda_graph, attn_scale, p_dropout, qkv_layout, o_format, + qkv_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size[0], + window_size[1], bottom_right_diagonal, workspace.data(), at::cuda::getCurrentCUDAStream()); }); // destroy tensor wrappers, but not allocated memory @@ -310,11 +329,13 @@ std::vector fused_attn_fwd( // fused attention BWD with separate Q, K and V std::vector fused_attn_bwd( size_t max_seqlen_q, size_t max_seqlen_kv, float attn_scale, float p_dropout, bool set_zero, - NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, + NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, + NVTE_QKV_Layout dqkv_layout, NVTE_QKV_Format qkv_scale_inv_format, + NVTE_QKV_Format do_scale_inv_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type attn_mask_type, NVTE_Softmax_Type softmax_type, const std::vector window_size, bool bottom_right_diagonal, bool deterministic, const at::Tensor cu_seqlens_q, const at::Tensor cu_seqlens_kv, const py::handle Q, const py::handle K, const py::handle V, - const py::handle O, const py::handle dO, const at::ScalarType fake_dtype, const DType dqkv_type, + const py::handle O, const py::handle dO, const at::ScalarType fake_dtype, const std::vector Aux_CTX_Tensors, const std::optional cu_seqlens_q_padded, const std::optional cu_seqlens_kv_padded, py::handle s_quantizer, @@ -343,25 +364,37 @@ std::vector fused_attn_bwd( std::vector q_shape = convertShape(te_Q.shape()); std::vector k_shape = convertShape(te_K.shape()); std::vector v_shape = convertShape(te_V.shape()); - auto h_q = q_shape[q_shape.size() - 2]; - auto h_kv = k_shape[k_shape.size() - 2]; - auto d_qk = q_shape[q_shape.size() - 1]; - const DType fake_dtype_te = GetTransformerEngineDType(fake_dtype); - + const DType dqkv_fake_dtype = GetTransformerEngineDType(fake_dtype); + size_t ndim_q = q_shape.size(); + size_t ndim_kv = k_shape.size(); + std::vector dQ_shape(ndim_q), dK_shape(ndim_kv), dV_shape(ndim_kv); + NVTE_QKV_Format q_format = nvte_get_q_format(qkv_layout); + NVTE_QKV_Format kv_format = nvte_get_kv_format(qkv_layout); + NVTE_QKV_Format dq_format = nvte_get_q_format(dqkv_layout); + NVTE_QKV_Format dkv_format = nvte_get_kv_format(dqkv_layout); + AttentionShape q_parsed(q_format, q_shape.data()); + size_t h_q = q_parsed.h(), d_qk = q_parsed.d(); + q_parsed.to_format(dq_format, dQ_shape.data()); + AttentionShape k_parsed(kv_format, k_shape.data()); + size_t h_kv = k_parsed.h(); + k_parsed.to_format(dkv_format, dK_shape.data()); + AttentionShape v_parsed(kv_format, v_shape.data()); + size_t d_v = v_parsed.d(); + v_parsed.to_format(dkv_format, dV_shape.data()); at::Tensor dQ, dK, dV, dQKV, dKV; - NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(qkv_layout); - std::vector tmp_shape; - auto options = torch::TensorOptions().dtype(GetATenDType(dqkv_type)).device(torch::kCUDA); - if (dqkv_type == DType::kFloat8E4M3 || dqkv_type == DType::kFloat8E5M2) { + // FP16/BF16: dqkv_fake_dtype = kFloat16/kBFloat16, dQ/dK/dV.dtype = torch.float16/torch.bfloat16 + // FP8DS: dqkv_fake_dtype = kFloat16/kBFloat16, dQ/dK/dV.dtype = torch.uint8 + // FP8CS/MXFP8: dqkv_fake_dtype = kFloat16/kBFloat16, dQ/dK/dV.dtype = torch.float16/torch.bfloat16 + auto options = torch::TensorOptions().dtype(fake_dtype).device(torch::kCUDA); + if (detail::IsFloat8Quantizers(dqkv_quantizer.ptr())) { options = options.dtype(torch::kUInt8); } - if (detail::IsFloat8CurrentScalingQuantizers(dqkv_quantizer.ptr())) { - options = options.dtype(fake_dtype); - } + NVTE_QKV_Layout_Group layout_group = nvte_get_qkv_layout_group(dqkv_layout); + std::vector tmp_shape; switch (layout_group) { case NVTE_QKV_Layout_Group::NVTE_3HD: - tmp_shape = std::vector{q_shape.begin(), q_shape.end()}; + tmp_shape = std::vector{dQ_shape.begin(), dQ_shape.end()}; tmp_shape.insert(tmp_shape.begin() + tmp_shape.size() - 2, int64_t(3)); dQKV = torch::empty(c10::IntArrayRef(tmp_shape), options); dQ = dQKV.index({"...", torch::indexing::Slice(0, 1, 1), @@ -378,7 +411,7 @@ std::vector fused_attn_bwd( .squeeze(tmp_shape.size() - 3); break; case NVTE_QKV_Layout_Group::NVTE_H3D: - tmp_shape = std::vector{q_shape.begin(), q_shape.end()}; + tmp_shape = std::vector{dQ_shape.begin(), dQ_shape.end()}; tmp_shape.insert(tmp_shape.begin() + tmp_shape.size() - 1, int64_t(3)); dQKV = torch::empty(c10::IntArrayRef(tmp_shape), options); dQ = dQKV.index({"...", torch::indexing::Slice(0, 1, 1), @@ -392,9 +425,9 @@ std::vector fused_attn_bwd( .squeeze(tmp_shape.size() - 2); break; case NVTE_QKV_Layout_Group::NVTE_HD_2HD: - tmp_shape = std::vector(q_shape.begin(), q_shape.end()); + tmp_shape = std::vector(dQ_shape.begin(), dQ_shape.end()); dQ = torch::empty(tmp_shape, options); - tmp_shape = std::vector{k_shape.begin(), k_shape.end()}; + tmp_shape = std::vector{dK_shape.begin(), dK_shape.end()}; tmp_shape.insert(tmp_shape.begin() + tmp_shape.size() - 2, int64_t(2)); dKV = torch::empty(c10::IntArrayRef(tmp_shape), options); dK = dKV.index({"...", torch::indexing::Slice(0, 1, 1), @@ -407,9 +440,9 @@ std::vector fused_attn_bwd( .squeeze(tmp_shape.size() - 3); break; case NVTE_QKV_Layout_Group::NVTE_HD_H2D: - tmp_shape = std::vector(q_shape.begin(), q_shape.end()); + tmp_shape = std::vector(dQ_shape.begin(), dQ_shape.end()); dQ = torch::empty(tmp_shape, options); - tmp_shape = std::vector{k_shape.begin(), k_shape.end()}; + tmp_shape = std::vector{dK_shape.begin(), dK_shape.end()}; tmp_shape.insert(tmp_shape.begin() + tmp_shape.size() - 1, int64_t(2)); dKV = torch::empty(c10::IntArrayRef(tmp_shape), options); dK = dKV.index({"...", torch::indexing::Slice(0, 1, 1), @@ -420,39 +453,51 @@ std::vector fused_attn_bwd( .squeeze(tmp_shape.size() - 2); break; case NVTE_QKV_Layout_Group::NVTE_HD_HD_HD: - tmp_shape = std::vector(q_shape.begin(), q_shape.end()); + case NVTE_QKV_Layout_Group::NVTE_SD_SD_SD: + tmp_shape = std::vector(dQ_shape.begin(), dQ_shape.end()); dQ = torch::empty(tmp_shape, options); - tmp_shape = std::vector(k_shape.begin(), k_shape.end()); + tmp_shape = std::vector(dK_shape.begin(), dK_shape.end()); dK = torch::empty(tmp_shape, options); - tmp_shape = std::vector(v_shape.begin(), v_shape.end()); + tmp_shape = std::vector(dV_shape.begin(), dV_shape.end()); dV = torch::empty(tmp_shape, options); break; default: NVTE_ERROR("QKV layout not supported!"); } - std::tie(te_dQ, py_dQ) = quantizer_helper(dqkv_quantizer, q_shape, fake_dtype_te, true, dQ); - std::tie(te_dK, py_dK) = quantizer_helper(dqkv_quantizer, k_shape, fake_dtype_te, true, dK); - std::tie(te_dV, py_dV) = quantizer_helper(dqkv_quantizer, v_shape, fake_dtype_te, true, dV); + std::tie(te_dQ, py_dQ) = quantizer_helper(dqkv_quantizer, dQ_shape, dqkv_fake_dtype, true, dQ); + std::tie(te_dK, py_dK) = quantizer_helper(dqkv_quantizer, dK_shape, dqkv_fake_dtype, true, dK); + std::tie(te_dV, py_dV) = quantizer_helper(dqkv_quantizer, dV_shape, dqkv_fake_dtype, true, dV); // construct NVTE tensors - if (dqkv_type == DType::kFloat8E4M3 || dqkv_type == DType::kFloat8E5M2) { + if (detail::IsFloat8Quantizers(dqkv_quantizer.ptr())) { // FP8 - if (set_zero && (nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD)) { - if (((h_q * d_qk) % block_size == 0) && ((h_kv * d_qk) % block_size == 0) && - dQ.is_contiguous() && dK.is_contiguous() && dV.is_contiguous()) { - mha_fill(te_dQ, cu_seqlens_q.index({torch::indexing::Slice(-1, torch::indexing::None)})); - mha_fill(te_dK, cu_seqlens_kv.index({torch::indexing::Slice(-1, torch::indexing::None)})); - mha_fill(te_dV, cu_seqlens_kv.index({torch::indexing::Slice(-1, torch::indexing::None)})); - } else { - dQ.fill_(0); - dK.fill_(0); - dV.fill_(0); + if (set_zero) { + if (dq_format == NVTE_QKV_Format::NVTE_THD) { + if (((h_q * d_qk) % block_size == 0) && dQ.is_contiguous()) { + mha_fill(te_dQ, cu_seqlens_q.index({torch::indexing::Slice(-1, torch::indexing::None)})); + } else { + dQ.fill_(0); + } + } + if (dkv_format == NVTE_QKV_Format::NVTE_THD) { + if (((h_kv * d_qk) % block_size == 0) && ((h_kv * d_v) % block_size == 0) && + dK.is_contiguous() && dV.is_contiguous()) { + mha_fill(te_dK, cu_seqlens_kv.index({torch::indexing::Slice(-1, torch::indexing::None)})); + mha_fill(te_dV, cu_seqlens_kv.index({torch::indexing::Slice(-1, torch::indexing::None)})); + } else { + dK.fill_(0); + dV.fill_(0); + } } } - } else if (dqkv_type == DType::kBFloat16 || dqkv_type == DType::kFloat16) { - if (nvte_get_qkv_format(qkv_layout) == NVTE_QKV_Format::NVTE_THD) { + } else if (dqkv_quantizer.is_none() || + detail::IsFloat8CurrentScalingQuantizers(dqkv_quantizer.ptr()) || + detail::IsMXFP8Quantizers(dqkv_quantizer.ptr())) { + if (dq_format == NVTE_QKV_Format::NVTE_THD) { dQ.fill_(0); + } + if (dkv_format == NVTE_QKV_Format::NVTE_THD) { dK.fill_(0); dV.fill_(0); } @@ -538,7 +583,8 @@ std::vector fused_attn_bwd( &nvte_aux_tensor_pack, te_dQ.data(), te_dK.data(), te_dV.data(), te_dBias.data(), te_dSoftmaxOffset.data(), te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), max_seqlen_q, max_seqlen_kv, - attn_scale, p_dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size[0], + attn_scale, p_dropout, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, + do_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size[0], window_size[1], bottom_right_diagonal, deterministic, cuda_graph, workspace.data(), at::cuda::getCurrentCUDAStream()); }); @@ -555,7 +601,8 @@ std::vector fused_attn_bwd( &nvte_aux_tensor_pack, te_dQ.data(), te_dK.data(), te_dV.data(), te_dBias.data(), te_dSoftmaxOffset.data(), te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), te_cu_seqlens_q_padded.data(), te_cu_seqlens_kv_padded.data(), max_seqlen_q, max_seqlen_kv, - attn_scale, p_dropout, qkv_layout, bias_type, attn_mask_type, softmax_type, window_size[0], + attn_scale, p_dropout, qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, + do_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size[0], window_size[1], bottom_right_diagonal, deterministic, cuda_graph, workspace.data(), at::cuda::getCurrentCUDAStream()); }); @@ -614,6 +661,135 @@ at::Tensor fa_prepare_bwd(at::Tensor q, at::Tensor k, at::Tensor v) { return qkv; } +std::vector> multi_tensor_transpose_to_bhsd( + std::vector> inputs, const std::string &original_format, + std::vector> outputs) { + NVTE_CHECK(original_format == "sbhd" || original_format == "bshd", + "multi_tensor_transpose_to_bhsd: only BSHD/SBHD -> BHSD is currently supported. " + "Got original_format=\"", + original_format, "\"."); + const auto original_format_enum = (original_format == "sbhd") ? NVTE_SBHD : NVTE_BSHD; + + if (inputs.empty()) return {}; + + const bool has_outputs = !outputs.empty(); + if (has_outputs) { + NVTE_CHECK(outputs.size() == inputs.size(), "multi_tensor_transpose_to_bhsd: outputs.size() (", + outputs.size(), ") != inputs.size() (", inputs.size(), ")."); + } + + std::vector te_ins, te_outs; + std::vector> result(inputs.size(), std::nullopt); + + for (size_t i = 0; i < inputs.size(); ++i) { + if (!inputs[i].has_value()) continue; + + auto &input = inputs[i].value(); + NVTE_CHECK(input.is_cuda() && input.dim() == 4, "multi_tensor_transpose_to_bhsd: input ", i, + " must be a 4D CUDA tensor."); + input = input.contiguous(); + NVTE_CHECK(input.scalar_type() == at::ScalarType::Half || + input.scalar_type() == at::ScalarType::BFloat16 || + input.scalar_type() == at::ScalarType::Byte, + "multi_tensor_transpose_to_bhsd: unsupported dtype at index ", i, "."); + + at::Tensor output; + if (has_outputs && outputs[i].has_value()) { + output = outputs[i].value(); + } else { + int64_t B, S, H, D; + if (original_format_enum == NVTE_SBHD) { + S = input.size(0); + B = input.size(1); + H = input.size(2); + D = input.size(3); + } else { + B = input.size(0); + S = input.size(1); + H = input.size(2); + D = input.size(3); + } + output = at::empty({B, H, S, D}, input.options()); + } + + te_ins.push_back(makeTransformerEngineTensor(input)); + te_outs.push_back(makeTransformerEngineTensor(output)); + result[i] = output; + } + + if (!te_ins.empty()) { + std::vector nvte_ins(te_ins.size()), nvte_outs(te_outs.size()); + for (size_t j = 0; j < te_ins.size(); ++j) { + nvte_ins[j] = te_ins[j].data(); + nvte_outs[j] = te_outs[j].data(); + } + nvte_multi_tensor_transpose_to_bhsd(nvte_ins.data(), nvte_outs.data(), te_ins.size(), + original_format_enum, at::cuda::getCurrentCUDAStream()); + } + + return result; +} + +std::vector multi_tensor_pad_last_dim(std::vector inputs, + int64_t alignment) { + const auto align = static_cast(alignment); + NVTE_CHECK(align > 0, "multi_tensor_pad_last_dim: alignment must be > 0."); + NVTE_CHECK(!inputs.empty(), "multi_tensor_pad_last_dim: inputs must not be empty."); + + auto stream = at::cuda::getCurrentCUDAStream(); + std::vector outputs; + outputs.reserve(inputs.size()); + + std::vector kernel_indices; + + for (size_t i = 0; i < inputs.size(); ++i) { + auto &input = inputs[i]; + + NVTE_CHECK(input.dim() == 2, "multi_tensor_pad_last_dim: expected 2D input at index ", i, + ", got ", input.dim(), "D."); + NVTE_CHECK(input.is_cuda(), "multi_tensor_pad_last_dim: input must be a CUDA tensor at index ", + i, "."); + input = input.contiguous(); + + const int64_t rows = input.size(0); + const int64_t in_cols = input.size(1); + const int64_t padded_cols = + static_cast(DIVUP_TO_MULTIPLE(static_cast(in_cols), align)); + + if (in_cols == padded_cols) { + outputs.push_back(input); + continue; + } + + at::Tensor output = at::empty({rows, padded_cols}, input.options()); + outputs.push_back(output); + kernel_indices.push_back(outputs.size() - 1); + } + + if (kernel_indices.empty()) return outputs; + + std::vector te_in_wrappers, te_out_wrappers; + te_in_wrappers.reserve(kernel_indices.size()); + te_out_wrappers.reserve(kernel_indices.size()); + + for (size_t idx : kernel_indices) { + te_in_wrappers.push_back(makeTransformerEngineTensor(inputs[idx])); + te_out_wrappers.push_back(makeTransformerEngineTensor(outputs[idx])); + } + + std::vector nvte_inputs(te_in_wrappers.size()); + std::vector nvte_outputs(te_out_wrappers.size()); + for (size_t i = 0; i < te_in_wrappers.size(); ++i) { + nvte_inputs[i] = te_in_wrappers[i].data(); + nvte_outputs[i] = te_out_wrappers[i].data(); + } + + nvte_multi_tensor_pad_last_dim(nvte_inputs.data(), nvte_outputs.data(), te_in_wrappers.size(), + stream); + + return outputs; +} + /*************************************************************************************************** * Support THD format for Context Parallel: Read the half of a THD tensor **************************************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 27d26d3da..eb7576d90 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -391,6 +391,15 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Fused Multi-tensor unpadding", py::call_guard()); m.def("swizzle_scales_for_gemm_", &transformer_engine::pytorch::inplace_swizzle_scale_for_gemm, "Convert tensor block scales into GEMM swizzled format"); + m.def("multi_tensor_swizzle_scales_for_gemm_", + &transformer_engine::pytorch::inplace_multi_tensor_swizzle_scales_for_gemm, + "Convert multiple tensors' block scales into GEMM swizzled format", py::arg("tensors"), + py::arg("rowwise_usage"), py::arg("columnwise_usage")); + m.def( + "multi_tensor_swizzle_scales_for_gemm_unchecked_", + &transformer_engine::pytorch::inplace_multi_tensor_swizzle_scales_for_gemm_unchecked, + "Convert multiple tensors' block scales into GEMM swizzled format (skip scale shape checks)", + py::arg("tensors"), py::arg("rowwise_usage"), py::arg("columnwise_usage")); m.def("grouped_swizzle_for_gemm", &transformer_engine::pytorch::grouped_swizzle_for_gemm, "In-place swizzle of grouped tensor scales for GEMM", py::arg("tensor"), py::arg("rowwise"), py::arg("columnwise")); @@ -401,6 +410,14 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("fa_prepare_bwd", &transformer_engine::pytorch::fa_prepare_bwd, "Backward of QKV preparation for Flash Attention", py::call_guard()); + m.def("multi_tensor_transpose_to_bhsd", + &transformer_engine::pytorch::multi_tensor_transpose_to_bhsd, + "Permute multiple tensors from BSHD/SBHD to BHSD.", py::arg("inputs"), + py::arg("original_format"), py::arg("outputs") = std::vector>{}, + py::call_guard()); + m.def("multi_tensor_pad_last_dim", &transformer_engine::pytorch::multi_tensor_pad_last_dim, + "Pad multiple tensors' last dimension to a common alignment.", py::arg("inputs"), + py::arg("alignment"), py::call_guard()); m.def("fused_attn_fwd", &transformer_engine::pytorch::fused_attn_fwd, "Fused Attention FP8/BF16/FP16 FWD with separate Q, K and V"); m.def("fused_attn_bwd", &transformer_engine::pytorch::fused_attn_bwd, diff --git a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp index a6b4e7569..cbaabaad1 100644 --- a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp +++ b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp @@ -141,9 +141,11 @@ std::tuple, std::optional> swizzle_scales_ return {std::move(rowwise_scales_pyt), std::move(columnwise_scales_pyt)}; } -std::optional multi_tensor_swizzle_scales_for_gemm( +namespace { + +std::optional multi_tensor_swizzle_scales_for_gemm_impl( std::vector &tensors, bool rowwise_usage, - bool columnwise_usage) { + bool columnwise_usage, bool check_scale_inv_shapes) { // Checks and trivial cases NVTE_CHECK(rowwise_usage != columnwise_usage, "Expect exactly one of rowwise_usage=", rowwise_usage, @@ -243,9 +245,15 @@ std::optional multi_tensor_swizzle_scales_for_gemm( // Launch kernel NVTE_SCOPED_GIL_RELEASE({ - nvte_multi_tensor_swizzle_scaling_factors(inputs_nvte_raw.data(), outputs_nvte_raw.data(), - inputs_nvte_raw.size(), - at::cuda::getCurrentCUDAStream()); + if (check_scale_inv_shapes) { + nvte_multi_tensor_swizzle_scaling_factors(inputs_nvte_raw.data(), outputs_nvte_raw.data(), + inputs_nvte_raw.size(), + at::cuda::getCurrentCUDAStream()); + } else { + nvte_multi_tensor_swizzle_scaling_factors_unchecked( + inputs_nvte_raw.data(), outputs_nvte_raw.data(), inputs_nvte_raw.size(), + at::cuda::getCurrentCUDAStream()); + } }); // Update tensors with swizzled scales @@ -269,6 +277,22 @@ std::optional multi_tensor_swizzle_scales_for_gemm( return std::move(output_scales_pyt); } +} // anonymous namespace + +std::optional multi_tensor_swizzle_scales_for_gemm( + std::vector &tensors, bool rowwise_usage, + bool columnwise_usage) { + return multi_tensor_swizzle_scales_for_gemm_impl(tensors, rowwise_usage, columnwise_usage, + /*check_scale_inv_shapes=*/true); +} + +std::optional multi_tensor_swizzle_scales_for_gemm_unchecked( + std::vector &tensors, bool rowwise_usage, + bool columnwise_usage) { + return multi_tensor_swizzle_scales_for_gemm_impl(tensors, rowwise_usage, columnwise_usage, + /*check_scale_inv_shapes=*/false); +} + at::Tensor convert_block_scaling_to_mxfp8_tensor(transformer_engine::TensorWrapper &input, bool rowwise) { // Check input tensor @@ -443,6 +467,105 @@ void grouped_swizzle_for_gemm(py::handle &tensor, bool rowwise, bool columnwise) } } +namespace { + +void inplace_multi_tensor_swizzle_scales_for_gemm_impl(std::vector &tensors, + bool rowwise_usage, bool columnwise_usage, + bool check_scale_inv_shapes) { + NVTE_CHECK(rowwise_usage != columnwise_usage, + "Expect exactly one of rowwise_usage and columnwise_usage."); + if (tensors.empty()) { + return; + } + + // Convert Python tensors to TensorWrappers, filtering those that need swizzling + std::vector swizzle_indices; + std::vector wrappers_to_swizzle; + + for (size_t i = 0; i < tensors.size(); ++i) { + auto tw = makeTransformerEngineTensor(tensors[i], py::none()); + + if (i == 0) { + switch (tw.scaling_mode()) { + case NVTE_MXFP8_1D_SCALING: + case NVTE_NVFP4_1D_SCALING: + break; + case NVTE_INVALID_SCALING: + NVTE_ERROR("Invalid scaling mode for swizzling scaling factors."); + default: + return; + } + } + + if (tw.get_with_gemm_swizzled_scales()) { + continue; + } + const auto scales_nvte = + rowwise_usage ? tw.get_rowwise_scale_inv() : tw.get_columnwise_scale_inv(); + if (scales_nvte.data_ptr == nullptr || + (scales_nvte.shape.ndim == 1 && scales_nvte.shape.data[0] == 0)) { + continue; + } + + swizzle_indices.push_back(i); + wrappers_to_swizzle.push_back(std::move(tw)); + } + + if (wrappers_to_swizzle.empty()) { + return; + } + + // Delegate to core C++ function + auto swizzle_fn = check_scale_inv_shapes ? multi_tensor_swizzle_scales_for_gemm + : multi_tensor_swizzle_scales_for_gemm_unchecked; + auto output_buffer = swizzle_fn(wrappers_to_swizzle, rowwise_usage, columnwise_usage); + if (!output_buffer.has_value()) { + return; + } + + // Update Python objects with properly-shaped views into the contiguous output buffer + const uint8_t *base = reinterpret_cast(output_buffer->data_ptr()); + for (size_t j = 0; j < wrappers_to_swizzle.size(); ++j) { + const auto scales_nvte = rowwise_usage ? wrappers_to_swizzle[j].get_rowwise_scale_inv() + : wrappers_to_swizzle[j].get_columnwise_scale_inv(); + + const size_t offset = reinterpret_cast(scales_nvte.data_ptr) - base; + const auto dtype = static_cast(scales_nvte.dtype); + const size_t num_elements = product(scales_nvte.shape, 0, scales_nvte.shape.ndim); + const size_t num_bytes = + ceildiv(num_elements * transformer_engine::pytorch::typeToNumBits(dtype), size_t(8)); + + std::vector torch_shape; + for (size_t d = 0; d < scales_nvte.shape.ndim; ++d) { + torch_shape.push_back(static_cast(scales_nvte.shape.data[d])); + } + auto scale_view = + output_buffer->narrow(0, static_cast(offset), static_cast(num_bytes)) + .view(torch_shape); + + if (rowwise_usage) { + tensors[swizzle_indices[j]].attr("_rowwise_scale_inv") = py::cast(scale_view); + } else { + tensors[swizzle_indices[j]].attr("_columnwise_scale_inv") = py::cast(scale_view); + } + } +} + +} // anonymous namespace + +void inplace_multi_tensor_swizzle_scales_for_gemm(std::vector &tensors, + bool rowwise_usage, bool columnwise_usage) { + inplace_multi_tensor_swizzle_scales_for_gemm_impl(tensors, rowwise_usage, columnwise_usage, + /*check_scale_inv_shapes=*/true); +} + +void inplace_multi_tensor_swizzle_scales_for_gemm_unchecked(std::vector &tensors, + bool rowwise_usage, + bool columnwise_usage) { + inplace_multi_tensor_swizzle_scales_for_gemm_impl(tensors, rowwise_usage, columnwise_usage, + /*check_scale_inv_shapes=*/false); +} + void inplace_swizzle_scale_for_gemm(py::handle &tensor) { // Convert Python tensor to C++ tensor auto tensor_nvte = makeTransformerEngineTensor(tensor, py::none()); diff --git a/transformer_engine/pytorch/csrc/util.h b/transformer_engine/pytorch/csrc/util.h index 88f76a7cb..132db4075 100644 --- a/transformer_engine/pytorch/csrc/util.h +++ b/transformer_engine/pytorch/csrc/util.h @@ -33,6 +33,9 @@ std::optional multi_tensor_swizzle_scales_for_gemm(std::vector multi_tensor_swizzle_scales_for_gemm_unchecked( + std::vector& tensors, bool rowwise_usage, bool columnwise_usage); + using SwizzledGroupedScales = std::pair, std::optional>; /*! \brief Swizzle grouped tensor scales for GEMM if needed. diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 7e2fea45f..5f12c3ed8 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -635,7 +635,7 @@ def make_grouped_tensor( total_columnwise_scale_elements = 0 columnwise_scale_inv_offsets = [0] for i, s in enumerate(shape): - scale_inv_shape = quantizer.get_scale_shape(s, False) + scale_inv_shape = quantizer.get_scale_shape(s, True) columnwise_scale_elements = math.prod(scale_inv_shape) total_columnwise_scale_elements += columnwise_scale_elements columnwise_scale_inv_offsets.append(total_columnwise_scale_elements) @@ -872,15 +872,25 @@ def split_into_quantized_tensors( # populate scale_inv_offsets from the tensor offsets if self.scale_inv is not None and self.scale_inv_offsets is None: - if recipe.nvfp4(): - self.scale_inv_offsets = self.tensor_offsets // 16 - if recipe.mxfp8(): - self.scale_inv_offsets = self.tensor_offsets // 32 + if recipe.nvfp4() or recipe.mxfp8() or recipe.float8_block_scaling(): + cum = 0 + scale_inv_offsets = [0] + for i in range(self.num_tensors): + tensor_shape = self.tensor_shapes[i] + scale_shape = self.quantizer.get_scale_shape(tensor_shape, False) + cum += math.prod(scale_shape) + scale_inv_offsets.append(cum) + self.scale_inv_offsets = scale_inv_offsets if self.columnwise_scale_inv is not None and self.columnwise_scale_inv_offsets is None: - if recipe.nvfp4(): - self.columnwise_scale_inv_offsets = self.tensor_offsets // 16 - if recipe.mxfp8(): - self.columnwise_scale_inv_offsets = self.tensor_offsets // 32 + if recipe.nvfp4() or recipe.mxfp8() or recipe.float8_block_scaling(): + cum = 0 + columnwise_scale_inv_offsets = [0] + for i in range(self.num_tensors): + tensor_shape = self.tensor_shapes[i] + scale_shape = self.quantizer.get_scale_shape(tensor_shape, True) + cum += math.prod(scale_shape) + columnwise_scale_inv_offsets.append(cum) + self.columnwise_scale_inv_offsets = columnwise_scale_inv_offsets for i in range(self.num_tensors): quantizer = self.quantizer From 0be9046db16d75bd41301750949a928aa667b47c Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Tue, 21 Apr 2026 17:15:26 -0700 Subject: [PATCH 006/180] Bias/Dbias Support for GroupedLinear (#2885) * starting grouped linear integration, not tested, grouped_bias_add optimized and uses scales now Signed-off-by: Varun Thumbe * all changes Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * optimize grouped bias add kernel to 4TB/s handling load imbalance Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * kernel optimized + review comments Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove uneecessary load from the kernel Remove unnecessary pointer reinterpretation for bias input. Signed-off-by: vthumbe1503 * No need for text for grouped bias add Signed-off-by: vthumbe1503 * Remove lru_cache from _is_deterministic_mode Removed lru_cache decorator from the _is_deterministic_mode function. Signed-off-by: vthumbe1503 * address more review comments Signed-off-by: Varun Thumbe * for better perf Signed-off-by: Varun Thumbe * remove unecessary comments Signed-off-by: vthumbe1503 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * rename a bit Signed-off-by: Varun Thumbe --------- Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/test_numerics.py | 40 ++- .../common/gemm/cublaslt_grouped_gemm.cu | 311 +++++++++++++----- .../common/include/transformer_engine/gemm.h | 12 +- .../common/triton/grouped_dbias_dscales.py | 69 ++-- .../pytorch/cpp_extensions/gemm.py | 5 + transformer_engine/pytorch/csrc/extensions.h | 7 +- .../pytorch/csrc/extensions/gemm.cpp | 40 ++- .../pytorch/ops/basic/grouped_linear.py | 35 +- .../pytorch/ops/fused/backward_grouped_mlp.py | 4 +- .../pytorch/triton/grouped_dbias_dscales.py | 149 ++++++--- 10 files changed, 482 insertions(+), 190 deletions(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 4bfe06095..5eef7f151 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -2848,6 +2848,27 @@ def _make_grouped_tensor_uniform( ) +def _apply_grouped_bias_ref( + base_outs: List[torch.Tensor], + bias: Optional[List[torch.Tensor]], + bias_scale: Optional[torch.Tensor], + m_sizes: List[int], + dtype: torch.dtype, +) -> List[torch.Tensor]: + """Reference: add (optionally per-row scaled) bias to each group's output, cast to ``dtype``.""" + if bias is None: + return list(base_outs) + if bias_scale is None: + return [(o.float() + b.float()).to(dtype) for o, b in zip(base_outs, bias)] + out = [] + offset = 0 + for i, ms in enumerate(m_sizes): + s = bias_scale[offset : offset + ms].unsqueeze(-1) + out.append((base_outs[i].float() + bias[i].float() * s).to(dtype)) + offset += ms + return out + + @pytest.mark.parametrize( "z, m, n, k", [ @@ -2860,7 +2881,8 @@ def _make_grouped_tensor_uniform( @pytest.mark.parametrize("case", ["no_discrete", "discrete_in", "discrete_out"]) @pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) @pytest.mark.parametrize("accumulate", [False, True]) -def test_grouped_gemm_grouped_tensor(z, m, n, k, case, layout, accumulate) -> None: +@pytest.mark.parametrize("use_bias_scale", [False, True]) +def test_grouped_gemm_grouped_tensor(z, m, n, k, case, layout, accumulate, use_bias_scale) -> None: if tex.get_cublasLt_version() < 130300: pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") if torch.cuda.get_device_capability() < (10, 0): @@ -2914,12 +2936,11 @@ def test_grouped_gemm_grouped_tensor(z, m, n, k, case, layout, accumulate) -> No if case != "discrete_out" else None ) + bias_scale = None + if use_bias_scale and bias is not None and layout != "NT": + bias_scale = torch.randn(m, device="cuda", dtype=torch.float32) # Bias add in grouped kernel accumulates in FP32 for BF16/FP16. - out_ref = ( - [(o.float() + b.float()).to(dtype) for o, b in zip(out_ref_no_bias, bias)] - if bias is not None - else out_ref_no_bias - ) + out_ref = _apply_grouped_bias_ref(out_ref_no_bias, bias, bias_scale, m_sizes, dtype) # Create grouped tensors based on case device = A[0].device grouped_A = A @@ -2983,6 +3004,7 @@ def test_grouped_gemm_grouped_tensor(z, m, n, k, case, layout, accumulate) -> No layout=layout, accumulate=accumulate, bias=grouped_bias, + bias_scale=bias_scale, ) out_grouped_no_bias = ( grouped_out_no_bias @@ -2995,10 +3017,8 @@ def test_grouped_gemm_grouped_tensor(z, m, n, k, case, layout, accumulate) -> No else grouped_out_bias.split_into_quantized_tensors() ) - out_grouped_manual_bias = ( - [(o.float() + b.float()).to(dtype) for o, b in zip(out_grouped_no_bias, bias)] - if bias is not None - else out_grouped_no_bias + out_grouped_manual_bias = _apply_grouped_bias_ref( + out_grouped_no_bias, bias, bias_scale, m_sizes, dtype ) tols = dtype_tols(dtype) for o, o_ref in zip(out_grouped_no_bias, out_ref_no_bias): diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 985c53f76..ed2275b44 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -12,6 +12,7 @@ #include #include +#include #include #include "../common.h" @@ -331,20 +332,20 @@ struct GroupedOperandSelection { bool trans = false; }; -constexpr int kMaxTensorsPerKernel = 64; +constexpr int kMaxGroups = 64; // Arguments for the grouped GEMM kernel that operates on multiple output tensors. struct MultiTensorGroupGemmOutputArgs { - void *data_ptrs[kMaxTensorsPerKernel]; - int rows[kMaxTensorsPerKernel]; - int cols[kMaxTensorsPerKernel]; + void *data_ptrs[kMaxGroups]; + int rows[kMaxGroups]; + int cols[kMaxGroups]; }; // Arguments for the grouped GEMM kernel that operates on multiple inputA tensors. struct MultiTensorGroupGemmInputArgs { - void *data_ptrs[kMaxTensorsPerKernel]; - void *scale_inv_ptrs[kMaxTensorsPerKernel]; - int rows[kMaxTensorsPerKernel]; - int cols[kMaxTensorsPerKernel]; + void *data_ptrs[kMaxGroups]; + void *scale_inv_ptrs[kMaxGroups]; + int rows[kMaxGroups]; + int cols[kMaxGroups]; }; struct MultiTensorListInfo { bool all_row = true; @@ -425,8 +426,8 @@ inline MultiTensorGroupGemmOutputArgs build_grouped_gemm_multi_out_args( "_tensors=", list_size); NVTE_CHECK(list_size == expected_num_tensors, "Grouped GEMM: ", name, "_list must have num_tensors (", expected_num_tensors, ") entries, got ", list_size); - NVTE_CHECK(list_size <= static_cast(kMaxTensorsPerKernel), "Grouped GEMM: ", name, - "_list supports up to ", kMaxTensorsPerKernel, " tensors per kernel, got ", list_size); + NVTE_CHECK(list_size <= static_cast(kMaxGroups), "Grouped GEMM: ", name, + "_list supports up to ", kMaxGroups, " tensors per kernel, got ", list_size); for (size_t i = 0; i < list_size; ++i) { const transformer_engine::Tensor *t = @@ -499,8 +500,8 @@ inline MultiTensorListInfo validate_grouped_gemm_multi_inputA_list(const NVTETen "_tensors=", list_size); NVTE_CHECK(list_size == expected_num_tensors, "Grouped GEMM: ", name, "_list must have num_tensors (", expected_num_tensors, ") entries, got ", list_size); - NVTE_CHECK(list_size <= static_cast(kMaxTensorsPerKernel), "Grouped GEMM: ", name, - "_list supports up to ", kMaxTensorsPerKernel, " tensors per kernel, got ", list_size); + NVTE_CHECK(list_size <= static_cast(kMaxGroups), "Grouped GEMM: ", name, + "_list supports up to ", kMaxGroups, " tensors per kernel, got ", list_size); const transformer_engine::Tensor *t0 = transformer_engine::convertNVTETensorCheck(tensor_list[0]); info.scaling_mode = t0->scaling_mode; @@ -845,43 +846,120 @@ __forceinline__ __device__ int64_t compute_grouped_tensor_offset(const TensorSha } } -// Kernel that performs bias addition to the Grouped GEMM output tensors. -// Bias itself is a grouped tensor with the collections of same number of tensors -// as the output tensors. -template -__global__ void grouped_bias_add_kernel(char *d_base, const char *bias_base, TensorShapeInfo d_meta, - TensorShapeInfo bias_meta, size_t num_tensors) { - const size_t tensor_idx = blockIdx.x; - if (tensor_idx >= num_tensors) return; +// Linear scan to find which tensor contains the given row. +// Returns the tensor index and writes the exclusive end-row of that tensor to *out_tensor_row_end. +__forceinline__ __device__ int find_tensor_for_row(const int64_t *first_dims, int64_t uniform_first, + int row, int num_tensors, + int *out_tensor_row_end) { + int offset = 0; + for (int i = 0; i < num_tensors; i++) { + int dim = first_dims ? static_cast(first_dims[i]) : static_cast(uniform_first); + offset += dim; + if (row < offset) { + *out_tensor_row_end = offset; + return i; + } + } + *out_tensor_row_end = offset; + return num_tensors - 1; +} - const int64_t m = d_meta.first_dims ? d_meta.first_dims[tensor_idx] : d_meta.uniform_first; - const int64_t n = d_meta.last_dims ? d_meta.last_dims[tensor_idx] : d_meta.uniform_last; +// Kernel that performs (optionally scaled) bias addition to Grouped GEMM output tensors. +// SM-filling grid with grid-stride over row chunks. +// 2D grid: blockIdx.x = SM-filling row blocks, blockIdx.y = column chunk. +// Each block grid-strides over kRowsPerBlock-sized row chunks, processing +// all chunks that map to it. Safe when sum(first_dims) <= total_rows. +template +__global__ void grouped_bias_add_kernel(char *__restrict__ d_base, + const char *__restrict__ bias_base, + const float *__restrict__ scale_base, + TensorShapeInfo d_meta, int n, int total_rows, + int num_tensors) { + using VecStorage = transformer_engine::VectorizedStorage; + using VecType = typename VecStorage::LType; + + const int tid = static_cast(threadIdx.x); + const int row_bid = static_cast(blockIdx.x); + const int col_bid = static_cast(blockIdx.y); + const int row_grid_stride = static_cast(gridDim.x); + + // Single-warp reduction to compute valid_rows = sum(first_dims). + // kMaxGroups <= 64 so warp 0 (32 lanes) covers it with <=2 loads each. + __shared__ int s_valid_rows; + if (tid < 32) { + int local_sum = 0; + for (int i = tid; i < num_tensors; i += 32) { + local_sum += d_meta.first_dims ? static_cast(d_meta.first_dims[i]) + : static_cast(d_meta.uniform_first); + } + for (int offset = 16; offset > 0; offset >>= 1) { + local_sum += __shfl_down_sync(0xffffffff, local_sum, offset); + } + if (tid == 0) s_valid_rows = local_sum; + } + __syncthreads(); + const int valid_rows = s_valid_rows; + + const int block_cols = kBlockDim * kVec; + const int col = col_bid * block_cols + tid * kVec; + if (col >= n) return; + + T *__restrict__ d = reinterpret_cast(d_base); + const T *__restrict__ bias = reinterpret_cast(bias_base); + + // Grid-stride loop over row chunks. + for (int chunk_start = row_bid * kRowsPerBlock; chunk_start < valid_rows; + chunk_start += row_grid_stride * kRowsPerBlock) { + const int row_start = chunk_start; + const int row_end = min(row_start + kRowsPerBlock, valid_rows); + + // Linear scan to find the starting row's tensor and its boundary. + int tensor_row_end; + int tensor_idx = find_tensor_for_row(d_meta.first_dims, d_meta.uniform_first, row_start, + num_tensors, &tensor_row_end); + int bias_idx = tensor_idx * n; + + VecStorage b_in; + + // Walk tensor segments within this chunk's row range. + int seg_start = row_start; + while (seg_start < row_end) { + while (tensor_idx < num_tensors - 1 && tensor_row_end <= seg_start) { + tensor_idx++; + bias_idx += n; + int dim = d_meta.first_dims ? static_cast(d_meta.first_dims[tensor_idx]) + : static_cast(d_meta.uniform_first); + tensor_row_end += dim; + } + b_in.scratch_.aligned = *reinterpret_cast(bias + bias_idx + col); + const int seg_end = min(tensor_row_end, row_end); - const int64_t d_offset = compute_grouped_tensor_offset(d_meta, tensor_idx); - const int64_t bias_offset = compute_grouped_tensor_offset(bias_meta, tensor_idx); + for (int row = seg_start; row < seg_end; row++) { + T *d_ptr = d + row * n + col; + VecStorage d_in; + d_in.scratch_.aligned = *reinterpret_cast(d_ptr); - auto *d_ptr = reinterpret_cast(d_base + d_offset * sizeof(T)); - const auto *bias_ptr = reinterpret_cast(bias_base + bias_offset * sizeof(T)); + [[maybe_unused]] float s_val; + if constexpr (UseScale) s_val = scale_base[row]; - const int64_t elements = m * n; - const int64_t vec_count = elements / kVec; - using VecStorage = transformer_engine::VectorizedStorage; - using VecType = typename VecStorage::LType; - transformer_engine::VectorizedLoader loader(d_ptr, elements); - transformer_engine::VectorizedStorer storer(d_ptr, elements); - const int64_t vec_id = static_cast(blockIdx.y) * blockDim.x + threadIdx.x; - if (vec_id >= vec_count) return; - const int64_t vec_start = vec_id * kVec; - const int64_t col = vec_start % n; - loader.load(vec_id, elements); - const auto *b_vec = reinterpret_cast(bias_ptr + col); - VecStorage b_in; - b_in.scratch_.aligned = *b_vec; #pragma unroll - for (int i = 0; i < kVec; ++i) { - storer.separate()[i] = loader.separate()[i] + b_in.scratch_.separate[i]; + for (int i = 0; i < kVec; ++i) { + if constexpr (UseScale) { + d_in.scratch_.separate[i] = + static_cast(fmaf(static_cast(b_in.scratch_.separate[i]), s_val, + static_cast(d_in.scratch_.separate[i]))); + } else { + d_in.scratch_.separate[i] = + static_cast(static_cast(d_in.scratch_.separate[i]) + + static_cast(b_in.scratch_.separate[i])); + } + } + *reinterpret_cast(d_ptr) = d_in.scratch_.aligned; + } + + seg_start = seg_end; + } } - storer.store(vec_id, elements); } // Single kernel that sets up all GEMM parameters. @@ -1307,54 +1385,121 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, workspace.cublas_workspace_ptr, stream, config_.sm_count); } -void nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, - cudaStream_t stream) { - NVTE_API_CALL(nvte_grouped_bias_add); - using namespace transformer_engine; +namespace { - const GroupedTensor *outputD = convertNVTEGroupedTensorCheck(output); - const GroupedTensor *bias_tensor = convertNVTEGroupedTensorCheck(bias); +void launch_grouped_bias_add(const transformer_engine::GroupedTensor *outputD, + const transformer_engine::GroupedTensor *bias_tensor, + const float *scale_ptr, bool use_scale, cudaStream_t stream) { + using namespace transformer_engine; - NVTE_CHECK(outputD->num_tensors >= 1, "Grouped bias add: number of tensors must be at least 1"); - NVTE_CHECK(outputD->num_tensors == bias_tensor->num_tensors, - "Grouped bias add: output and bias must have the same number of tensors"); - NVTE_CHECK(outputD->has_data(), "Grouped bias add: output is missing row-wise data"); - NVTE_CHECK(bias_tensor->has_data(), "Grouped bias add: bias is missing row-wise data"); - NVTE_CHECK(outputD->dtype() == bias_tensor->dtype(), - "Grouped bias add: output and bias must have matching dtypes"); - NVTE_CHECK(bias_tensor->all_same_first_dim(), - "Grouped bias add: bias must have uniform first dim (expected 1)"); - NVTE_CHECK(bias_tensor->get_common_first_dim() == 1, - "Grouped bias add: bias first dim must be 1"); - NVTE_CHECK(outputD->all_same_last_dim() && bias_tensor->all_same_last_dim(), - "Grouped bias add requires uniform last dim for output and bias"); - NVTE_CHECK(outputD->get_common_last_dim() == bias_tensor->get_common_last_dim(), - "Grouped bias add: output and bias last dims must match"); - constexpr int kVec = 4; - NVTE_CHECK(outputD->get_common_last_dim() % kVec == 0, - "Grouped bias add requires last dim divisible by ", kVec); + const char *api_name = use_scale ? "Grouped scaled bias add" : "Grouped bias add"; + + NVTE_CHECK(outputD->num_tensors >= 1, api_name, ": number of tensors must be at least 1"); + NVTE_CHECK(outputD->num_tensors == bias_tensor->num_tensors, api_name, + ": output and bias must have the same number of tensors"); + NVTE_CHECK(outputD->has_data(), api_name, ": output is missing row-wise data"); + NVTE_CHECK(bias_tensor->has_data(), api_name, ": bias is missing row-wise data"); + NVTE_CHECK(outputD->dtype() == bias_tensor->dtype(), api_name, + ": output and bias must have matching dtypes"); + NVTE_CHECK(bias_tensor->all_same_first_dim(), api_name, + ": bias must have uniform first dim (expected 1)"); + NVTE_CHECK(bias_tensor->get_common_first_dim() == 1, api_name, ": bias first dim must be 1"); + NVTE_CHECK(outputD->all_same_last_dim() && bias_tensor->all_same_last_dim(), api_name, + ": requires uniform last dim for output and bias"); + NVTE_CHECK(outputD->get_common_last_dim() == bias_tensor->get_common_last_dim(), api_name, + ": output and bias last dims must match"); const TensorShapeInfo d_meta = TensorShapeInfo::from_tensor(outputD); - const TensorShapeInfo bias_meta = TensorShapeInfo::from_tensor(bias_tensor); const DType dtype = outputD->dtype(); - constexpr int kThreads = 256; - const size_t total_elements = static_cast(outputD->logical_shape.data[0]) * - static_cast(outputD->logical_shape.data[1]); - const size_t total_vec_count = (total_elements + kVec - 1) / kVec; - int blocks_per_tensor = static_cast((total_vec_count + kThreads - 1) / kThreads); - const dim3 grid(outputD->num_tensors, blocks_per_tensor); + constexpr int kThreads = 128; + + const int num_tensors = static_cast(outputD->num_tensors); + NVTE_CHECK(num_tensors <= kMaxGroups, api_name, " supports at most ", kMaxGroups, + " tensors, got ", num_tensors); + const int total_rows = static_cast(outputD->logical_shape.data[0]); + const int n = static_cast(outputD->get_common_last_dim()); + + const size_t elem_size = typeToSize(dtype); + const int kVec = (elem_size <= 2) ? 8 : 4; + NVTE_CHECK(n % kVec == 0, api_name, ": requires last dim divisible by ", kVec); + + constexpr int kRowsPerBlock = 16; + constexpr int kBlocksPerSM = 32; + + const int num_sms = transformer_engine::cuda::sm_count(); + + const int block_cols = kThreads * kVec; + const int col_blocks = (n + block_cols - 1) / block_cols; + const int max_row_chunks = (total_rows + kRowsPerBlock - 1) / kRowsPerBlock; + const int row_blocks = std::min(max_row_chunks, num_sms * kBlocksPerSM / col_blocks); + const dim3 grid(std::max(1, row_blocks), col_blocks); const dim3 block(kThreads); - TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(dtype, T, { - grouped_bias_add_kernel<<>>( - static_cast(outputD->data.dptr), static_cast(bias_tensor->data.dptr), - d_meta, bias_meta, outputD->num_tensors); - }); + auto launch = [&](auto use_scale_tag) { + constexpr bool kUseScale = decltype(use_scale_tag)::value; + if (elem_size <= 2) { + constexpr int kV = 8; + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(dtype, T, { + grouped_bias_add_kernel + <<>>(static_cast(outputD->data.dptr), + static_cast(bias_tensor->data.dptr), + scale_ptr, d_meta, n, total_rows, num_tensors); + }); + } else { + constexpr int kV = 4; + TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(dtype, T, { + grouped_bias_add_kernel + <<>>(static_cast(outputD->data.dptr), + static_cast(bias_tensor->data.dptr), + scale_ptr, d_meta, n, total_rows, num_tensors); + }); + } + }; + + if (use_scale) { + launch(std::true_type{}); + } else { + launch(std::false_type{}); + } NVTE_CHECK_CUDA(cudaGetLastError()); } +} // namespace + +void nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, + cudaStream_t stream) { + NVTE_API_CALL(nvte_grouped_bias_add); + using namespace transformer_engine; + const GroupedTensor *outputD = convertNVTEGroupedTensorCheck(output); + const GroupedTensor *bias_tensor = convertNVTEGroupedTensorCheck(bias); + launch_grouped_bias_add(outputD, bias_tensor, nullptr, false, stream); +} + +void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, + const NVTETensor scale, cudaStream_t stream) { + NVTE_API_CALL(nvte_grouped_scaled_bias_add); + using namespace transformer_engine; + const GroupedTensor *outputD = convertNVTEGroupedTensorCheck(output); + const GroupedTensor *bias_tensor = convertNVTEGroupedTensorCheck(bias); + const Tensor *scale_tensor = convertNVTETensorCheck(scale); + + NVTE_CHECK(scale_tensor->data.dptr != nullptr, + "Grouped scaled bias add: scale tensor must not be null"); + NVTE_CHECK(scale_tensor->dtype() == DType::kFloat32, + "Grouped scaled bias add: scale must be float32"); + NVTE_CHECK(scale_tensor->data.shape.size() == 1, + "Grouped scaled bias add: scale must be 1D, got ", scale_tensor->data.shape.size(), + "D"); + const size_t total_rows = static_cast(outputD->logical_shape.data[0]); + NVTE_CHECK(scale_tensor->data.shape[0] == total_rows, "Grouped scaled bias add: scale size (", + scale_tensor->data.shape[0], ") must equal total rows (", total_rows, ")"); + + const float *scale_ptr = static_cast(scale_tensor->data.dptr); + launch_grouped_bias_add(outputD, bias_tensor, scale_ptr, true, stream); +} + #else // CUBLAS_VERSION < CUBLAS_GROUPED_GEMM_VERSION void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedTensor B, int transb, @@ -1397,6 +1542,14 @@ void nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTens CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); } +void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, + const NVTETensor scale, cudaStream_t stream) { + NVTE_ERROR( + "nvte_grouped_scaled_bias_add requires cuBLAS 13.3+, but compile-time cuBLAS version " + "is ", + CUBLAS_VERSION, ". Please upgrade to cuBLAS 13.3 (shipped with CUDA 13.2) or newer."); +} + size_t nvte_get_grouped_gemm_setup_workspace_size(size_t num_tensors) { NVTE_ERROR( "nvte_get_grouped_gemm_setup_workspace_size requires cuBLAS 13.3+, but compile-time cuBLAS " diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index fcd08a40a..bf9394c98 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -429,13 +429,23 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, NVTETensor workspace_setup, NVTETensor workspace_cublas, NVTEGroupedMatmulConfig config, cudaStream_t stream); -/*! \brief Grouped bias add for grouped GEMM outputs. +/*! \brief Grouped Bias add for grouped GEMM outputs. * +* output[row,col] += bias[col]. * Requires uniform last-dimension across all output tensors and bias tensors. */ void nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, cudaStream_t stream); +/*! \brief Grouped Scaled Bias add for grouped GEMM outputs. +* +* output[row,col] += bias[col] * scale[row], where biases are per-group +* and scales are per-token (per-row across all groups). +* Requires uniform last-dimension across all output tensors and bias tensors. +*/ +void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, + const NVTETensor scale, cudaStream_t stream); + #ifdef __cplusplus } // extern "C" #endif // __cplusplus diff --git a/transformer_engine/common/triton/grouped_dbias_dscales.py b/transformer_engine/common/triton/grouped_dbias_dscales.py index f5ddda259..592d5ace9 100644 --- a/transformer_engine/common/triton/grouped_dbias_dscales.py +++ b/transformer_engine/common/triton/grouped_dbias_dscales.py @@ -2,38 +2,49 @@ # # See LICENSE for license information. -"""Fused grouped dbias + dscales Triton kernel.""" +"""Fused grouped-dbias (+optional dscales) Triton kernel.""" import triton import triton.language as tl @triton.jit -def _grouped_dbias_dscales_kernel( +def _grouped_dbias_kernel( dy_ptr, + dbias_ptr, + offsets_ptr, scales_ptr, bias_ptr, - dbias_ptr, dscales_ptr, - offsets_ptr, hidden, + HAS_SCALES: tl.constexpr, N_ROW_SPLITS: tl.constexpr, BLOCK_M: tl.constexpr, BLOCK_H: tl.constexpr, ): - """Fused kernel: dbias[g] = sum_i(dy[i]*scales[i]), dscales[i] = dot(dy[i], bias[g]). + """Grouped dbias, optionally fused with dscales. + + For tokens i in group g(i), with s_i = scales[i] if HAS_SCALES else 1:: + + dbias[g, j] += sum_{i in g} dy[i, j] * s_i + + When HAS_SCALES is True, additionally:: + + dscales[i] += sum_j dy[i, j] * bias[g(i), j] Grid: (num_groups, N_ROW_SPLITS, cdiv(hidden, BLOCK_H)). - Each CTA computes the actual group size from device-side offsets, - divides row tiles evenly among N_ROW_SPLITS, and loops only over - its share. The loop bound is dynamic (no constexpr) so it adapts - to each group's size -- no wasted iterations, no host-device sync. + Each CTA computes its group's actual size from device-side offsets, + splits the row tiles evenly across N_ROW_SPLITS, and loops over only + its share -- dynamic loop bound, no host-device sync, no wasted iters. - - dbias: accumulated in registers, one atomic-add at the end + - dbias: register-accumulated, one atomic-add per CTA at the end (N_ROW_SPLITS contributors per group). - - dscales: atomic-add per iteration across column tiles + - dscales (if enabled): atomic-add per column-tile iteration (cdiv(hidden, BLOCK_H) contributors per element). + + When HAS_SCALES is False, ``scales_ptr``, ``bias_ptr`` and + ``dscales_ptr`` are unused and may be passed as dummy pointers. """ group_idx = tl.program_id(0) row_split = tl.program_id(1) @@ -46,41 +57,41 @@ def _grouped_dbias_dscales_kernel( total_tiles = (group_rows + BLOCK_M - 1) // BLOCK_M tiles_per_split = (total_tiles + N_ROW_SPLITS - 1) // N_ROW_SPLITS my_tile_start = row_split * tiles_per_split - col_offs = col_block * BLOCK_H + tl.arange(0, BLOCK_H) col_mask = col_offs < hidden - bias_vals = tl.load( - bias_ptr + group_idx * hidden + col_offs, - mask=col_mask, - other=0.0, - ).to(tl.float32) + if HAS_SCALES: + bias_vals = tl.load( + bias_ptr + group_idx * hidden + col_offs, + mask=col_mask, + other=0.0, + ).to(tl.float32) dbias_acc = tl.zeros([BLOCK_H], dtype=tl.float32) row_offs = tl.arange(0, BLOCK_M) - for local_tile in range(tiles_per_split): tile_idx = my_tile_start + local_tile global_rows = row_start + tile_idx * BLOCK_M + row_offs row_mask = global_rows < row_end tile_mask = row_mask[:, None] & col_mask[None, :] - dy_tile = tl.load( dy_ptr + global_rows[:, None] * hidden + col_offs[None, :], mask=tile_mask, other=0.0, ).to(tl.float32) - scales_vals = tl.load(scales_ptr + global_rows, mask=row_mask, other=0.0) - - dbias_acc += tl.sum(dy_tile * scales_vals[:, None], axis=0) - - dscales_partial = tl.sum(dy_tile * bias_vals[None, :], axis=1) - tl.atomic_add( - dscales_ptr + global_rows, - dscales_partial, - mask=row_mask, - ) + if HAS_SCALES: + scales_vals = tl.load(scales_ptr + global_rows, mask=row_mask, other=0.0) + dbias_acc += tl.sum(dy_tile * scales_vals[:, None], axis=0) + + dscales_partial = tl.sum(dy_tile * bias_vals[None, :], axis=1) + tl.atomic_add( + dscales_ptr + global_rows, + dscales_partial, + mask=row_mask, + ) + else: + dbias_acc += tl.sum(dy_tile, axis=0) tl.atomic_add( dbias_ptr + group_idx * hidden + col_offs, diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 82891ca83..6f3553bf9 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -327,6 +327,7 @@ def general_grouped_gemm_for_grouped_tensor( accumulate: bool = False, use_split_accumulator: bool = False, bias=None, + bias_scale: Optional[torch.Tensor] = None, grad: bool = False, alpha: Optional[torch.Tensor] = None, beta: Optional[torch.Tensor] = None, @@ -365,6 +366,9 @@ def general_grouped_gemm_for_grouped_tensor( "Apply bias manually after the GEMM." ) + if bias_scale is not None and bias is None: + raise ValueError("bias_scale requires bias to be provided.") + num_tensors = B.num_tensors rowwise = B.rowwise_data device = rowwise.device if rowwise is not None else B.columnwise_data.device @@ -401,6 +405,7 @@ def general_grouped_gemm_for_grouped_tensor( transb, out, bias, + bias_scale, alpha, beta, workspace_setup, diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 929be8906..4a2ea7412 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -160,11 +160,13 @@ std::optional> te_general_grouped_gemm( py::object te_general_grouped_gemm_for_grouped_tensor( py::handle A, bool transa, py::handle B, bool transb, py::handle D, py::object bias, - at::Tensor alpha, at::Tensor beta, at::Tensor workspace_setup, at::Tensor workspace_cublas, - bool use_split_accumulator, int math_sm_count); + std::optional bias_scale, at::Tensor alpha, at::Tensor beta, + at::Tensor workspace_setup, at::Tensor workspace_cublas, bool use_split_accumulator, + int math_sm_count); py::object te_general_grouped_gemm_for_discrete_in(py::handle A, bool transa, py::handle B, bool transb, py::handle D, py::object bias, + std::optional bias_scale, at::Tensor alpha, at::Tensor beta, at::Tensor workspace_setup, at::Tensor workspace_cublas, @@ -172,6 +174,7 @@ py::object te_general_grouped_gemm_for_discrete_in(py::handle A, bool transa, py py::object te_general_grouped_gemm_for_discrete_out(py::handle A, bool transa, py::handle B, bool transb, py::handle D, py::object bias, + std::optional bias_scale, at::Tensor alpha, at::Tensor beta, at::Tensor workspace_setup, at::Tensor workspace_cublas, diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index 08470962f..427eb7934 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -610,8 +610,9 @@ std::optional> te_general_grouped_gemm( py::object te_general_grouped_gemm_for_grouped_tensor( py::handle A, bool transa, py::handle B, bool transb, py::handle D, py::object bias, - at::Tensor alpha, at::Tensor beta, at::Tensor workspace_setup, at::Tensor workspace_cublas, - bool use_split_accumulator, int math_sm_count) { + std::optional bias_scale, at::Tensor alpha, at::Tensor beta, + at::Tensor workspace_setup, at::Tensor workspace_cublas, bool use_split_accumulator, + int math_sm_count) { using namespace transformer_engine::pytorch::detail; init_extension(); @@ -652,10 +653,18 @@ py::object te_general_grouped_gemm_for_grouped_tensor( if (!bias.is_none()) { auto grouped_bias = GroupedTensorFromPyTorchGroupedTensor(bias); - NVTE_SCOPED_GIL_RELEASE({ - nvte_grouped_bias_add(grouped_D.data(), grouped_bias.data(), - at::cuda::getCurrentCUDAStream()); - }); + if (bias_scale.has_value()) { + auto te_bias_scale = makeTransformerEngineTensor(*bias_scale); + NVTE_SCOPED_GIL_RELEASE({ + nvte_grouped_scaled_bias_add(grouped_D.data(), grouped_bias.data(), te_bias_scale.data(), + at::cuda::getCurrentCUDAStream()); + }); + } else { + NVTE_SCOPED_GIL_RELEASE({ + nvte_grouped_bias_add(grouped_D.data(), grouped_bias.data(), + at::cuda::getCurrentCUDAStream()); + }); + } } return py::reinterpret_borrow(D); @@ -663,6 +672,7 @@ py::object te_general_grouped_gemm_for_grouped_tensor( py::object te_general_grouped_gemm_for_discrete_in(py::handle A, bool transa, py::handle B, bool transb, py::handle D, py::object bias, + std::optional bias_scale, at::Tensor alpha, at::Tensor beta, at::Tensor workspace_setup, at::Tensor workspace_cublas, @@ -720,10 +730,18 @@ py::object te_general_grouped_gemm_for_discrete_in(py::handle A, bool transa, py if (!bias.is_none()) { auto grouped_bias = GroupedTensorFromPyTorchGroupedTensor(bias); - NVTE_SCOPED_GIL_RELEASE({ - nvte_grouped_bias_add(grouped_D.data(), grouped_bias.data(), - at::cuda::getCurrentCUDAStream()); - }); + if (bias_scale.has_value()) { + auto te_bias_scale = makeTransformerEngineTensor(*bias_scale); + NVTE_SCOPED_GIL_RELEASE({ + nvte_grouped_scaled_bias_add(grouped_D.data(), grouped_bias.data(), te_bias_scale.data(), + at::cuda::getCurrentCUDAStream()); + }); + } else { + NVTE_SCOPED_GIL_RELEASE({ + nvte_grouped_bias_add(grouped_D.data(), grouped_bias.data(), + at::cuda::getCurrentCUDAStream()); + }); + } } return py::reinterpret_borrow(D); @@ -731,6 +749,7 @@ py::object te_general_grouped_gemm_for_discrete_in(py::handle A, bool transa, py py::object te_general_grouped_gemm_for_discrete_out(py::handle A, bool transa, py::handle B, bool transb, py::handle D, py::object bias, + std::optional bias_scale, at::Tensor alpha, at::Tensor beta, at::Tensor workspace_setup, at::Tensor workspace_cublas, @@ -788,5 +807,4 @@ py::object te_general_grouped_gemm_for_discrete_out(py::handle A, bool transa, p return py::reinterpret_borrow(D); } - } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index a1d40a30e..fe5997a71 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -35,7 +35,10 @@ from .._common import is_quantized_tensor, maybe_dequantize from ..op import BasicOperation, OperationContext from ...tensor import GroupedTensor -from ...triton.grouped_dbias_dscales import _compute_grouped_dbias_dscales +from ...triton.grouped_dbias_dscales import ( + compute_grouped_dbias, + compute_grouped_dbias_dscales, +) class GroupedLinear(BasicOperation): @@ -890,27 +893,25 @@ def fuser_backward( columnwise=ctx.weight_requires_grad, ) dys = tex.split_quantize(dy, split_sizes_int, ctx.grad_output_quantizers) - if has_bias and not self._scale_bias: - dy_splits = list(torch.split(grad_output, split_sizes_int)) - grad_biases = [dy_s.reshape(-1, dy_s.size(-1)).sum(dim=0) for dy_s in dy_splits] else: dys = torch.split(dy, split_sizes_int) - if has_bias and not self._scale_bias: - grad_biases = [dy_s.reshape(-1, dy_s.size(-1)).sum(dim=0) for dy_s in dys] - if self._scale_bias and has_bias: - bias_packed = torch.stack(self._get_bias_tensors(ctx.dtype)) - scales_f32 = scales.to(dtype=torch.float32) + if has_bias: + dy_2d = dy.reshape(-1, dy.size(-1)) offsets = torch.zeros(num_groups + 1, dtype=torch.int64, device=device) offsets[1:] = split_sizes.cumsum(0) - dy_2d = dy.reshape(-1, dy.size(-1)) - dbias_packed, grad_scales = _compute_grouped_dbias_dscales( - dy_2d, - scales_f32, - bias_packed, - offsets=offsets, - ) - grad_biases = [dbias_packed[idx] for idx in range(num_groups)] + if self._scale_bias: + bias_packed = torch.stack(self._get_bias_tensors(ctx.dtype)) + scales_f32 = scales.to(dtype=torch.float32) + dbias_packed, grad_scales = compute_grouped_dbias_dscales( + dy_2d, + scales_f32, + bias_packed, + offsets=offsets, + ) + else: + dbias_packed = compute_grouped_dbias(dy_2d, offsets, num_groups) + grad_biases = [dbias_packed[idx].to(dtype=ctx.dtype) for idx in range(num_groups)] # Initialize grad weight buffers accumulate_into_main_grad = self._accumulate_into_main_grad diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index fc69b522d..aca49e986 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -32,7 +32,7 @@ ) from ...cpp_extensions import general_grouped_gemm_for_grouped_tensor from ...module.base import _2X_ACC_WGRAD -from ...triton.grouped_dbias_dscales import _compute_grouped_dbias_dscales +from ...triton.grouped_dbias_dscales import compute_grouped_dbias_dscales def _cudnn_compute_wgrad( @@ -593,7 +593,7 @@ def fuser_backward( if scale_bias: fc2_biases = fc2_op._get_bias_tensors(dtype) bias_packed = torch.stack(fc2_biases) - fc2_dbias_packed_result, grad_scales = _compute_grouped_dbias_dscales( + fc2_dbias_packed_result, grad_scales = compute_grouped_dbias_dscales( fc2_dy, scales_f32, bias_packed, diff --git a/transformer_engine/pytorch/triton/grouped_dbias_dscales.py b/transformer_engine/pytorch/triton/grouped_dbias_dscales.py index f87130b7c..836510007 100644 --- a/transformer_engine/pytorch/triton/grouped_dbias_dscales.py +++ b/transformer_engine/pytorch/triton/grouped_dbias_dscales.py @@ -2,19 +2,79 @@ # # See LICENSE for license information. -"""PyTorch wrapper for the fused grouped dbias + dscales Triton kernel.""" +"""PyTorch wrappers for the fused grouped-dbias (+optional dscales) Triton kernel.""" +import os from typing import Optional, Tuple import torch import triton -from transformer_engine.common.triton.grouped_dbias_dscales import ( - _grouped_dbias_dscales_kernel, -) +from transformer_engine.common.triton.grouped_dbias_dscales import _grouped_dbias_kernel -def _compute_grouped_dbias_dscales( +def _is_deterministic_mode() -> bool: + """Return True if TE is currently requesting deterministic execution.""" + return not bool(int(os.getenv("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1"))) + + +def _launch_grouped_dbias( + dy: torch.Tensor, + offsets: torch.Tensor, + dbias: torch.Tensor, + scales: Optional[torch.Tensor], + bias: Optional[torch.Tensor], + dscales: Optional[torch.Tensor], +) -> None: + """Launch the unified grouped-dbias kernel. + + If ``scales`` / ``bias`` / ``dscales`` are all None, runs the dbias-only + specialization; otherwise runs the fused dbias+dscales specialization + (all three must be provided together). + """ + if _is_deterministic_mode(): + raise RuntimeError( + "grouped_dbias Triton kernel uses non-deterministic atomic adds " + "and cannot be used when deterministic execution is requested " + "(NVTE_ALLOW_NONDETERMINISTIC_ALGO=0). " + "Disable determinism or use a deterministic fallback." + ) + + BLOCK_M = 128 + BLOCK_H = 128 + N_ROW_SPLITS = 4 + has_scales = scales is not None + assert ( + has_scales == (bias is not None) == (dscales is not None) + ), "_launch_grouped_dbias: scales, bias and dscales must be provided together" + + hidden = dy.shape[1] + num_groups = dbias.shape[0] + + # Triton requires real pointers; reuse dy as a harmless dummy when unused. + scales_arg = scales if has_scales else dy + bias_arg = bias if has_scales else dy + dscales_arg = dscales if has_scales else dy + + grid = (num_groups, N_ROW_SPLITS, triton.cdiv(hidden, BLOCK_H)) + _grouped_dbias_kernel[grid]( + dy, + dbias, + offsets, + scales_arg, + bias_arg, + dscales_arg, + hidden, + HAS_SCALES=has_scales, + N_ROW_SPLITS=N_ROW_SPLITS, + BLOCK_M=BLOCK_M, + BLOCK_H=BLOCK_H, + num_warps=4, + num_stages=2, + ) + + +def compute_grouped_dbias_dscales( dy: torch.Tensor, scales: torch.Tensor, bias: torch.Tensor, @@ -22,11 +82,11 @@ def _compute_grouped_dbias_dscales( dbias: Optional[torch.Tensor] = None, dscales: Optional[torch.Tensor] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - """Compute dbias and dscales via a single fused Triton kernel. + """Compute scaled grouped dbias and dscales via a single fused Triton kernel. - Computes the following, where token *i* belongs to group *g(i)*: + For tokens i in group g(i):: - dbias[g, j] += sum_{i in group g} dy[i, j] * scales[i] + dbias[g, j] += sum_{i in g} dy[i, j] * scales[i] dscales[i] += sum_j dy[i, j] * bias[g(i), j] Both outputs use fp32 atomic adds, so pre-populated tensors are @@ -38,12 +98,10 @@ def _compute_grouped_dbias_dscales( bias: (num_groups, hidden) -- per-group FC2 biases. offsets: (num_groups+1,) int64 -- cumulative row offsets ``[0, s0, s0+s1, ..., total_tokens]``. - dbias: optional (num_groups, hidden) float32 -- if provided, - the kernel accumulates into this tensor; otherwise a - zero tensor is allocated. - dscales: optional (total_tokens,) float32 -- if provided, - the kernel accumulates into this tensor; otherwise a - zero tensor is allocated. + dbias: optional (num_groups, hidden) float32 -- accumulated into + if provided, otherwise a fresh zero tensor is allocated. + dscales: optional (total_tokens,) float32 -- accumulated into + if provided, otherwise a fresh zero tensor is allocated. Returns: dbias: (num_groups, hidden) float32 @@ -58,37 +116,50 @@ def _compute_grouped_dbias_dscales( else: assert ( dbias.dtype == torch.float32 - ), f"_compute_grouped_dbias_dscales: dbias must be float32, got {dbias.dtype}" + ), f"compute_grouped_dbias_dscales: dbias must be float32, got {dbias.dtype}" if dscales is None: dscales = torch.zeros(total_tokens, dtype=torch.float32, device=dy.device) else: assert ( dscales.dtype == torch.float32 - ), f"_compute_grouped_dbias_dscales: dscales must be float32, got {dscales.dtype}" + ), f"compute_grouped_dbias_dscales: dscales must be float32, got {dscales.dtype}" - BLOCK_M = 128 - BLOCK_H = 128 - N_ROW_SPLITS = 4 + _launch_grouped_dbias(dy, offsets, dbias, scales, bias, dscales) + return dbias, dscales - grid = ( - num_groups, - N_ROW_SPLITS, - triton.cdiv(hidden, BLOCK_H), - ) - _grouped_dbias_dscales_kernel[grid]( - dy, - scales, - bias, - dbias, - dscales, - offsets, - hidden, - N_ROW_SPLITS=N_ROW_SPLITS, - BLOCK_M=BLOCK_M, - BLOCK_H=BLOCK_H, - num_warps=4, - num_stages=2, - ) +def compute_grouped_dbias( + dy: torch.Tensor, + offsets: torch.Tensor, + num_groups: int, + dbias: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Compute grouped dbias = per-group sum of dy via the fused Triton kernel. - return dbias, dscales + For tokens i in group g:: + + dbias[g, j] += sum_{i in g} dy[i, j] + + Args: + dy: (total_tokens, hidden) -- output grad. + offsets: (num_groups+1,) int64 -- cumulative row offsets + ``[0, s0, s0+s1, ..., total_tokens]``. + num_groups: number of groups (``offsets`` has ``num_groups + 1`` + entries). + dbias: optional (num_groups, hidden) float32 -- accumulated into + if provided, otherwise a fresh zero tensor is allocated. + + Returns: + dbias: (num_groups, hidden) float32 + """ + hidden = dy.shape[1] + + if dbias is None: + dbias = torch.zeros(num_groups, hidden, dtype=torch.float32, device=dy.device) + else: + assert ( + dbias.dtype == torch.float32 + ), f"compute_grouped_dbias: dbias must be float32, got {dbias.dtype}" + + _launch_grouped_dbias(dy, offsets, dbias, scales=None, bias=None, dscales=None) + return dbias From f2ed86bbf9eeddeecc8bc34a61cbff142d3edf7b Mon Sep 17 00:00:00 2001 From: Chase Block Date: Wed, 22 Apr 2026 00:18:54 -0500 Subject: [PATCH 007/180] Add better ordering enforcment to split_overlap_rs gemms. (#2056) * Add better ordering enforcment to split_overlap_rs gemms. This adds a short delay kernel to the split_overlap_rs function, which ensures that the gemms are properly ordered when run with cuda graphs. Signed-off-by: Chase Block * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Chase Block Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Przemyslaw Tredak --- .../common/comm_gemm_overlap/comm_gemm_overlap.cpp | 11 +++++++++++ .../comm_gemm_overlap/userbuffers/userbuffers.cu | 8 ++++++++ .../comm_gemm_overlap/userbuffers/userbuffers.h | 1 + 3 files changed, 20 insertions(+) diff --git a/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp b/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp index aad2ec068..133f1a09e 100644 --- a/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp +++ b/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp @@ -1157,6 +1157,10 @@ void CommOverlapP2PBase::split_overlap_rs(const TensorWrapper &A, bool transa, NVTE_CHECK_CUDA(cudaStreamWaitEvent(_stream_compute[i], _start_compute, 0)); } + // Launch the tiny delay kernel + userbuffers_tiny_delay(_stream_send[0]); + NVTE_CHECK_CUDA(cudaEventRecord(_start_compute, _stream_send[0])); + // GEMM and send/recv chunks for (int i = 0; i < _tp_size; i++) { // GEMM chunk @@ -1169,6 +1173,13 @@ void CommOverlapP2PBase::split_overlap_rs(const TensorWrapper &A, bool transa, auto workspace_chunk = get_tensor_chunk(workspace, stream_id * workspace_size_chunk, {workspace_size_chunk}); + if (i == 1) { + NVTE_CHECK_CUDA(cudaStreamWaitEvent(_stream_compute[stream_id], _start_compute)); + } else if (i > 1) { + NVTE_CHECK_CUDA( + cudaEventRecord(_start_compute, _stream_compute[(i - 2) % _stream_compute.size()])); + NVTE_CHECK_CUDA(cudaStreamWaitEvent(_stream_compute[stream_id], _start_compute)); + } nvte_cublas_gemm(A.data(), input_b_chunk.data(), output_chunk.data(), bias.data(), pre_gelu_out.data(), transa, transb, grad, workspace_chunk.data(), accumulate, use_split_accumulator, _math_sms, _stream_compute[stream_id]); diff --git a/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.cu b/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.cu index 3d8848d95..4bbbfc3c1 100644 --- a/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.cu +++ b/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.cu @@ -2304,6 +2304,14 @@ __global__ void __launch_bounds__(MAX_THREADS) kuserbuffers_pushsendrecv_multiat (index) * NVTE_MAX_NVLINK * NVTE_MAX_REGIONS) * \ sizeof(int))) +static __global__ void tiny_delay_kern() { + // Although a loop could be used to add a larger delay here, for + // the purpose of enforcing proper kernel ordering when using CG, + // an empty kernel seems to work well enough. +} + +void userbuffers_tiny_delay(cudaStream_t stream) { tiny_delay_kern<<<1, 1, 0, stream>>>(); } + void userbuffers_send(const int srchandler, const size_t srcoffset, const int dsthandler, const size_t dstoffset, const size_t bytes, communicator *comm, const int peer, cudaStream_t stream) { diff --git a/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.h b/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.h index c8d7c8731..52be0af53 100644 --- a/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.h +++ b/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers.h @@ -266,6 +266,7 @@ output is strided: row starts separated by stride elements*/ // push model: data arrived and visible at receiver(barrier enforced) // pull model: data ready to be pulled by receiver(no barrier needed) +void userbuffers_tiny_delay(cudaStream_t stream); void userbuffers_send(const int srchandler, const size_t srcoffset, const int dsthandler, const size_t dstoffset, const size_t bytes, communicator *comm, const int peer, cudaStream_t stream = 0); From 4014f7f4a477ed93e8afce0af78fc070a0991333 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Buschk=C3=A4mper?= Date: Wed, 22 Apr 2026 07:19:27 +0200 Subject: [PATCH 008/180] Fix flash attention version check. (#2910) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Björn Buschkämper --- .../attention/dot_product_attention/backends.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 60a6f655b..4104820a1 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -981,14 +981,14 @@ def forward( batch_size * context_len, ) - use_flash_attn_4 = False - if flash_attention_backend is not None and flash_attention_backend > PkgVersion("4.0.0b"): - use_flash_attn_4 = True - use_flash_attn_3 = False - if flash_attention_backend is not None and PkgVersion( - "3.0.0b" - ) < flash_attention_backend < PkgVersion("4.0.0"): - use_flash_attn_3 = True + # FA4 prereleases such as 4.0.0b8 sort below 4.0.0, so key off the major + # version instead of a stable-version range check when selecting the API. + use_flash_attn_4 = ( + flash_attention_backend is not None and flash_attention_backend.major == 4 + ) + use_flash_attn_3 = ( + flash_attention_backend is not None and flash_attention_backend.major == 3 + ) if context_parallel and all( not isinstance(x, Float8Tensor) for x in [query_layer, key_layer, value_layer] ): From 0a088c1d5059a2e0df44f9bd9df19d071bc5712a Mon Sep 17 00:00:00 2001 From: "Peter St. John" Date: Wed, 22 Apr 2026 09:48:06 -0600 Subject: [PATCH 009/180] [PyT] Fix FSDP2 memory leaks for FP8 weight workspaces and transpose caches (#2805) * fixing mem leaks Signed-off-by: Peter St. John * update xfail message Signed-off-by: Peter St. John * addressing greptile comments Signed-off-by: Peter St. John * remove fsdp_safe Signed-off-by: Peter St. John * fix Float8BlockScaling backward override, unused imports, xfail MXFP8 HSDP Fix three issues: 1. LayerNormLinear weight quantizer was requesting columnwise usage even when backward_override is set. The FSDP2 refactor on this branch changed the columnwise condition to `is_grad_enabled and not is_fsdp2` but omitted the `backward_override is None` guard. This caused the weight to carry unnecessary columnwise data, which the dequantized backward path rejects. Also harden Float8BlockQuantizer::create_tensor to pass py::none() based on usage flags rather than relying on at::Tensor::defined() which is unreliable for default-constructed tensors in some PyTorch builds. 2. Remove unused imports: Float8Tensor from linear.py, and the entire float8_tensor import line from layernorm_linear.py (Float8Quantizer, Float8CurrentScalingQuantizer, Float8Tensor were all unused). 3. xfail MXFP8BlockScaling + fp8_init + HSDP in FSDP2 model tests. Pre-existing bug (confirmed on main) where fsdp_post_all_gather receives fewer output tensors than fsdp_pre_all_gather sent when the HSDP shard dimension is trivial (size 1). Signed-off-by: Peter St. John * address review comments Signed-off-by: Peter St. John * address greptile review Signed-off-by: Peter St. John * no need Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * revert grouped linear changes Signed-off-by: Varun Thumbe * fix all tests Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix cudnn commit Signed-off-by: Varun Thumbe * fix lint Signed-off-by: Varun Thumbe * revert Signed-off-by: Varun Thumbe * fix lint Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Peter St. John Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 Co-authored-by: Varun Thumbe Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../fsdp2_tests/run_fsdp2_fused_adam.py | 28 +++--- .../fsdp2_tests/run_fsdp2_mem_leak.py | 24 +---- .../fsdp2_tests/run_fsdp2_model.py | 10 ++ transformer_engine/pytorch/module/base.py | 13 +++ .../pytorch/module/layernorm_linear.py | 44 ++++++++- .../pytorch/module/layernorm_mlp.py | 91 +++++++++++++++++-- transformer_engine/pytorch/module/linear.py | 40 +++++++- transformer_engine/pytorch/tensor/utils.py | 15 +++ 8 files changed, 214 insertions(+), 51 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index ac38bc4aa..83c3f5b56 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -70,16 +70,14 @@ def _build_model( """Build a Sequential of TransformerLayers, optionally with FP8 init. When fp8_init=True and use_meta_device=True (the default), the model is - created on the meta device to avoid FSDP2 incompatibility with - QuantizedTensor wrapper subclasses (e.g. MXFP8Tensor) whose storage is - inaccessible via data_ptr(). Parameters are materialized after FSDP2 - sharding via reset_parameters() in _shard_model(). + created on the meta device so parameters are materialized after FSDP2 + sharding via reset_parameters() in _shard_model(). This ensures the + sharded parameter format is compatible with the FSDP2 all-gather hooks. When use_meta_device=False, the model is created directly on CUDA. - This is the legacy path that does NOT work for block-scaling quantized - tensors (MXFP8, Float8Blockwise, NVFP4) because FSDP2's - reset_sharded_param() crashes on wrapper subclass tensors with - data_ptr() == 0. + This only works for per-tensor FP8 (DelayedScaling, Float8CurrentScaling). + Block-scaling types (MXFP8, Float8Blockwise, NVFP4) fail because their + FSDP2 all-gather hooks do not support CUDA-initialized parameters. """ if fp8_init: ctx = te.quantized_model_init( @@ -220,18 +218,20 @@ def test_fused_adam_fp8_master_weights_no_meta(recipe_name): """FusedAdam with master_weights + FSDP2 + quantized_model_init WITHOUT meta device. This is the legacy path that creates quantized params directly on CUDA. - FSDP2's reset_sharded_param() crashes on block-scaling QuantizedTensor - wrapper subclasses (data_ptr() == 0). This test documents that failure. + FSDP2's forward-time all-gather hooks for block-scaling QuantizedTensor + subclasses fail when parameters are initialized directly on CUDA rather + than on the meta device. NVFP4Tensor does not implement the FSDP all-gather + hooks at all. - For per-tensor FP8 (DelayedScaling, Float8CurrentScaling) this works - because Float8Tensor's storage is accessible via data_ptr(). + For per-tensor FP8 (DelayedScaling, Float8CurrentScaling) the all-gather + hooks handle CUDA-initialized Float8Tensor parameters correctly. """ recipe = get_recipe_from_string(recipe_name) if recipe_name in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): pytest.xfail( - f"{recipe_name}: FSDP2 without meta-device init crashes on block-scaling " - "QuantizedTensor wrapper subclasses (data_ptr() == 0). " + f"{recipe_name}: FSDP2 all-gather hooks for block-scaling QuantizedTensor " + "subclasses fail when parameters are initialized on CUDA. " "Use device='meta' + reset_parameters() after sharding." ) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py index 387d3a964..b5436e270 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_mem_leak.py @@ -253,14 +253,6 @@ def test_bf16_no_excess_forward_memory(): ) -@pytest.mark.xfail( - strict=False, - reason=( - "Issue #2681: Quantized weights created during forward pass are not " - "deallocated between layers. Each layer's FP8 copies accumulate, " - "adding per-layer memory overhead beyond what bf16 autograd saves require." - ), -) def test_fp8_temp_accumulation_across_layers(recipe_name, quantized_model_init): """Detect FP8 weight temporaries accumulating across layers during forward. @@ -381,15 +373,6 @@ def test_bf16_no_excess_backward_memory(): ) -@pytest.mark.xfail( - strict=False, - reason=( - "Issue #2717: _create_transpose tensor allocated in " - "float8_tensor_storage.py persists after backward pass until the next " - "forward pass frees it. These tensors should be released when backward " - "completes, not retained across step boundaries." - ), -) def test_transpose_cache_retained_after_backward(recipe_name, quantized_model_init): """Detect transpose caches persisting after backward completes. @@ -456,9 +439,10 @@ def test_transpose_cache_retained_after_backward(recipe_name, quantized_model_in # significantly more positive than bf16. excess = fp8_bwd_delta - bf16_bwd_delta - # Allow 256 KiB total for FP8 scale/amax bookkeeping. - # Transpose caches (~3 MiB for this 8-layer model) should NOT persist. - tolerance = 256 * 1024 + # Allow 1 MiB for FP8 scale/amax bookkeeping and temporary workspace + # re-creation during backward. The key check is that transpose caches + # (~3 MiB for this 8-layer model) do NOT persist across steps. + tolerance = 1024 * 1024 assert excess <= tolerance, ( f"FP8 backward retains {excess/1024**2:.2f} MiB more than bf16 baseline. " diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 5a8c903c7..6342e63e7 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -370,6 +370,16 @@ def _train(args): @pytest.mark.parametrize("fp8_init", [False, True]) @pytest.mark.parametrize("layer_type", ["LayerNormLinear", "TransformerLayer"]) def test_distributed(recipe_name, fp8_init, sharding_dims, layer_type): + if recipe_name == "MXFP8BlockScaling" and fp8_init and len(sharding_dims) == 2: + pytest.xfail( + "MXFP8BlockScaling + fp8_init + HSDP: fsdp_post_all_gather receives fewer " + "all_gather_outputs than the number of tensors sent by fsdp_pre_all_gather " + "when the HSDP shard dimension is trivial (size 1). MXFP8 sends 2 tensors " + "(data + scale_inv, both uint8) but gets back 1. Float8Tensor avoids this by " + "sending only 1 tensor (scale is per-tensor metadata). Fix: concatenate MXFP8 " + "data and scale_inv into a single buffer in pre_all_gather, split in post." + ) + if recipe_name == "Float8BlockScaling" and fp8_init: pytest.xfail( "Float8BlockScaling + fp8_init: scale inverse padding is not handled " diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 83781ca3f..ebfb98b2d 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -809,6 +809,19 @@ def module_setattr(self, name: str, value: Any) -> None: """ super().__setattr__(name, value) + @property + def is_fsdp2(self) -> bool: + """Whether this module is wrapped with FSDP2.""" + if not hasattr(self, "_is_fsdp2"): + try: + from ..distributed import _get_module_fsdp_state + + _get_module_fsdp_state(self) + self._is_fsdp2 = True + except (RuntimeError, ImportError): + self._is_fsdp2 = False + return self._is_fsdp2 + def adjust_amax_history_length(self, length: int, fwd: Optional[bool] = None) -> None: """ Delayed scaling only. diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index f26faade0..d69e643c4 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -17,7 +17,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version -from transformer_engine.pytorch.tensor.utils import is_custom +from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom from .base import ( fill_userbuffers_buffer_for_all_gather, get_ub, @@ -142,6 +142,7 @@ def forward( skip_fp8_weight_update, symmetric_ar_type, debug, + is_fsdp2, ) = non_tensor_args if fp8: backward_override = FP8GlobalStateManager.get_fp8_recipe().backward_override @@ -303,14 +304,17 @@ def forward( is_weight_param_quantized = isinstance(weight, QuantizedTensorStorage) # Configure quantizer - # If weight is already quantized, no need to set quantizer states + # If weight is already quantized, weight._quantizer is its true quantizer. # for debug mode we create quantizer every iteration, thus we need to set the quantizer states if is_weight_param_quantized and not debug: weight_quantizer = weight._quantizer elif weight_quantizer is not None: + # FSDP2: Skip columnwise/transpose creation during forward + # to avoid accumulating caches across layers. Backward's + # FSDP2 all-gather will recreate them. (Issue #2681) weight_quantizer.set_usage( rowwise=True, - columnwise=is_grad_enabled and backward_override is None, + columnwise=is_grad_enabled and not is_fsdp2 and backward_override is None, ) # Get quantized weight @@ -325,6 +329,7 @@ def forward( workspace_dtype=activation_dtype, cache=cache_weight, ) + weightmat.update_usage(rowwise_usage=True) else: @@ -471,9 +476,15 @@ def forward( ln_bias, ) + # FSDP2: Don't save FP8 workspace for non-quantized weights. + # Backward will re-quantize from FSDP2 all-gathered weight. + # (Issue #2681) + wt_save = weightmat + if is_fsdp2 and weightmat is not weight: + wt_save = None tensors_to_save, tensor_objects = prepare_for_saving( inputmat, - weightmat, + wt_save, weight, bias, ln_weight, @@ -486,6 +497,7 @@ def forward( ctx.requires_dgrad = inp_requires_grad ctx.requires_wgrad = weight.requires_grad ctx.is_weight_param_quantized = is_weight_param_quantized + ctx.is_fsdp2 = is_fsdp2 if fuse_wgrad_accumulation and weight.requires_grad: # Keep weakref to weight to preserve attributes like main_grad # when we need to modify the weight python object @@ -740,6 +752,19 @@ def backward( # Note: Gradient w.r.t. GEMM input (i.e. norm output). # -------------------------------------------------- + # FSDP2: Re-create workspace from all-gathered weight when + # workspace was not saved. (Issue #2681) + # Use saved_weight (the original weight parameter) since + # origin_weight is only set when fuse_wgrad_accumulation=True. + if weight is None: + if isinstance(saved_weight, QuantizedTensorStorage): + # saved weight is already set to right usages by + # fsdp2 quantized-tensor hooks when workspace was not saved. + weight = saved_weight + elif ctx.weight_quantizer is not None: + ctx.weight_quantizer.set_usage(rowwise=True, columnwise=True) + weight = ctx.weight_quantizer(saved_weight) + # Make sure required data is available if isinstance(grad_output, QuantizedTensorStorage): grad_output.update_usage(rowwise_usage=True) @@ -800,6 +825,14 @@ def backward( ) nvtx_range_pop(f"{nvtx_label}.dgrad_gemm") + # FSDP2 only handles deallocation all-gathered weights that it allocates. + # Columnwise data is derived from rowwise data after allgather for fp8 + # and 2d block-scaled weights in TE managed memory. So we need to clear + # it here. + # (Issues #2681, #2717) + if getattr(ctx, "is_fsdp2", False) and isinstance(weight, QuantizedTensorStorage): + clear_columnwise_cache(weight) + # Prepare grad input tensor # Note: Perform tensor-parallel communication dgrad = None @@ -1601,7 +1634,7 @@ def forward( else: fwd_fn = _LayerNormLinear.forward autograd_ctx = [None] - cache_name = None if is_first_microbatch is None else "weight" + cache_name = None if (is_first_microbatch is None or self.is_fsdp2) else "weight" weight_workspace = ( self._fp8_workspaces.get(cache_name) if cache_name is not None else None ) @@ -1645,6 +1678,7 @@ def forward( skip_fp8_weight_update, self.symmetric_ar_type, debug, + self.is_fsdp2, ) out, ln_out, new_weight_workspace = fwd_fn( *autograd_ctx, diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index a8d6e2e60..4fa7eb285 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -18,7 +18,7 @@ from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.torch_version import torch_version -from transformer_engine.pytorch.tensor.utils import is_custom +from transformer_engine.pytorch.tensor.utils import clear_columnwise_cache, is_custom from .base import ( fill_userbuffers_buffer_for_all_gather, _ub_communicators, @@ -237,6 +237,7 @@ def _forward( symmetric_ar_type, checkpoint, debug, + is_fsdp2, recompute_for_bwd, ) = non_tensor_args if fp8: @@ -338,6 +339,7 @@ def _forward( "symmetric_ar_type": symmetric_ar_type, "checkpoint": checkpoint, "debug": debug, + "is_fsdp2": is_fsdp2, "recompute_for_bwd": True, # set this to true for recomputation phase } # Make sure input dimensions are compatible @@ -480,19 +482,29 @@ def _forward( new_fc2_weight_workspace = None fc1_weight_final = fc1_weight fc2_weight_final = fc2_weight + # FSDP2: Skip columnwise/transpose creation during forward (not + # recompute) to avoid accumulating FP8 caches across layers. + # Backward's FSDP2 all-gather will recreate them. (Issue #2681) + fsdp2_skip_columnwise = is_fsdp2 and not is_recomputation if fp8 or debug: update_ws = is_first_microbatch is None or is_first_microbatch - # No need to set the quantizer states if weights are already quantized + # If weight is already quantized, weight._quantizer is its true quantizer. # for debug mode we create quantizer every iteration, thus we need to set the quantizer states if isinstance(fc1_weight, QuantizedTensorStorage) and not debug: fc1_weight_quantizer = fc1_weight._quantizer elif fc1_weight_quantizer is not None: - fc1_weight_quantizer.set_usage(rowwise=True, columnwise=is_grad_enabled) + fc1_weight_quantizer.set_usage( + rowwise=True, + columnwise=is_grad_enabled and not fsdp2_skip_columnwise, + ) if isinstance(fc2_weight, QuantizedTensorStorage) and not debug: fc2_weight_quantizer = fc2_weight._quantizer elif fc2_weight_quantizer is not None: - fc2_weight_quantizer.set_usage(rowwise=True, columnwise=is_grad_enabled) + fc2_weight_quantizer.set_usage( + rowwise=True, + columnwise=is_grad_enabled and not fsdp2_skip_columnwise, + ) fc1_weight_final, new_fc1_weight_workspace = quantize_weight( tensor=fc1_weight, @@ -757,16 +769,28 @@ def _forward( fc2_weight, fc2_bias, ) + # FSDP2: Don't save FP8 workspace copies for non-quantized + # weights. Backward will re-quantize from the FSDP2 + # all-gathered weight parameter. (Issue #2681) + fc1_wt_save = fc1_weight_final + fc2_wt_save = fc2_weight_final + if fsdp2_skip_columnwise: + if fc1_weight_final is not fc1_weight: + fc1_wt_save = None + if fc2_weight_final is not fc2_weight: + fc2_wt_save = None tensors_to_save, tensor_objects = prepare_for_saving( inputmat, ln_weight, ln_out, - fc1_weight_final, + fc1_wt_save, + fc1_weight, fc1_bias, fc1_out, fc1_out_without_bias, act_out, - fc2_weight_final, + fc2_wt_save, + fc2_weight, fc2_bias, mu, rsigma, @@ -818,6 +842,13 @@ def _forward( ctx.fc1_weight_requires_grad = fc1_weight.requires_grad ctx.fc2_weight_requires_grad = fc2_weight.requires_grad + ctx.fc1_weight = fc1_weight + ctx.fc2_weight = fc2_weight + ctx.fsdp2_skip_columnwise = fsdp2_skip_columnwise + # Store raw is_fsdp2 flag for backward cleanup — must not be + # gated on is_recomputation since backward cleanup runs after + # the real backward, not the recomputation forward. + ctx.is_fsdp2 = is_fsdp2 ctx.device = device ctx.activation_dtype = activation_dtype @@ -870,11 +901,13 @@ def _forward( ln_weight, ln_out, fc1_weight_final, + fc1_weight, fc1_bias, fc1_out, fc1_out_without_bias, act_out, fc2_weight_final, + fc2_weight, fc2_bias, mu, rsigma, @@ -1006,11 +1039,13 @@ def backward( ln_weight, ln_out, fc1_weight, + origin_fc1_weight, fc1_bias, fc1_out, fc1_out_without_bias, act_out, fc2_weight, + origin_fc2_weight, fc2_bias, mu, rsigma, @@ -1163,6 +1198,16 @@ def backward( and (not ctx.debug) ) + # FSDP2: Re-create workspace from all-gathered weight when + # workspace was not saved to avoid forward memory + # accumulation. (Issue #2681) + if fc2_weight is None: + if isinstance(origin_fc2_weight, QuantizedTensorStorage): + fc2_weight = origin_fc2_weight + elif ctx.fc2_weight_quantizer is not None: + ctx.fc2_weight_quantizer.set_usage(rowwise=True, columnwise=True) + fc2_weight = ctx.fc2_weight_quantizer(origin_fc2_weight) + # Make sure required data is available if isinstance(grad_output, QuantizedTensorStorage): grad_output.update_usage(rowwise_usage=True) @@ -1190,6 +1235,14 @@ def backward( ub_type=tex.CommOverlapType.AG if ctx.ub_overlap_ag else None, ) + # FSDP2: Clear columnwise/transpose caches after FC2 dgrad GEMM + # to prevent them from persisting on the all-gathered buffer. + # Uses is_fsdp2 (not fsdp2_skip_columnwise) so cleanup runs + # even when backward follows gradient-checkpoint recomputation. + # (Issues #2681, #2717) + if getattr(ctx, "is_fsdp2", False) and isinstance(fc2_weight, QuantizedTensorStorage): + clear_columnwise_cache(fc2_weight) + # Prepare input grad tensor dact = None fc2_dgrad = None @@ -1417,6 +1470,15 @@ def fc2_wgrad_gemm( # FC1 DGRAD # -------------------------------------------------- + # FSDP2: Re-create workspace from all-gathered weight when + # workspace was not saved. (Issue #2681) + if fc1_weight is None: + if isinstance(origin_fc1_weight, QuantizedTensorStorage): + fc1_weight = origin_fc1_weight + elif ctx.fc1_weight_quantizer is not None: + ctx.fc1_weight_quantizer.set_usage(rowwise=True, columnwise=True) + fc1_weight = ctx.fc1_weight_quantizer(origin_fc1_weight) + # Make sure required data is available if ctx.fc1_weight_quantizer is not None and isinstance( fc1_weight, QuantizedTensorStorage @@ -1449,6 +1511,14 @@ def fc2_wgrad_gemm( bulk_overlap=ctx.ub_bulk_dgrad, ) + # FSDP2: Clear columnwise/transpose caches after FC1 dgrad GEMM + # to prevent them from persisting on the all-gathered buffer. + # Uses is_fsdp2 (not fsdp2_skip_columnwise) so cleanup runs + # even when backward follows gradient-checkpoint recomputation. + # (Issues #2681, #2717) + if getattr(ctx, "is_fsdp2", False) and isinstance(fc1_weight, QuantizedTensorStorage): + clear_columnwise_cache(fc1_weight) + # Prepare grad input tensor # Note: Perform tensor-parallel communication fc1_dgrad = None @@ -2164,8 +2234,12 @@ def forward( fwd_fn = _LayerNormMLP.forward autograd_ctx = [None] - cache_name_fc1 = None if is_first_microbatch is None else "fc1_weight" - cache_name_fc2 = None if is_first_microbatch is None else "fc2_weight" + cache_name_fc1 = ( + None if (is_first_microbatch is None or self.is_fsdp2) else "fc1_weight" + ) + cache_name_fc2 = ( + None if (is_first_microbatch is None or self.is_fsdp2) else "fc2_weight" + ) fc1_weight_workspace = ( self._fp8_workspaces.get(cache_name_fc1) if cache_name_fc1 is not None else None ) @@ -2222,6 +2296,7 @@ def forward( self.symmetric_ar_type, self.checkpoint, debug, + self.is_fsdp2, ) out, ln_out, new_fc1_ws, new_fc2_ws = fwd_fn( *autograd_ctx, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 12339e777..7498760af 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -67,7 +67,7 @@ ) from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer -from ..tensor.utils import is_custom +from ..tensor.utils import clear_columnwise_cache, is_custom from ..export import is_in_onnx_export_mode, assert_warmed_up from ..cpu_offload import ( is_cpu_offload_enabled, @@ -137,6 +137,7 @@ def _linear_forward_impl( backward_override, custom, backward_input_needs_gather, + is_fsdp2, ) = non_tensor_args if backward_override == "high_precision": save_original_input = True @@ -263,7 +264,7 @@ def _linear_forward_impl( # No need to set the quantizer states if weight is already quantized # for debug mode we create quantizer every iteration, thus we need to set the quantizer states if weight_quantizer is not None and (not isinstance(weight, QuantizedTensor) or debug): - columnwise_usage = is_grad_enabled and inp.requires_grad + columnwise_usage = is_grad_enabled and inp.requires_grad and not is_fsdp2 if backward_override is not None: columnwise_usage = False if not columnwise_usage: @@ -427,9 +428,15 @@ def _linear_forward_impl( mark_not_offload(weight, weightmat, bias) # TODO(ksivamani): Check memory usage + # FSDP2: Don't save FP8 workspace for non-quantized weights. + # Backward will re-quantize from the FSDP2 all-gathered weight. + # (Issue #2681) + wt_save = weightmat + if is_fsdp2 and weightmat is not weight: + wt_save = None tensors_to_save, tensor_objects = prepare_for_saving( saved_inputmat, - weightmat, + wt_save, weight, bias, ) @@ -440,6 +447,7 @@ def _linear_forward_impl( "weight_quantizer": weight_quantizer, "fsdp_shapes": fsdp_shapes, "owns_input": owns_input, + "is_fsdp2": is_fsdp2, } return out, new_weight_workspace, tensors_to_save, tensor_objects, ctx_attrs @@ -494,6 +502,7 @@ def _linear_setup_ctx( backward_override, custom, backward_input_needs_gather, + _is_fsdp2, ) = non_tensor_args # Values derived from input tensors @@ -549,6 +558,7 @@ def _linear_setup_ctx( ctx.weight_quantizer = ctx_attrs["weight_quantizer"] ctx.fsdp_shapes = ctx_attrs["fsdp_shapes"] ctx.owns_input = ctx_attrs["owns_input"] + ctx.is_fsdp2 = ctx_attrs["is_fsdp2"] # backward overrides if backward_override is not None: @@ -765,6 +775,19 @@ def _linear_backward( dgrad_work = None if ctx.requires_dgrad: + # FSDP2: Re-create workspace from all-gathered weight when + # workspace was not saved. (Issue #2681) + # Use saved_weight (the original weight parameter) since + # weight_fp8 is only set when workspace was saved. + if weight_fp8 is None: + if isinstance(saved_weight, QuantizedTensorStorage): + # saved weight is already set to right usages by + # fsdp2 quantized-tensor hooks when workspace was not saved. + weight_fp8 = saved_weight + elif ctx.weight_quantizer is not None: + ctx.weight_quantizer.set_usage(rowwise=True, columnwise=True) + weight_fp8 = ctx.weight_quantizer(saved_weight) + # Make sure required data is available if isinstance(grad_output, QuantizedTensorStorage): grad_output.update_usage(rowwise_usage=True) @@ -826,6 +849,14 @@ def _linear_backward( ) nvtx_range_pop(f"{nvtx_label}.dgrad_gemm") + # FSDP2 only handles deallocation all-gathered weights that it allocates. + # Columnwise data is derived from rowwise data after allgather for fp8 + # and 2d block-scaled weights in TE managed memory. So we need to clear + # it here. + # (Issues #2681, #2717) + if getattr(ctx, "is_fsdp2", False) and isinstance(weight_fp8, QuantizedTensorStorage): + clear_columnwise_cache(weight_fp8) + # Prepare grad input tensor # Note: Perform tensor-parallel communication if ctx.ub_overlap_rs_dgrad: @@ -1597,7 +1628,7 @@ def forward( linear_fn = _Linear.forward autograd_ctx = [None] - cache_name = None if is_first_microbatch is None else "weight" + cache_name = None if (is_first_microbatch is None or self.is_fsdp2) else "weight" weight_workspace = ( self._fp8_workspaces.get(cache_name) if cache_name is not None else None ) @@ -1659,6 +1690,7 @@ def forward( backward_override, custom, backward_input_needs_gather, + self.is_fsdp2, ) out, new_weight_workspace = linear_fn( *autograd_ctx, diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index ba44c7a61..1802b7fcc 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -1043,6 +1043,21 @@ def _nvfp4_2d_multi_tensor_transpose(nvfp4_tensors: List[NVFP4Tensor]): ) +def clear_columnwise_cache(tensor: QuantizedTensorStorage) -> None: + """Clear the columnwise cache of a quantized tensor. + Use-case: FSDP2, where TE allocates allgathered + columnwise data(by deriving it out of allgathered rowwise data) + in fsdp2 hooks. And so FSDP2 cant deallocate it when it's done with it""" + if hasattr(tensor, "_columnwise_data"): + tensor._columnwise_data = None + if hasattr(tensor, "_columnwise_scale_inv"): + tensor._columnwise_scale_inv = None + if hasattr(tensor, "_transpose"): + tensor._transpose = None + if hasattr(tensor, "_transpose_invalid"): + tensor._transpose_invalid = True + + def is_custom(x: Optional[Union[Quantizer, QuantizedTensorStorage]] = None) -> bool: """Check if an object is custom. From 3c62f42725ef57a5ddda90104a77dcd349693169 Mon Sep 17 00:00:00 2001 From: vcherepanov-nv Date: Wed, 22 Apr 2026 09:27:39 -0700 Subject: [PATCH 010/180] Make NS coefficients parameter 2D in Python API (#2904) Signed-off-by: Vladimir Cherepanov --- .../pytorch/distributed/run_newton_schulz.py | 7 +++-- transformer_engine/pytorch/newton_schulz.py | 28 ++++++++++++------- 2 files changed, 22 insertions(+), 13 deletions(-) diff --git a/tests/pytorch/distributed/run_newton_schulz.py b/tests/pytorch/distributed/run_newton_schulz.py index bbd073344..712d83bd1 100644 --- a/tests/pytorch/distributed/run_newton_schulz.py +++ b/tests/pytorch/distributed/run_newton_schulz.py @@ -21,11 +21,12 @@ ) -def newton_schulz_reference(in_x: torch.Tensor, coefficients: list[float]) -> torch.Tensor: +def newton_schulz_reference( + in_x: torch.Tensor, coefficients: list[tuple[float, float, float]] +) -> torch.Tensor: """Local Newton-Schulz reference mirroring the provided Octave update.""" x = in_x.clone() - for i in range(len(coefficients) // 3): - a, b, c = coefficients[3 * i : 3 * (i + 1)] + for a, b, c in coefficients: xxt = x @ x.mT x = a * x + b * xxt @ x + c * xxt @ xxt @ x return x diff --git a/transformer_engine/pytorch/newton_schulz.py b/transformer_engine/pytorch/newton_schulz.py index 236789756..1cbe6ebfb 100644 --- a/transformer_engine/pytorch/newton_schulz.py +++ b/transformer_engine/pytorch/newton_schulz.py @@ -5,7 +5,7 @@ """Distributed Newton-Schulz matrix orthogonalization via cuSolverMp.""" from itertools import chain, cycle, islice, repeat -from typing import Iterator, List, Literal, Optional, Sequence +from typing import Iterator, Literal, Optional, Sequence import torch import torch.distributed as dist @@ -63,13 +63,14 @@ NSCoeffT = Literal[_COEFFICIENT_SETS.keys()] CoeffIterMode = Literal["cycle", "repeat_last"] +CoeffT = tuple[float, float, float] def get_coefficient_iterator( steps: int, - coefficient_sets: Sequence[tuple[float, float, float]], + coefficient_sets: Sequence[CoeffT], mode: CoeffIterMode = "cycle", -) -> Iterator[tuple[float, float, float]]: +) -> Iterator[CoeffT]: """Iterate through coefficient sets with configurable end behavior using itertools. Args: @@ -89,7 +90,7 @@ def get_coefficient_iterator( if not coefficient_sets: raise ValueError("coefficient_sets must be non-empty.") - base: Iterator[tuple[float, float, float]] + base: Iterator[CoeffT] if mode == "cycle": base = cycle(coefficient_sets) elif mode == "repeat_last": @@ -101,7 +102,7 @@ def get_coefficient_iterator( return islice(base, steps) -def get_coefficients(steps: int, coefficient_type: NSCoeffT = "quintic") -> List[float]: +def get_coefficients(steps: int, coefficient_type: NSCoeffT = "quintic") -> list[CoeffT]: """Return the coefficient schedule for Newton-Schulz. Parameter ``coefficient_type`` can be one of the following @@ -119,7 +120,7 @@ def get_coefficients(steps: int, coefficient_type: NSCoeffT = "quintic") -> List coeff_iter = get_coefficient_iterator( steps, _COEFFICIENT_SETS[coefficient_type], mode=iter_mode ) - return list(chain.from_iterable(coeff_iter)) + return list(coeff_iter) class CusolverMpCtx: @@ -159,7 +160,7 @@ def newton_schulz( x: torch.Tensor, ctx: CusolverMpCtx, num_iterations: int = 5, - coefficients: Optional[List[float]] = None, + coefficients: Optional[Sequence[CoeffT]] = None, ) -> None: """Compute Newton-Schulz matrix orthogonalization in-place on a distributed matrix. @@ -173,16 +174,23 @@ def newton_schulz( cuSolverMp context created by :func:`cusolvermp_ctx_create`. num_iterations : int, optional Number of Newton-Schulz iterations. Default: 5. - coefficients : list of float, optional + coefficients : sequence of tuple[float, float, float], optional Polynomial coefficients for the Newton-Schulz iteration. """ if coefficients is None: coefficients = get_coefficients(num_iterations) - if len(coefficients) != num_iterations * 3: + if len(coefficients) != num_iterations: raise ValueError( f"Unexpected number of coefficients: {len(coefficients)} for" f" {num_iterations} iterations" ) + flat_coefficients: list[float] = [] + for i, coeff in enumerate(coefficients): + if len(coeff) != 3: + raise ValueError( + f"Expected coefficient tuple of length 3 at iteration {i}, got {len(coeff)}" + ) + flat_coefficients.extend(coeff) if x.dim() != 2: raise ValueError(f"Expected 2D tensor, got {x.dim()}D") @@ -197,4 +205,4 @@ def newton_schulz( m = x.size(0) n = x.size(1) * ctx.nranks - tex.newton_schulz(ctx._ptr, m, n, x, num_iterations, coefficients) + tex.newton_schulz(ctx._ptr, m, n, x, num_iterations, flat_coefficients) From a5164fe1f3cac3fd41ee5677f9884a8a01180cb4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Thu, 23 Apr 2026 10:52:47 +0200 Subject: [PATCH 011/180] [PyTorch] [torch.compile] Remove internal tensor state from Float8CurrentScalingQuantizer (#2816) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Make Float8CurrentScalingQuantizer stateless (no amax/scale members) Remove amax, scale, and use_existing_amax from Float8CurrentScalingQuantizer on both C++ and Python sides. All amax/scale allocations are now ad-hoc at quantization time: - quantize() allocates a combined 2-element tensor for amax+scale - quantize_with_amax() accepts amax as a parameter - create_unquantized_tensor_with_amax() returns amax in a tuple - set_quantization_params() is now a no-op Update all call sites in activation.cpp, bias.cpp, normalization.cpp, attention.cpp, and cast.cpp to propagate the amax buffer. For FusedAdam FP8 kernel: when scale_ptr is null (current scaling), derive scale from scale_inv and skip writing amax/scale_inv metadata. Python side passes empty(0) tensors for scale/amax to signal this. Signed-off-by: Pawel Gadzinski Made-with: Cursor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Guard quantizer.amax/scale access with recipe.delayed() checks Add delayed() guards in context_parallel.py for amax writeback and scale cloning, since Float8CurrentScalingQuantizer no longer has these attributes. Allocate scratch scale buffer in _cast_master_weights_to_fp8_current_scaling instead of reading quantizer.scale. Signed-off-by: Pawel Gadzinski Made-with: Cursor * Fix pylint unused-argument warnings in Float8CurrentScalingQuantizer The device, use_existing_amax, scale, and amax parameters are kept for backward compatibility but not used internally. Signed-off-by: Pawel Gadzinski Made-with: Cursor * Clear dangling scale/amax pointers in Float8CurrentScalingQuantizer::quantize_impl amax_buf and scale_buf are caller-owned local tensors whose storage is released as soon as quantize() returns, leaving raw pointers stored in `out` dangling. No current caller reads out.scale()/out.amax() after quantize_impl returns, so this is currently safe, but it is a silent invariant that could turn into a use-after-free if new callers are added. Defensively clear both pointers at the end of quantize_impl (and in the empty-input early return), mirroring the existing set_amax(nullptr, ...) call already present before nvte_quantize_v2. Signed-off-by: Pawel Gadzinski Made-with: Cursor * Route current-scaling Float8Tensor through FP32 master path in FusedAdam Float8Tensor with Float8CurrentScalingQuantizer now goes through the FP32 master + requantize path (same as MXFP8/NVFP4/blockwise) instead of the fused FP8 Adam kernel. The fused FP8 kernel stays for delayed- scaling Float8Tensor only. Also revert adam.cu to upstream — current scaling no longer needs the scale_inv-derived path in the kernel. Signed-off-by: Pawel Gadzinski Made-with: Cursor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Restore intent/perf comments in FusedAdam quantized dispatch - Note about a possible fused Adam+requantize kernel removing the FP32 round-trip for the QuantizedTensor path. - Justification for casting BF16/FP16 grads to FP32 before the FP32 Adam kernel. Also drop a stray f-string prefix on a literal-only RuntimeError part. Signed-off-by: Pawel Gadzinski Made-with: Cursor * tests skip Signed-off-by: Pawel Gadzinski * tests skip Signed-off-by: Pawel Gadzinski * Warn on deprecated kwargs in Float8CurrentScalingQuantizer Signed-off-by: Pawel Gadzinski * Guard remaining quantizer.scale/amax accesses with recipe.delayed() Followup to e06253d7. Three more spots accessed Float8CurrentScaling- Quantizer.scale/.amax (which no longer exist after the stateless refactor): - AttnFuncWithCPAndKVAllGather.forward used `not mxfp8()` instead of `delayed()` when cloning .scale, causing AttributeError under Float8CurrentScaling. - AttnFuncWithCPAndKVP2P fwd/bwd assigned .amax on per-step quantizer copies under `not mxfp8()`; harmless for current scaling (Python attaches a dynamic attribute that nobody reads) but inconsistent with the delayed()-guarded amax aggregation that follows. - print_quantizers (debug-only) read .scale/.amax for both DS and CS; restrict to DS only. Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../fsdp2_tests/run_fsdp2_fused_adam.py | 50 ++++++++++++++--- .../dot_product_attention/context_parallel.py | 14 ++--- .../attention/dot_product_attention/utils.py | 2 +- transformer_engine/pytorch/csrc/common.h | 18 +++--- .../pytorch/csrc/extensions/activation.cpp | 8 +-- .../pytorch/csrc/extensions/attention.cpp | 40 +++++++------- .../pytorch/csrc/extensions/bias.cpp | 4 +- .../pytorch/csrc/extensions/cast.cpp | 18 +----- .../pytorch/csrc/extensions/normalization.cpp | 10 ++-- transformer_engine/pytorch/csrc/quantizer.cpp | 55 ++++++++++--------- .../pytorch/optimizers/fused_adam.py | 29 ++++------ .../pytorch/tensor/float8_tensor.py | 32 +++++------ transformer_engine/pytorch/tensor/utils.py | 2 +- 13 files changed, 147 insertions(+), 135 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index 83c3f5b56..ecda481ed 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -433,10 +433,15 @@ def test_fused_adam_fp8_no_master(recipe_name): """ recipe = get_recipe_from_string(recipe_name) - if recipe_name in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): + if recipe_name in ( + "MXFP8BlockScaling", + "Float8BlockScaling", + "NVFP4BlockScaling", + "Float8CurrentScaling", + ): pytest.xfail( f"{recipe_name}: FusedAdam without master_weights does not support " - "block-scaling quantized tensors. Use master_weights=True." + "this quantized tensor type. Use master_weights=True." ) world_size, device = _get_dist_info() @@ -825,11 +830,21 @@ def test_dcp_output_parity(recipe_name, async_save): with te.autocast(enabled=True, recipe=recipe): loaded_output = model2(x) - if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): - # DelayedScaling stores amax history and scaling factors in _extra_state, - # which cannot be saved via DCP due to non-deterministic pickle sizes - # across ranks. The fresh model therefore uses default scaling factors, - # producing small numerical differences from FP8 re-quantization. + # DelayedScaling: amax history and scaling factors live in _extra_state, + # which cannot be saved via DCP due to non-deterministic pickle sizes + # across ranks; the fresh model uses default scaling factors, producing + # small numerical differences from FP8 re-quantization. + # Float8CurrentScaling: Float8Tensor._scale_inv is passed via + # fsdp_pre_all_gather metadata rather than as a sharded tensor, so DCP + # saves it cast to the model's param_dtype (bf16) instead of fp32; the + # precision loss in the reloaded scale_inv prevents bitwise parity. + if isinstance( + recipe, + ( + transformer_engine.common.recipe.DelayedScaling, + transformer_engine.common.recipe.Float8CurrentScaling, + ), + ): torch.testing.assert_close( loaded_output, ref_output, @@ -861,7 +876,13 @@ def test_dcp_output_parity(recipe_name, async_save): loss2.backward() optimizer2.step() - if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + if isinstance( + recipe, + ( + transformer_engine.common.recipe.DelayedScaling, + transformer_engine.common.recipe.Float8CurrentScaling, + ), + ): torch.testing.assert_close( out2, out1, @@ -1023,7 +1044,18 @@ def test_dcp_resharding_load(recipe_name): if rank == 0: ref_output = torch.load(ref_output_path, weights_only=True) - if isinstance(recipe, transformer_engine.common.recipe.DelayedScaling): + # DelayedScaling and Float8CurrentScaling use loose tolerance because + # Float8Tensor._scale_inv is passed via fsdp_pre_all_gather metadata + # rather than as a sharded tensor, so DCP saves it cast to the model's + # param_dtype (bf16) instead of fp32. The resulting precision loss in + # the reloaded scale_inv prevents bitwise-identical output parity. + if isinstance( + recipe, + ( + transformer_engine.common.recipe.DelayedScaling, + transformer_engine.common.recipe.Float8CurrentScaling, + ), + ): torch.testing.assert_close( loaded_output, ref_output, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index dfc15cc6c..cd7ce8c98 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -1556,7 +1556,7 @@ def forward( for i in range(cp_size): S_quantizer_per_step[i] = S_quantizer.copy() if S_quantizer is not None else None O_quantizer_per_step[i] = O_quantizer.copy() - if not fp8_recipe.mxfp8(): + if fp8_recipe.delayed(): S_quantizer_per_step[i].amax = amax_per_step[0][i].reshape((1,)) O_quantizer_per_step[i].amax = amax_per_step[1][i].reshape((1,)) else: @@ -2042,7 +2042,7 @@ def forward( ) # update FP8 quantizers: amax across cp_size steps - if fp8 and use_fused_attention and not fp8_recipe.mxfp8(): + if fp8 and use_fused_attention and fp8_recipe.delayed(): amax_cp_fwd = amax_per_step.amax(dim=1) S_quantizer.amax.copy_(amax_cp_fwd[0]) O_quantizer.amax.copy_(amax_cp_fwd[1]) @@ -2182,7 +2182,7 @@ def forward( ctx.QKV_quantizer = QKV_quantizer.copy() ctx.O_quantizer = O_quantizer.copy() ctx.S_quantizer = S_quantizer.copy() if S_quantizer is not None else None - if not ctx.fp8_recipe.mxfp8(): + if fp8_recipe.delayed(): ctx.QKV_quantizer.scale = QKV_quantizer.scale.clone() ctx.O_quantizer.scale = O_quantizer.scale.clone() ctx.S_quantizer.scale = S_quantizer.scale.clone() @@ -2380,7 +2380,7 @@ def backward(ctx, dout, *_args): ctx.dP_quantizer.copy() if ctx.dP_quantizer is not None else None ) dQKV_quantizer_per_step[i] = ctx.dQKV_quantizer.copy() - if not ctx.fp8_recipe.mxfp8(): + if ctx.fp8_recipe.delayed(): dP_quantizer_per_step[i].amax = amax_per_step[0][i].reshape((1,)) dQKV_quantizer_per_step[i].amax = amax_per_step[1][i].reshape((1,)) else: @@ -2793,7 +2793,7 @@ def backward(ctx, dout, *_args): # sum up all cp_size for dq, dk, dv if ctx.fp8 and ctx.use_fused_attention: - if not ctx.fp8_recipe.mxfp8(): + if ctx.fp8_recipe.delayed(): amax_cp_bwd = amax_per_step.amax(dim=1) ctx.dP_quantizer.amax.copy_(amax_cp_bwd[0]) ctx.dQKV_quantizer.amax.copy_(amax_cp_bwd[1]) @@ -3405,7 +3405,7 @@ def forward( ctx.QKV_quantizer = QKV_quantizer.copy() ctx.O_quantizer = O_quantizer.copy() ctx.S_quantizer = S_quantizer.copy() if S_quantizer is not None else None - if not ctx.fp8_recipe.mxfp8(): + if ctx.fp8_recipe.delayed(): ctx.QKV_quantizer.scale = QKV_quantizer.scale.clone() ctx.O_quantizer.scale = O_quantizer.scale.clone() ctx.S_quantizer.scale = S_quantizer.scale.clone() @@ -4232,7 +4232,7 @@ def forward( ctx.QKV_quantizer = QKV_quantizer.copy() ctx.O_quantizer = O_quantizer.copy() ctx.S_quantizer = S_quantizer.copy() if S_quantizer is not None else None - if not ctx.fp8_recipe.mxfp8(): + if fp8_recipe.delayed(): ctx.QKV_quantizer.scale = QKV_quantizer.scale.clone() ctx.O_quantizer.scale = O_quantizer.scale.clone() ctx.S_quantizer.scale = S_quantizer.scale.clone() diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index c416e49da..0fb1a2e3f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2452,7 +2452,7 @@ def print_quantizers( type_str = "CS" elif isinstance(q, MXFP8Quantizer): type_str = "MXFP8" - if type_str in ["DS", "CS"]: + if type_str == "DS": print( f"{label} >> {names[i]:14s}: {type_str}, {q.scale.item():.4e} x" f" {q.amax.item():.4e} = {q.scale.item()*q.amax.item():.4e}" diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index e40d39ee2..8e3bcdd5b 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -195,9 +195,7 @@ class Float8Quantizer : public Quantizer { class Float8CurrentScalingQuantizer : public Quantizer { public: - at::Tensor scale; - at::Tensor scale_inv; - at::Tensor amax; + DType dtype; bool with_amax_reduction; c10::intrusive_ptr amax_reduction_group; bool force_pow_2_scales = false; @@ -217,12 +215,13 @@ class Float8CurrentScalingQuantizer : public Quantizer { py::object quantizer, const std::optional& first_dims, size_t logical_first_dim, size_t logical_last_dim) const override; - /*! @brief Construct an unquantized tensor that shares the quantizer's amax pointer. + /*! @brief Construct an unquantized tensor with a freshly allocated amax buffer. * * The amax is zeroed out. Most TE kernels that output amax expect - * amax to be initialized to zero. + * amax to be initialized to zero. The amax tensor is returned as + * the third element to keep it alive in the caller's scope. */ - std::pair create_unquantized_tensor_with_amax( + std::tuple create_unquantized_tensor_with_amax( const std::vector& shape, DType dtype, std::optional data = std::nullopt); std::pair convert_and_update_tensor(py::object shape) const override; @@ -232,16 +231,17 @@ class Float8CurrentScalingQuantizer : public Quantizer { /*! @brief Quantize to FP8, skipping local amax computation * - * The quantizer's amax pointer is assumed to already hold the local + * The provided amax tensor is assumed to already hold the local * amax. The amax may still be reduced across the amax reduction * group. */ - void quantize_with_amax(TensorWrapper& input, TensorWrapper& out, + void quantize_with_amax(TensorWrapper& input, TensorWrapper& out, at::Tensor amax, const std::optional& noop_flag = std::nullopt); private: void quantize_impl(const TensorWrapper& input, TensorWrapper& out, - const std::optional& noop_flag, bool compute_amax); + const std::optional& noop_flag, bool compute_amax, + at::Tensor amax_buf, at::Tensor scale_buf); }; class Float8BlockQuantizer : public Quantizer { diff --git a/transformer_engine/pytorch/csrc/extensions/activation.cpp b/transformer_engine/pytorch/csrc/extensions/activation.cpp index 99b9c1fef..2df3b6655 100644 --- a/transformer_engine/pytorch/csrc/extensions/activation.cpp +++ b/transformer_engine/pytorch/csrc/extensions/activation.cpp @@ -86,7 +86,7 @@ py::object activation_helper(const at::Tensor& input, py::handle quantizer, int { auto fp8_quantizer_cpp = dynamic_cast(quantizer_cpp.get()); NVTE_CHECK(fp8_quantizer_cpp != nullptr, "Could not cast to FP8 current scaling quantizer"); - auto [temp_nvte, _] = + auto [temp_nvte, _, amax_buf] = fp8_quantizer_cpp->create_unquantized_tensor_with_amax(output_shape, fake_dtype); NVTE_SCOPED_GIL_RELEASE({ if constexpr (act_func == nullptr) { @@ -96,7 +96,7 @@ py::object activation_helper(const at::Tensor& input, py::handle quantizer, int act_func(input_nvte.data(), temp_nvte.data(), stream); } }); - fp8_quantizer_cpp->quantize_with_amax(temp_nvte, out_nvte); + fp8_quantizer_cpp->quantize_with_amax(temp_nvte, out_nvte, amax_buf); } break; case Impl::FUSED_ACTIVATION_AMAX_NVFP4: @@ -198,7 +198,7 @@ py::object dactivation_helper(const at::Tensor& grad_output, const at::Tensor& i { auto fp8_quantizer_cpp = dynamic_cast(quantizer_cpp.get()); NVTE_CHECK(fp8_quantizer_cpp != nullptr, "Could not cast to FP8 current scaling quantizer"); - auto [temp_nvte, _] = + auto [temp_nvte, _, amax_buf] = fp8_quantizer_cpp->create_unquantized_tensor_with_amax(input_shape, fake_dtype); NVTE_SCOPED_GIL_RELEASE({ if constexpr (dact_func == nullptr) { @@ -208,7 +208,7 @@ py::object dactivation_helper(const at::Tensor& grad_output, const at::Tensor& i dact_func(grad_output_nvte.data(), input_nvte.data(), temp_nvte.data(), stream); } }); - fp8_quantizer_cpp->quantize_with_amax(temp_nvte, grad_input_nvte); + fp8_quantizer_cpp->quantize_with_amax(temp_nvte, grad_input_nvte, amax_buf); } break; case Impl::FUSED_ACTIVATION_AMAX_NVFP4: diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 8a2e54a73..e6781bd58 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -55,13 +55,13 @@ NVTE_Fused_Attn_Backend get_fused_attn_backend( } // helper function for S and dP quantizers -std::pair quantizer_helper(py::handle quantizer, - const std::vector &shape, DType dtype, - bool create_hp_tensor, - std::optional data) { +std::tuple> quantizer_helper( + py::handle quantizer, const std::vector &shape, DType dtype, bool create_hp_tensor, + std::optional data) { std::unique_ptr T_quantizer = convert_quantizer(quantizer); TensorWrapper te_T; py::object py_T; + std::optional amax_buf; if (quantizer.is_none()) { // high precision auto *none_quantizer = dynamic_cast(T_quantizer.get()); @@ -80,10 +80,11 @@ std::pair quantizer_helper(py::handle quantizer, auto *T_quantizer_fp8 = dynamic_cast(T_quantizer.get()); if (create_hp_tensor) { if (data.has_value()) { - std::tie(te_T, py_T) = + std::tie(te_T, py_T, amax_buf) = T_quantizer_fp8->create_unquantized_tensor_with_amax(shape, dtype, data.value()); } else { - std::tie(te_T, py_T) = T_quantizer_fp8->create_unquantized_tensor_with_amax(shape, dtype); + std::tie(te_T, py_T, amax_buf) = + T_quantizer_fp8->create_unquantized_tensor_with_amax(shape, dtype); } } else { std::tie(te_T, py_T) = T_quantizer_fp8->create_tensor(shape, dtype); @@ -106,7 +107,7 @@ std::pair quantizer_helper(py::handle quantizer, "MXFP8Quantizer::create_tensor() does not take data tensor as input!"); } } - return {std::move(te_T), std::move(py_T)}; + return {std::move(te_T), std::move(py_T), std::move(amax_buf)}; } // fused attention FWD with separate Q, K and V tensors @@ -138,13 +139,9 @@ std::vector fused_attn_fwd( const DType qkv_type = te_Q.dtype(); // create S tensor - TensorWrapper te_S; - py::object py_S; - std::tie(te_S, py_S) = quantizer_helper(s_quantizer, {0}, DType::kFloat32, false, std::nullopt); + auto [te_S, py_S, _] = quantizer_helper(s_quantizer, {0}, DType::kFloat32, false, std::nullopt); // create O tensor - TensorWrapper te_O; - py::object py_O; std::unique_ptr O_quantizer = convert_quantizer(o_quantizer); std::vector q_shape = convertShape(te_Q.shape()); std::vector v_shape = convertShape(te_V.shape()); @@ -156,7 +153,8 @@ std::vector fused_attn_fwd( size_t h = o_parsed.h(), d = o_parsed.d(); o_parsed.to_format(o_format, o_shape.data()); const DType fake_dtype_te = GetTransformerEngineDType(fake_dtype); - std::tie(te_O, py_O) = quantizer_helper(o_quantizer, o_shape, fake_dtype_te, true, std::nullopt); + auto [te_O, py_O, o_amax_buf] = + quantizer_helper(o_quantizer, o_shape, fake_dtype_te, true, std::nullopt); // construct NVTE tensors TensorWrapper te_Bias; @@ -351,15 +349,14 @@ std::vector fused_attn_bwd( te_dO = makeTransformerEngineTensor(dO, none); // create S and dP tensors - TensorWrapper te_S, te_dP; - py::object py_S, py_dP; - std::tie(te_S, py_S) = quantizer_helper(s_quantizer, {0}, DType::kFloat32, false, std::nullopt); - std::tie(te_dP, py_dP) = + auto [te_S, py_S, _s] = quantizer_helper(s_quantizer, {0}, DType::kFloat32, false, std::nullopt); + auto [te_dP, py_dP, _dp] = quantizer_helper(dp_quantizer, {0}, DType::kFloat32, false, std::nullopt); // create dQ, dK, dV tensors TensorWrapper te_dQ, te_dK, te_dV; py::object py_dQ, py_dK, py_dV; + std::optional dq_amax_buf, dk_amax_buf, dv_amax_buf; std::unique_ptr dQKV_quantizer = convert_quantizer(dqkv_quantizer); std::vector q_shape = convertShape(te_Q.shape()); std::vector k_shape = convertShape(te_K.shape()); @@ -465,9 +462,12 @@ std::vector fused_attn_bwd( NVTE_ERROR("QKV layout not supported!"); } - std::tie(te_dQ, py_dQ) = quantizer_helper(dqkv_quantizer, dQ_shape, dqkv_fake_dtype, true, dQ); - std::tie(te_dK, py_dK) = quantizer_helper(dqkv_quantizer, dK_shape, dqkv_fake_dtype, true, dK); - std::tie(te_dV, py_dV) = quantizer_helper(dqkv_quantizer, dV_shape, dqkv_fake_dtype, true, dV); + std::tie(te_dQ, py_dQ, dq_amax_buf) = + quantizer_helper(dqkv_quantizer, dQ_shape, dqkv_fake_dtype, true, dQ); + std::tie(te_dK, py_dK, dk_amax_buf) = + quantizer_helper(dqkv_quantizer, dK_shape, dqkv_fake_dtype, true, dK); + std::tie(te_dV, py_dV, dv_amax_buf) = + quantizer_helper(dqkv_quantizer, dV_shape, dqkv_fake_dtype, true, dV); // construct NVTE tensors if (detail::IsFloat8Quantizers(dqkv_quantizer.ptr())) { diff --git a/transformer_engine/pytorch/csrc/extensions/bias.cpp b/transformer_engine/pytorch/csrc/extensions/bias.cpp index c59e3c4f6..0cf2025f1 100644 --- a/transformer_engine/pytorch/csrc/extensions/bias.cpp +++ b/transformer_engine/pytorch/csrc/extensions/bias.cpp @@ -208,14 +208,14 @@ std::vector dact_dbias( dynamic_cast(quantizer_cpp.get()); NVTE_CHECK(fp8_quantizer_cpp != nullptr, "Invalid quantizer for fused dact-amax kernel impl"); - auto [temp_nvte, temp_py] = + auto [temp_nvte, temp_py, amax_buf] = fp8_quantizer_cpp->create_unquantized_tensor_with_amax(input_shape, grad_output_dtype); NVTE_SCOPED_GIL_RELEASE({ dact_func(grad_output_nvte.data(), act_input_nvte.data(), temp_nvte.data(), stream); }); const auto temp_torch = temp_py.cast(); at::sum_out(grad_bias_torch, temp_torch.reshape({-1, bias_size}), {0}); - fp8_quantizer_cpp->quantize_with_amax(temp_nvte, grad_input_nvte); + fp8_quantizer_cpp->quantize_with_amax(temp_nvte, grad_input_nvte, amax_buf); break; } case Impl::FUSED_DACT_AMAX_NVFP4: diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 5fb162c72..50fe4c109 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -42,17 +42,6 @@ py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::ob auto input_contiguous = tensor.contiguous(); auto input_cpp = makeTransformerEngineTensor(input_contiguous); - // Set amax if use_existing_amax = true (only valid for CS) - bool use_existing_amax = false; - if (detail::IsFloat8CurrentScalingQuantizers(quantizer.ptr())) { - use_existing_amax = quantizer.attr("use_existing_amax").cast(); - if (use_existing_amax) { - const at::Tensor &amax = quantizer.attr("amax").cast(); - input_cpp.set_amax(amax.data_ptr(), GetTransformerEngineDType(amax.scalar_type()), - getTensorShape(amax)); - } - } - // Initialize output tensor TensorWrapper output_cpp; py::object output_py; @@ -71,12 +60,7 @@ py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::ob } // Perform quantization - if (use_existing_amax) { - auto *quantizer_cs = dynamic_cast(quantizer_cpp.get()); - quantizer_cs->quantize_with_amax(input_cpp, output_cpp, noop_flag_cpp); - } else { - quantizer_cpp->quantize(input_cpp, output_cpp, noop_flag_cpp); - } + quantizer_cpp->quantize(input_cpp, output_cpp, noop_flag_cpp); return output_py; } diff --git a/transformer_engine/pytorch/csrc/extensions/normalization.cpp b/transformer_engine/pytorch/csrc/extensions/normalization.cpp index 3214c3a9d..fb4c7aa1c 100644 --- a/transformer_engine/pytorch/csrc/extensions/normalization.cpp +++ b/transformer_engine/pytorch/csrc/extensions/normalization.cpp @@ -145,6 +145,7 @@ std::vector layernorm_fwd(py::handle input, py::handle weight, Maybe // Construct unquantized output tensor if needed TensorWrapper unquantized_out_nvte; py::object unquantized_out; + at::Tensor amax_buf; TensorWrapper *kernel_out_nvte = &out_nvte; switch (impl) { case Impl::UNFUSED: { @@ -154,7 +155,7 @@ std::vector layernorm_fwd(py::handle input, py::handle weight, Maybe } break; case Impl::FUSED_NORM_AMAX_FP8: { auto fp8_quantizer_cpp = static_cast(quantizer_cpp.get()); - std::tie(unquantized_out_nvte, unquantized_out) = + std::tie(unquantized_out_nvte, unquantized_out, amax_buf) = fp8_quantizer_cpp->create_unquantized_tensor_with_amax(shape, out_dtype); kernel_out_nvte = &unquantized_out_nvte; } break; @@ -199,7 +200,7 @@ std::vector layernorm_fwd(py::handle input, py::handle weight, Maybe } break; case Impl::FUSED_NORM_AMAX_FP8: { auto fp8_quantizer_cpp = static_cast(quantizer_cpp.get()); - fp8_quantizer_cpp->quantize_with_amax(unquantized_out_nvte, out_nvte); + fp8_quantizer_cpp->quantize_with_amax(unquantized_out_nvte, out_nvte, amax_buf); } break; case Impl::FUSED_NORM_AMAX_NVFP4: { auto nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); @@ -381,6 +382,7 @@ std::vector rmsnorm_fwd(const py::handle &input, const py::handle &w // Construct unquantized output tensor if needed TensorWrapper unquantized_out_nvte; py::object unquantized_out; + at::Tensor amax_buf; TensorWrapper *kernel_out_nvte = &out_nvte; switch (impl) { case Impl::UNFUSED: { @@ -390,7 +392,7 @@ std::vector rmsnorm_fwd(const py::handle &input, const py::handle &w } break; case Impl::FUSED_NORM_AMAX_FP8: { auto fp8_quantizer_cpp = static_cast(quantizer_cpp.get()); - std::tie(unquantized_out_nvte, unquantized_out) = + std::tie(unquantized_out_nvte, unquantized_out, amax_buf) = fp8_quantizer_cpp->create_unquantized_tensor_with_amax(shape, out_dtype); kernel_out_nvte = &unquantized_out_nvte; } break; @@ -433,7 +435,7 @@ std::vector rmsnorm_fwd(const py::handle &input, const py::handle &w } break; case Impl::FUSED_NORM_AMAX_FP8: { auto fp8_quantizer_cpp = static_cast(quantizer_cpp.get()); - fp8_quantizer_cpp->quantize_with_amax(unquantized_out_nvte, out_nvte); + fp8_quantizer_cpp->quantize_with_amax(unquantized_out_nvte, out_nvte, amax_buf); } break; case Impl::FUSED_NORM_AMAX_NVFP4: { auto nvfp4_quantizer_cpp = static_cast(quantizer_cpp.get()); diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index b59f3fa3c..da91e5c17 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -532,12 +532,7 @@ void Float8Quantizer::quantize(const TensorWrapper& input, TensorWrapper& out, Float8CurrentScalingQuantizer::Float8CurrentScalingQuantizer(const py::handle& quantizer) : Quantizer(quantizer) { - const at::Tensor& scale = quantizer.attr("scale").cast(); - const at::Tensor& amax = quantizer.attr("amax").cast(); - const DType type = quantizer.attr("dtype").cast(); - this->amax = amax; - this->scale = scale; - this->dtype = type; + this->dtype = quantizer.attr("dtype").cast(); // Get amax reduction group if needed const bool with_amax_reduction = quantizer.attr("with_amax_reduction").cast(); @@ -556,14 +551,7 @@ Float8CurrentScalingQuantizer::Float8CurrentScalingQuantizer(const py::handle& q this->amax_epsilon = quantizer.attr("amax_epsilon").cast(); } -void Float8CurrentScalingQuantizer::set_quantization_params(TensorWrapper* tensor) const { - // transfer amax and scale pointer from quantizer to output tensor (only as gpu buffer, no meaningful data in them) - tensor->set_scale(scale.data_ptr(), GetTransformerEngineDType(scale.scalar_type()), - getTensorShape(scale)); - at::TensorOptions opts = opts.dtype(torch::kFloat32).device(torch::kCUDA); - tensor->set_amax(amax.data_ptr(), GetTransformerEngineDType(amax.scalar_type()), - getTensorShape(amax)); -} +void Float8CurrentScalingQuantizer::set_quantization_params(TensorWrapper* tensor) const {} std::pair Float8CurrentScalingQuantizer::create_tensor( const std::vector& shape, DType dtype) const { @@ -748,18 +736,18 @@ std::pair Float8CurrentScalingQuantizer::creat return {std::move(out_cpp), std::move(out_py)}; } -std::pair +std::tuple Float8CurrentScalingQuantizer::create_unquantized_tensor_with_amax(const std::vector& shape, DType dtype, std::optional data) { - amax.zero_(); + const auto opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + at::Tensor amax_buf = at::zeros({1}, opts); auto out = data.has_value() ? NoneQuantizer(py::none()).create_tensor(shape, dtype, data.value()) : NoneQuantizer(py::none()).create_tensor(shape, dtype); TensorWrapper out_cpp = std::move(out.first); py::object out_py = std::move(out.second); - out_cpp.set_amax(amax.data_ptr(), GetTransformerEngineDType(amax.scalar_type()), - getTensorShape(amax)); - return {std::move(out_cpp), std::move(out_py)}; + out_cpp.set_amax(amax_buf.data_ptr(), DType::kFloat32, std::vector{1}); + return {std::move(out_cpp), std::move(out_py), std::move(amax_buf)}; } std::pair Float8CurrentScalingQuantizer::convert_and_update_tensor( @@ -856,11 +844,20 @@ std::pair Float8CurrentScalingQuantizer::convert_and_ void Float8CurrentScalingQuantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& out, const std::optional& noop_flag, - bool compute_amax) { + bool compute_amax, at::Tensor amax_buf, + at::Tensor scale_buf) { + out.set_amax(amax_buf.data_ptr(), DType::kFloat32, std::vector{1}); + out.set_scale(scale_buf.data_ptr(), DType::kFloat32, std::vector{1}); + auto stream = at::cuda::getCurrentCUDAStream(); // Nothing to be done if input is empty if (input.numel() == 0) { + // Clear amax/scale pointers defensively: amax_buf/scale_buf are caller-owned + // locals that may be released right after this call, leaving dangling raw + // pointers in `out`. + out.set_amax(nullptr, DType::kFloat32, out.defaultShape); + out.set_scale(nullptr, DType::kFloat32, out.defaultShape); return; } @@ -883,7 +880,7 @@ void Float8CurrentScalingQuantizer::quantize_impl(const TensorWrapper& input, Te // allreduce amax tensor c10d::AllreduceOptions opts; opts.reduceOp = c10d::ReduceOp::MAX; - std::vector tensors = {amax}; + std::vector tensors = {amax_buf}; NVTE_SCOPED_GIL_RELEASE({ amax_reduction_group->allreduce(tensors, opts)->wait(); }); } @@ -893,19 +890,25 @@ void Float8CurrentScalingQuantizer::quantize_impl(const TensorWrapper& input, Te // Cast to FP8 out.set_amax(nullptr, DType::kFloat32, out.defaultShape); // Avoid atomic amax updates NVTE_SCOPED_GIL_RELEASE({ nvte_quantize_v2(input.data(), out.data(), quant_config, stream); }); + + // Clear scale pointer defensively: amax_buf/scale_buf are caller-owned locals + // that may be released right after this call, leaving a dangling raw pointer in `out`. + out.set_scale(nullptr, DType::kFloat32, out.defaultShape); } void Float8CurrentScalingQuantizer::quantize(const TensorWrapper& input, TensorWrapper& out, const std::optional& noop_flag) { - this->quantize_impl(input, out, noop_flag, true); + const auto opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + at::Tensor amax_and_scale = at::empty({2}, opts); + this->quantize_impl(input, out, noop_flag, true, amax_and_scale[0], amax_and_scale[1]); } void Float8CurrentScalingQuantizer::quantize_with_amax( - TensorWrapper& input, TensorWrapper& out, const std::optional& noop_flag) { - NVTE_CHECK(input.get_amax().data_ptr == amax.data_ptr(), - "Input does not use the appropriate amax tensor"); + TensorWrapper& input, TensorWrapper& out, at::Tensor amax, + const std::optional& noop_flag) { + const auto opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); input.set_amax(nullptr, DType::kFloat32, input.defaultShape); - this->quantize_impl(input, out, noop_flag, false); + this->quantize_impl(input, out, noop_flag, false, std::move(amax), at::empty({1}, opts)); } Float8BlockQuantizer::Float8BlockQuantizer(const py::handle& quantizer) : Quantizer(quantizer) { diff --git a/transformer_engine/pytorch/optimizers/fused_adam.py b/transformer_engine/pytorch/optimizers/fused_adam.py index 437dfa829..828c34f53 100644 --- a/transformer_engine/pytorch/optimizers/fused_adam.py +++ b/transformer_engine/pytorch/optimizers/fused_adam.py @@ -656,13 +656,18 @@ def step(self, closure=None, grad_scaler=None): unscaled_lists[name].append(unscaled) scaled_lists[name].append(state_tensor) state_scales[name].append(self._scales[p][name]) - if isinstance(p, Float8Tensor) or ( - isinstance(p, DTensor) and isinstance(p._local_tensor, Float8Tensor) + local_p = p._local_tensor if isinstance(p, DTensor) else p + # Only delayed-scaling Float8Tensor uses the fused FP8 Adam kernel. + # Everything else (MXFP8/NVFP4/blockwise, current-scaling Float8) goes + # through the FP32 master + requantize path. A fused Adam+requantize + # kernel (like multi_tensor_adam_fp8 for delayed-scaling Float8Tensor) + # would avoid the FP32 round-trip in that path. + if isinstance(local_p, Float8Tensor) and isinstance( + local_p._quantizer, Float8Quantizer ): - p = p._local_tensor if isinstance(p, DTensor) else p - out_dtype = p._fp8_dtype - p_fp8_model.append(p._data.data) - scale, amax, scale_inv = get_fp8_meta(p) + out_dtype = local_p._fp8_dtype + p_fp8_model.append(local_p._data.data) + scale, amax, scale_inv = get_fp8_meta(local_p) scales.append(scale) amaxes.append(amax) scale_invs.append(scale_inv) @@ -671,22 +676,12 @@ def step(self, closure=None, grad_scaler=None): g_of_fp8_model.append(p_grad.data) m_of_fp8_model.append(unscaled_state["exp_avg"]) v_of_fp8_model.append(unscaled_state["exp_avg_sq"]) - elif isinstance(p, QuantizedTensor) or ( - isinstance(p, DTensor) and isinstance(p._local_tensor, QuantizedTensor) - ): - # Block-scaling quantized params (MXFP8Tensor, Float8BlockwiseQTensor, - # NVFP4Tensor). Operate on FP32 master weights, requantize back after - # Adam update. - # Note: a fused Adam+requantize kernel (like multi_tensor_adam_fp8 - # for Float8Tensor) would avoid the FP32 round-trip here. + elif isinstance(local_p, QuantizedTensor): if not self.master_weights: - local_p = p._local_tensor if isinstance(p, DTensor) else p raise RuntimeError( "FusedAdam without master_weights does not support " f"{type(local_p).__name__} parameters. Use master_weights=True." ) - # Route to the FP32 master-weight path: Adam updates the FP32 master, - # then we write back to the quantized param after kernels run. # Gradients may be BF16/FP16 from the backward pass — cast to FP32 # to match the FP32 Adam kernel expectations. p_f32_model.append(unscaled_state["master_param"].data) diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index 2c828aaaa..ed6091c85 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -233,23 +233,21 @@ class Float8CurrentScalingQuantizer(Quantizer): high-precision tensor, without the need of any history window. Unlike delayed scaling, scale and amax tensors are not needed to initialize the - quantizer, becuse they are simply GPU buffers that will be filled by current + quantizer, because they are simply GPU buffers that will be filled by current scaling quantization kernels, instead of using values taken from delayed scaling - history window. Therefore, device parameter is needed for tensor allocation. + history window. Both Float8CurrentScalingQuantizer and Float8Quantizer produces Float8Tensor, because they are both per-tensor scaling, ie. one scaling factor per tensor. + Note: The ``device``, ``use_existing_amax``, ``scale``, and ``amax`` + parameters are accepted but unused. They are kept for backward + compatibility with existing callers. + """ - """Workspace buffer for FP8 scaling factor""" - scale: torch.Tensor - """Workspace buffer for max-abs value""" - amax: torch.Tensor """FP8 datatype""" dtype: TE_DType - """amax update options""" - use_existing_amax: bool """amax reduction options""" with_amax_reduction: bool amax_reduction_group: Optional[dist_group_type] @@ -273,14 +271,15 @@ def __init__( amax: Optional[torch.Tensor] = None, ) -> None: super().__init__(rowwise=rowwise, columnwise=columnwise) - if scale is None: - scale = torch.empty(1, dtype=torch.float32, device=device) - if amax is None: - amax = torch.empty(1, dtype=torch.float32, device=device) - self.scale = scale - self.amax = amax + if use_existing_amax or scale is not None or amax is not None: + warnings.warn( + "Float8CurrentScalingQuantizer ignores `use_existing_amax`, `scale`, " + "and `amax`; kept for backward compatibility and will be removed.", + DeprecationWarning, + stacklevel=2, + ) + del device, use_existing_amax, scale, amax # Kept for backward compatibility self.dtype = fp8_dtype - self.use_existing_amax = use_existing_amax self.with_amax_reduction = with_amax_reduction self.amax_reduction_group = amax_reduction_group self.force_pow_2_scales = force_pow_2_scales @@ -302,11 +301,8 @@ def copy(self) -> Float8CurrentScalingQuantizer: columnwise=self.columnwise_usage, with_amax_reduction=self.with_amax_reduction, amax_reduction_group=self.amax_reduction_group, - use_existing_amax=self.use_existing_amax, force_pow_2_scales=self.force_pow_2_scales, amax_epsilon=self.amax_epsilon, - scale=self.scale, - amax=self.amax, ) quantizer.internal = self.internal quantizer.optimize_for_gemm = self.optimize_for_gemm diff --git a/transformer_engine/pytorch/tensor/utils.py b/transformer_engine/pytorch/tensor/utils.py index 1802b7fcc..8b22097f7 100644 --- a/transformer_engine/pytorch/tensor/utils.py +++ b/transformer_engine/pytorch/tensor/utils.py @@ -352,7 +352,7 @@ def _cast_master_weights_to_fp8_current_scaling( f"expected {amax_epsilon} but got {quantizer.amax_epsilon}" ) - scales.append(quantizer.scale.view(1)) + scales.append(torch.empty(1, dtype=torch.float32, device=device)) scale_invs.append(model_weight._scale_inv.view(1)) # Compute amax of the master weight and store it in packed_amaxes. From 424b031954bcaa05a9088ceadfe6cd8452235e08 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 23 Apr 2026 07:56:03 -0700 Subject: [PATCH 012/180] [PyTorch] Fix CP A2A F16 when NVTE_FP8_DPA_BWD=1 (#2917) fix fp8 and is_bwd_fp8 relationship Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../dot_product_attention/context_parallel.py | 28 +++++++++---------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index cd7ce8c98..7b10593ac 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -1469,7 +1469,8 @@ def forward( fwd_nominal_dtype = q.dtype is_input_fp8 = isinstance(q, QuantizedTensorStorage) is_output_fp8 = fp8_output - is_bwd_fp8 = int(os.getenv("NVTE_FP8_DPA_BWD", "1")) + _use_fp8_dpa_bwd = bool(int(os.getenv("NVTE_FP8_DPA_BWD", "1"))) + is_bwd_fp8 = fp8 and _use_fp8_dpa_bwd # recipe passed in through autocast or set by NVTE_DPA_FP8_RECIPE; # may be different from fp8_meta["recipe"] fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() @@ -2063,20 +2064,17 @@ def forward( # prepare for return and ctx saves out_fp8 = None out_f16 = out.to(fwd_nominal_dtype) - if fp8 and ( - is_output_fp8 - or ( - is_bwd_fp8 - and not (fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16) - and not fp8_recipe.mxfp8() - ) + if (fp8 and is_output_fp8) or ( + is_bwd_fp8 + and not (fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16) + and not fp8_recipe.mxfp8() ): out_fp8 = O_quantizer(out_f16) out_ret = out_fp8 if (fp8 and is_output_fp8) else out_f16 ctx.layer_number = layer_number ctx.fp8_recipe = fp8_recipe - ctx.fp8 = fp8 and is_bwd_fp8 + ctx.fp8 = is_bwd_fp8 kv_fp8 = None kv = p2p_comm_buffers[-1] @@ -3063,7 +3061,8 @@ def forward( ), "q, k, v must be of the same class, e.g. torch.Tensor or QuantizedTensorStorage." is_input_fp8 = isinstance(q, QuantizedTensorStorage) is_output_fp8 = fp8_output - is_bwd_fp8 = int(os.getenv("NVTE_FP8_DPA_BWD", "1")) + _use_fp8_dpa_bwd = bool(int(os.getenv("NVTE_FP8_DPA_BWD", "1"))) + is_bwd_fp8 = fp8 and _use_fp8_dpa_bwd fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None: fp8_recipe = fp8_meta["local_recipes"][0] @@ -3306,12 +3305,12 @@ def forward( or (fp8_recipe.float8_current_scaling() and not _dpa_fp8_cs_o_in_f16) ) ) - if fp8 and (is_output_fp8 or bwd_requires_o_fp8): + if (fp8 and is_output_fp8) or bwd_requires_o_fp8: out_fp8 = O_quantizer(out_f16) out_ret = out_fp8 if is_output_fp8 else out_f16 # save tensors for backward - ctx.fp8 = fp8 and is_bwd_fp8 + ctx.fp8 = is_bwd_fp8 ctx.fp8_recipe = fp8_recipe fp8_tensors = (None, None, None, None) f16_tensors = (None, None, None, None) @@ -3931,7 +3930,8 @@ def forward( ), "q, k, v must be of the same class, e.g. torch.Tensor or QuantizedTensorStorage." is_input_fp8 = isinstance(q, QuantizedTensorStorage) is_output_fp8 = fp8_output - is_bwd_fp8 = int(os.getenv("NVTE_FP8_DPA_BWD", "1")) + _use_fp8_dpa_bwd = bool(int(os.getenv("NVTE_FP8_DPA_BWD", "1"))) + is_bwd_fp8 = fp8 and _use_fp8_dpa_bwd # recipe passed in through autocast or set by NVTE_DPA_FP8_RECIPE; # may be different from fp8_meta["recipe"] fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() @@ -4161,7 +4161,7 @@ def forward( ctx.orig_o_shape = orig_o_shape # save tensors for backward - ctx.fp8 = fp8 and is_bwd_fp8 + ctx.fp8 = is_bwd_fp8 fp8_tensors = (None, None, None, None) f16_tensors = (None, None, None, None) if is_training: From ab60f4c3cf9ecc160d0d866b7786d704800c56f1 Mon Sep 17 00:00:00 2001 From: Dongmin Ra Date: Fri, 24 Apr 2026 01:04:57 +0900 Subject: [PATCH 013/180] fix: scope get_full_cu_seqlens cache key by device and inference mode (#2728) * fix: scope get_full_cu_seqlens cache key by device and inference mode Signed-off-by: Dongmin Ra * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Dongmin Ra Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --- .../attention/test_cu_seqlens_cache.py | 97 +++++++++++++++++++ .../attention/dot_product_attention/utils.py | 11 ++- 2 files changed, 103 insertions(+), 5 deletions(-) create mode 100644 tests/pytorch/attention/test_cu_seqlens_cache.py diff --git a/tests/pytorch/attention/test_cu_seqlens_cache.py b/tests/pytorch/attention/test_cu_seqlens_cache.py new file mode 100644 index 000000000..be4895199 --- /dev/null +++ b/tests/pytorch/attention/test_cu_seqlens_cache.py @@ -0,0 +1,97 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import pytest +import torch + +from transformer_engine.pytorch import DotProductAttention +from transformer_engine.pytorch.attention.dot_product_attention import utils as dpa_utils +from transformer_engine.pytorch.utils import get_cudnn_version + + +pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is required.") + + +@pytest.fixture(autouse=True) +def clear_cu_seqlens_cache(): + dpa_utils._cu_seqlens_cache.clear() + yield + dpa_utils._cu_seqlens_cache.clear() + + +def _make_dpa(device: torch.device) -> DotProductAttention: + return DotProductAttention( + num_attention_heads=2, + kv_channels=16, + attention_dropout=0.0, + qkv_format="bshd", + attn_mask_type="no_mask", + attention_type="self", + ).to(device=device, dtype=torch.float16) + + +def _make_qkv(device: torch.device, requires_grad: bool = False): + shape = (2, 8, 2, 16) + q = torch.randn(*shape, device=device, dtype=torch.float16, requires_grad=requires_grad) + k = torch.randn(*shape, device=device, dtype=torch.float16, requires_grad=requires_grad) + v = torch.randn(*shape, device=device, dtype=torch.float16, requires_grad=requires_grad) + return q, k, v + + +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") +def test_cu_seqlens_cache_isolated_across_devices_for_forward(): + if torch.cuda.device_count() < 2: + pytest.skip("Requires at least 2 CUDA devices.") + + dev0 = torch.device("cuda:0") + dev1 = torch.device("cuda:1") + + dpa0 = _make_dpa(dev0).eval() + dpa1 = _make_dpa(dev1).eval() + + with torch.no_grad(): + q0, k0, v0 = _make_qkv(dev0) + out0 = dpa0(q0, k0, v0, attn_mask_type="no_mask") + + q1, k1, v1 = _make_qkv(dev1) + out1 = dpa1(q1, k1, v1, attn_mask_type="no_mask") + + assert out0.device == dev0 + assert out1.device == dev1 + + expected_key_0 = (2, 8, dev0, False) + expected_key_1 = (2, 8, dev1, False) + assert expected_key_0 in dpa_utils._cu_seqlens_cache + assert expected_key_1 in dpa_utils._cu_seqlens_cache + + assert dpa_utils._cu_seqlens_cache[expected_key_0].device == dev0 + assert dpa_utils._cu_seqlens_cache[expected_key_1].device == dev1 + + +@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") +def test_cu_seqlens_cache_isolated_between_inference_and_train_forward(): + dev = torch.device("cuda:0") + dpa = _make_dpa(dev) + + dpa.eval() + with torch.inference_mode(): + q_inf, k_inf, v_inf = _make_qkv(dev) + out_inf = dpa(q_inf, k_inf, v_inf, attn_mask_type="no_mask") + + inf_key = (2, 8, dev, True) + assert inf_key in dpa_utils._cu_seqlens_cache + assert dpa_utils._cu_seqlens_cache[inf_key].device == dev + + dpa.train() + q_tr, k_tr, v_tr = _make_qkv(dev, requires_grad=True) + out_tr = dpa(q_tr, k_tr, v_tr, attn_mask_type="no_mask") + out_tr.sum().backward() + + train_key = (2, 8, dev, False) + assert train_key in dpa_utils._cu_seqlens_cache + assert dpa_utils._cu_seqlens_cache[train_key].device == dev + + assert out_inf.device == dev + assert out_tr.device == dev + assert dpa_utils._cu_seqlens_cache[inf_key] is not dpa_utils._cu_seqlens_cache[train_key] diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 0fb1a2e3f..3a0322a1c 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1846,11 +1846,12 @@ def _get_cu_seqlens(batch_size, max_seqlen, device): if is_in_onnx_export_mode(): return _get_cu_seqlens(batch_size, max_seqlen, device) - if (batch_size, max_seqlen) not in _cu_seqlens_cache: - _cu_seqlens_cache[(batch_size, max_seqlen)] = _get_cu_seqlens( - batch_size, max_seqlen, device - ) - return _cu_seqlens_cache[(batch_size, max_seqlen)] + + is_inference = torch.is_inference_mode_enabled() + cu_seqlens_cache_key = (batch_size, max_seqlen, device, is_inference) + if cu_seqlens_cache_key not in _cu_seqlens_cache: + _cu_seqlens_cache[cu_seqlens_cache_key] = _get_cu_seqlens(batch_size, max_seqlen, device) + return _cu_seqlens_cache[cu_seqlens_cache_key] @jit_fuser From 9e55a255dd2d63bbf6d2c6ec788d0fd27965b42b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Bj=C3=B6rn=20Buschk=C3=A4mper?= Date: Thu, 23 Apr 2026 18:55:23 +0200 Subject: [PATCH 014/180] [PyTorch] Fix FA4 selection when FA3 is unavailable. (#2909) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Fix FA4 selection when FA3 is unavailable. Signed-off-by: Björn Buschkämper --- .../pytorch/attention/dot_product_attention/utils.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 3a0322a1c..ed8742353 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -473,9 +473,14 @@ def get_attention_backend( # On SM90, prefer FA3 over FA4 when FA3 is available. # FA3 is more mature on Hopper; FA4's SM90 backward has limitations # (MLA, non-standard head dims, SplitKV). - if use_flash_attention_4 and use_flash_attention_3 and device_compute_capability == (9, 0): - if FlashAttentionUtils.v4_is_installed: - logger.debug("Disabling FlashAttention 4 to prefer FlashAttention 3 on SM90") + if ( + device_compute_capability == (9, 0) + and use_flash_attention_3 + and FlashAttentionUtils.v3_is_installed + and use_flash_attention_4 + and FlashAttentionUtils.v4_is_installed + ): + logger.debug("Disabling FlashAttention 4 to prefer FlashAttention 3 on SM90") use_flash_attention_4 = False # Filter: Data type From 0c2e7b09c6d33109803ea089fbf80421a326e0a4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Daniel=20C=C3=A1mpora?= <961215+dcampora@users.noreply.github.com> Date: Thu, 23 Apr 2026 19:00:19 +0200 Subject: [PATCH 015/180] Add optimised top-k kernel AIR. (#2890) * Add AIR TopK support to TE JAX extension Adds a custom AIR TopK implementation (header-only, vendored into transformer_engine/common/util/) exposed as a JAX FFI custom call via the TE JAX extension. Key changes: - transformer_engine/common/util/air_topk.cu: AIR TopK CUDA kernel - transformer_engine/common/util/standalone_air_topk.cuh: vendored header - transformer_engine/common/include/transformer_engine/air_topk.h: C API - transformer_engine/jax/csrc/extensions/air_topk.cpp: JAX FFI binding - transformer_engine/jax/cpp_extensions/air_topk.py: Python wrapper - CMakeLists.txt: compile new kernel; use CCCL from CUDA toolkit - CMakeLists.txt: fix SM100 arch handling when all arches are special-cased Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: dcampora <961215+dcampora@users.noreply.github.com> Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> * Address PR review comments: fix namespace pollution, unused var, missing export, cache sm_cnt - Move WARP_SIZE/WARP_BITS/FULL_WARP_MASK/VECTORIZED_READ_SIZE into namespace nv - Remove unused keys_element_bytes variable in AirTopkFFI; collapse switch to dtype validation - Add missing `from .air_topk import *` export in jax/cpp_extensions/__init__.py - Cache sm_cnt per device with static vars to avoid repeated cudaGetDevice/cudaDeviceGetAttribute calls - Add CMAKE_BUILD_WITH_INSTALL_RPATH=ON to build_ext.py Signed-off-by: dcampora <961215+dcampora@users.noreply.github.com> Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> * Rename air_topk -> topk throughout JAX extension Remove the `air_` prefix from all TopK-related identifiers: file names, C API functions (nvte_air_topk -> nvte_topk), FFI handler/primitive names (te_air_topk_ffi -> te_topk_ffi), Python symbols, and the internal `air_topk` namespace in standalone_topk.cuh. No functional changes. Signed-off-by: Diego Campora Signed-off-by: dcampora <961215+dcampora@users.noreply.github.com> Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> * Address ptrendx review comments and fix CI lint issues - Follow TE workspace convention: remove nvte_get_topk_workspace_bytes() and implement empty-workspace size-query pattern in nvte_topk() instead - Remove unnecessary nv_detail::float_to_T helper; replace usages with static_cast() directly - Remove unrelated CMAKE_CUDA_ARCHITECTURES OFF block from CMakeLists.txt - Fix cpplint errors in standalone_topk.cuh: replace unsigned long long int with uint64_t, add NOLINT for constexpr-sized arrays and else-with-comment Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> * Add assertion for 2D input in topk Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> --------- Signed-off-by: dcampora <961215+dcampora@users.noreply.github.com> Signed-off-by: Daniel Campora <961215+dcampora@users.noreply.github.com> Signed-off-by: Diego Campora Co-authored-by: Claude Sonnet 4.6 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/jax/test_custom_call_compute.py | 81 ++ transformer_engine/common/CMakeLists.txt | 1 + .../common/include/transformer_engine/topk.h | 50 + .../common/util/standalone_topk.cuh | 1266 +++++++++++++++++ transformer_engine/common/util/topk.cu | 60 + .../jax/cpp_extensions/__init__.py | 1 + transformer_engine/jax/cpp_extensions/topk.py | 136 ++ transformer_engine/jax/csrc/extensions.h | 4 + .../jax/csrc/extensions/pybind.cpp | 4 + .../jax/csrc/extensions/topk.cpp | 104 ++ 10 files changed, 1707 insertions(+) create mode 100644 transformer_engine/common/include/transformer_engine/topk.h create mode 100644 transformer_engine/common/util/standalone_topk.cuh create mode 100644 transformer_engine/common/util/topk.cu create mode 100644 transformer_engine/jax/cpp_extensions/topk.py create mode 100644 transformer_engine/jax/csrc/extensions/topk.cpp diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index 3e5529c07..d08f5cc11 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -49,6 +49,7 @@ from transformer_engine.jax.activation import activation from transformer_engine.jax.dense import dense, grouped_dense from transformer_engine.jax.layernorm_dense import layernorm_dense +from transformer_engine.jax.cpp_extensions.topk import topk GEMM_CASES = [ (256, 256, 512), @@ -2035,3 +2036,83 @@ def f(x): actual = load_array_dump("my_tensor_gpu0.bin", shape, dtype) assert_allclose(actual, expected, dtype=dtype) + + +@pytest.mark.parametrize("dtype", [jnp.bfloat16, jnp.float32]) +@pytest.mark.parametrize( + "problem_size", + [ + (1, 10000, 100), + (1, 50000, 200), + (4, 16384, 256), + (8, 65536, 512), + (1, 1000000, 1000), + ], +) +class TestTopK: + """Correctness tests for the TopK JAX primitive. + + Each test generates an input whose top-k entries lie in a known value range + so that correctness can be verified without a full sort, then cross-checks + against jax.lax.top_k as a reference. + """ + + def test_topk_1d(self, dtype, problem_size): + """1-D input: single row.""" + _bs, n, k = problem_size + + prng_key = jax.random.PRNGKey(0) + keys = jax.random.split(prng_key, 3) + topk_vals = jax.random.uniform(keys[0], shape=(k,), dtype=dtype, minval=1.5, maxval=2.5) + bottom_vals = jax.random.uniform( + keys[1], shape=(n - k,), dtype=dtype, minval=0.0, maxval=1.0 + ) + x = jax.random.permutation(keys[2], jnp.concatenate([topk_vals, bottom_vals])) + + ref_vals, ref_idx = jax.jit(jax.lax.top_k, static_argnums=(1,))(x, k) + prim_vals, prim_idx = jax.jit(topk, static_argnums=(1,))(x, k) + + # AIR TopK output is unordered; sort before comparing. + ref_vals, ref_idx = jax.lax.sort_key_val(ref_vals, ref_idx) + prim_vals, prim_idx = jax.lax.sort_key_val(prim_vals, prim_idx) + + assert_allclose(prim_vals, ref_vals, dtype=dtype) + + sorted_x = jax.lax.sort(x) + assert prim_vals[0] >= sorted_x[-(k + 1)] + + # Values at returned indices must match reference. + assert_allclose(x[prim_idx], x[ref_idx], dtype=dtype) + + def test_topk_2d(self, dtype, problem_size): + """2-D input: each row is an independent top-k problem.""" + bs, n, k = problem_size + + prng_key = jax.random.PRNGKey(42) + keys = jax.random.split(prng_key, 3) + topk_vals = jax.random.uniform(keys[0], shape=(bs, k), dtype=dtype, minval=1.5, maxval=2.5) + bottom_vals = jax.random.uniform( + keys[1], shape=(bs, n - k), dtype=dtype, minval=0.0, maxval=1.0 + ) + x_unsorted = jnp.concatenate([topk_vals, bottom_vals], axis=1) + # Shuffle columns independently per row. + col_perm = jax.random.permutation(keys[2], n) + x = x_unsorted[:, col_perm] + + ref_vals, ref_idx = jax.jit(jax.lax.top_k, static_argnums=(1,))(x, k) + prim_vals, prim_idx = jax.jit(topk, static_argnums=(1,))(x, k) + + # Sort each row independently for comparison. + ref_vals, ref_idx = jax.vmap(jax.lax.sort_key_val)(ref_vals, ref_idx) + prim_vals, prim_idx = jax.vmap(jax.lax.sort_key_val)(prim_vals, prim_idx) + + assert_allclose(prim_vals, ref_vals, dtype=dtype) + + # For each row, the smallest selected value must be >= the (k+1)-th largest in that row. + sorted_x = jnp.sort(x, axis=1) + assert jnp.all(prim_vals[:, 0] >= sorted_x[:, -(k + 1)]) + + # Values at returned indices must match reference values. + prim_gathered = jnp.take_along_axis(x, prim_idx, axis=1) + ref_gathered = jnp.take_along_axis(x, ref_idx, axis=1) + assert_allclose(prim_gathered, ref_gathered, dtype=dtype) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 53f9773a7..781fe4881 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -194,6 +194,7 @@ list(APPEND transformer_engine_cuda_sources permutation/permutation.cu util/utils.cu util/padding.cu + util/topk.cu swizzle/swizzle.cu swizzle/swizzle_block_scaling.cu fused_softmax/scaled_masked_softmax.cu diff --git a/transformer_engine/common/include/transformer_engine/topk.h b/transformer_engine/common/include/transformer_engine/topk.h new file mode 100644 index 000000000..3fe9c9447 --- /dev/null +++ b/transformer_engine/common/include/transformer_engine/topk.h @@ -0,0 +1,50 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#ifndef TRANSFORMER_ENGINE_TOPK_H_ +#define TRANSFORMER_ENGINE_TOPK_H_ + +#include "transformer_engine.h" + +#ifdef __cplusplus +extern "C" { +#endif + +/*! \brief Compute the top-K (key, index) pairs using the AIR radix algorithm. + * + * Operates on a batch of rows: each row of length \p seq_len is processed + * independently and the \p k largest entries are selected. + * + * Calling this function with workspace set to an empty tensor will not perform + * the operation, but instead set the shape and type of the workspace tensor to + * the required values. + * + * \param[in] stream CUDA stream used for the operation. + * \param[in] keys_in Input keys tensor, flat storage for + * batch_size rows of seq_len elements. + * \param[in] lengths_in Per-row lengths, shape (batch_size,); int32. + * Fill with seq_len for uniform-length batches. + * \param[in,out] keys_out Output top-k keys, flat storage for + * batch_size rows of k elements. + * \param[in,out] indices_out Output top-k indices (within each row), + * flat storage for batch_size rows of k int32 elements. + * \param[in,out] workspace Workspace tensor. + * \param[in] batch_size Number of rows. + * \param[in] seq_len Number of elements per row. + * \param[in] k Number of top-K entries to select per row. + * + * Supported key dtypes: float32, bfloat16. + * Index dtype: int32. + */ +void nvte_topk(cudaStream_t stream, const NVTETensor keys_in, const NVTETensor lengths_in, + NVTETensor keys_out, NVTETensor indices_out, NVTETensor workspace, int batch_size, + int seq_len, int k); + +#ifdef __cplusplus +} // extern "C" +#endif + +#endif // TRANSFORMER_ENGINE_TOPK_H_ diff --git a/transformer_engine/common/util/standalone_topk.cuh b/transformer_engine/common/util/standalone_topk.cuh new file mode 100644 index 000000000..3d19cbfcf --- /dev/null +++ b/transformer_engine/common/util/standalone_topk.cuh @@ -0,0 +1,1266 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#pragma once + +#include +#include +#include + +#include +#include +#include +namespace cg = cooperative_groups; + +// Workspace pointer-alignment helpers. +inline size_t calc_aligned_size(const std::vector &sizes) { + const size_t ALIGN_BYTES = 256; + const size_t ALIGN_MASK = ~(ALIGN_BYTES - 1); + size_t total = 0; + for (auto sz : sizes) total += (sz + ALIGN_BYTES - 1) & ALIGN_MASK; + return total + ALIGN_BYTES - 1; +} +inline std::vector calc_aligned_pointers(const void *p, const std::vector &sizes) { + const size_t ALIGN_BYTES = 256; + const size_t ALIGN_MASK = ~(ALIGN_BYTES - 1); + char *ptr = + reinterpret_cast((reinterpret_cast(p) + ALIGN_BYTES - 1) & ALIGN_MASK); + std::vector ptrs; + ptrs.reserve(sizes.size()); + for (auto sz : sizes) { + ptrs.push_back(ptr); + ptr += (sz + ALIGN_BYTES - 1) & ALIGN_MASK; + } + return ptrs; +} + +namespace nv { + +constexpr int VECTORIZED_READ_SIZE = 16; +constexpr int WARP_SIZE = 32; +constexpr int WARP_BITS = 5; +constexpr unsigned FULL_WARP_MASK = 0xffffffff; + +namespace topk { +using WideT = float4; + +#ifdef __CUDA_ARCH__ +using ::atomicAdd; +inline __device__ size_t atomicAdd(size_t *address, size_t value) { + static_assert(sizeof(size_t) == sizeof(uint64_t)); + return atomicAdd(reinterpret_cast(address), static_cast(value)); +} +#endif + +template +__host__ __device__ constexpr int calc_num_buckets() { + return 1 << BitsPerPass; +} + +/** + * @brief Provide a ceiling division operation ie. ceil(a / b) + * @tparam IntType supposed to be only integers for now! + */ +template +constexpr __host__ __device__ IntType ceildiv(IntType a, IntType b) { + return (a + b - 1) / b; +} + +/** + * @brief Provide an alignment function ie. ceil(a / b) * b + * @tparam IntType supposed to be only integers for now! + */ +template +constexpr __host__ __device__ IntType alignTo(IntType a, IntType b) { + return ceildiv(a, b) * b; +} + +template +__host__ __device__ constexpr int calc_num_passes() { + return ceildiv(sizeof(T) * 8, BitsPerPass); +} + +__host__ __device__ int round(int num, int round_value) { + return ((num - 1) / round_value + 1) * round_value; +} + +/** + * Bit 0 is the least significant (rightmost); + * this implementation processes input from the most to the least significant + * bit. This way, we can skip some passes in the end at the cost of having an + * unsorted output. + * + * NB: Use pass=-1 for calc_mask(). + */ +template +__device__ constexpr int calc_start_bit(int pass) { + int start_bit = static_cast(sizeof(T) * 8) - (pass + 1) * BitsPerPass; + if (start_bit < 0) { + start_bit = 0; + } + return start_bit; +} + +template +__device__ constexpr unsigned calc_mask(int pass) { + static_assert(BitsPerPass <= 31); + int num_bits = calc_start_bit(pass - 1) - calc_start_bit(pass); + return (1 << num_bits) - 1; +} + +/** + * Use CUB to twiddle bits - so that we can correctly compare bits of + * floating-point values as well as of integers. + */ +template +__device__ typename cub::Traits::UnsignedBits twiddle_in(T key, bool select_min) { + auto bits = reinterpret_cast::UnsignedBits &>(key); + bits = cub::Traits::TwiddleIn(bits); + if (!select_min) { + bits = ~bits; + } + return bits; +} + +template +__device__ T twiddle_out(typename cub::Traits::UnsignedBits bits, bool select_min) { + if (!select_min) { + bits = ~bits; + } + bits = cub::Traits::TwiddleOut(bits); + return reinterpret_cast(bits); +} + +template +__device__ int calc_bucket(T x, int start_bit, unsigned mask, bool select_min) { + static_assert(BitsPerPass <= sizeof(int) * 8 - 1, + "BitsPerPass is too large that the result type could not be int"); + return (twiddle_in(x, select_min) >> start_bit) & mask; +} + +template +__host__ __device__ IdxT calc_buf_len(IdxT len) { + // When writing is skipped, only read `in`(type T). + // When writing is not skipped, read `in_buf`(T) and `in_idx_buf`(IdxT), and + // write `out_buf`(T) and `out_idx_buf`(IdxT). The ratio between these cases + // determines whether to skip writing and hence the buffer size. constexpr + // float ratio = 2 + sizeof(IdxT) * 2.0 / sizeof(T); + constexpr float ratio = 128; + return len / ratio; + // return len; +} + +/** + * Map a Func over the input data, using vectorized load instructions if + * possible. + * + * NB: in future, we should move this to + * cpp/include/raft/linalg/detail/unary_op.cuh, which currently does not support + * the second lambda argument (index of an element) + * + * @tparam T element type + * @tparam IdxT indexing type + * @tparam Func void (T x, IdxT idx) + * + * @param thread_rank rank of the calling thread among all participating threads + * @param num_threads number of the threads that participate in processing + * @param in the input data + * @param len the number of elements to read + * @param f the lambda taking two arguments (T x, IdxT idx) + */ +template +__device__ void vectorized_process(size_t thread_rank, size_t num_threads, const T *in, idxT len, + Func f) { + if constexpr (sizeof(T) >= sizeof(WideT)) { + for (idxT i = thread_rank; i < len; i += num_threads) { + f(in[i], i); + } + } else { + static_assert(sizeof(WideT) % sizeof(T) == 0); + constexpr int items_per_scalar = sizeof(WideT) / sizeof(T); + // TODO: it's UB + union { + WideT scalar; + T array[items_per_scalar]; // NOLINT(runtime/arrays) + } wide; + + int skip_cnt = + (reinterpret_cast(in) % sizeof(WideT)) + ? ((sizeof(WideT) - reinterpret_cast(in) % sizeof(WideT)) / sizeof(T)) + : 0; + if (skip_cnt > len) { + skip_cnt = len; + } + const WideT *in_cast = reinterpret_cast(in + skip_cnt); + const idxT len_cast = (len - skip_cnt) / items_per_scalar; + + for (idxT i = thread_rank; i < len_cast; i += num_threads) { + wide.scalar = in_cast[i]; + const idxT real_i = skip_cnt + i * items_per_scalar; +#pragma unroll + for (int j = 0; j < items_per_scalar; ++j) { + f(wide.array[j], real_i + j); + } + } + + static_assert(WARP_SIZE >= items_per_scalar); + // and because items_per_scalar > skip_cnt, WARP_SIZE > skip_cnt + // no need to use loop + if (thread_rank < skip_cnt) { + f(in[thread_rank], thread_rank); + } + // because len_cast = (len - skip_cnt) / items_per_scalar, + // len_cast * items_per_scalar + items_per_scalar > len - skip_cnt; + // and so + // len - (skip_cnt + len_cast * items_per_scalar) < items_per_scalar <= + // WARP_SIZE no need to use loop + const idxT remain_i = skip_cnt + len_cast * items_per_scalar + thread_rank; + if (remain_i < len) { + f(in[remain_i], remain_i); + } + } +} + +// sync_width should >= WARP_SIZE +template +__device__ void vectorized_process(const T *in, idxT len, Func f, int sync_width) { + const idxT stride = blockDim.x * gridDim.x; + const idxT tid = blockIdx.x * blockDim.x + threadIdx.x; + if constexpr (sizeof(T) >= sizeof(WideT)) { + for (idxT i = tid; i < len; i += stride) { + f(in[i], i, true); + } + } else { + static_assert(sizeof(WideT) % sizeof(T) == 0); + constexpr int items_per_scalar = sizeof(WideT) / sizeof(T); + union { + WideT scalar; + T array[items_per_scalar]; // NOLINT(runtime/arrays) + } wide; + + int skip_cnt = + (reinterpret_cast(in) % sizeof(WideT)) + ? ((sizeof(WideT) - reinterpret_cast(in) % sizeof(WideT)) / sizeof(T)) + : 0; + if (skip_cnt > len) { + skip_cnt = len; + } + const WideT *in_cast = reinterpret_cast(in + skip_cnt); + const idxT len_cast = (len - skip_cnt) / items_per_scalar; + + const idxT len_cast_for_sync = ((len_cast - 1) / sync_width + 1) * sync_width; + for (idxT i = tid; i < len_cast_for_sync; i += stride) { + bool valid = i < len_cast; + if (valid) { + wide.scalar = in_cast[i]; + } + const idxT real_i = skip_cnt + i * items_per_scalar; +#pragma unroll + for (int j = 0; j < items_per_scalar; ++j) { + f(wide.array[j], real_i + j, valid); + } + } + + static_assert(WARP_SIZE >= items_per_scalar); + // need at most one warp for skipped and remained elements, + // and sync_width >= WARP_SIZE + if (tid < sync_width) { + bool valid = tid < skip_cnt; + T value = valid ? in[tid] : T(); + f(value, tid, valid); + + const idxT remain_i = skip_cnt + len_cast * items_per_scalar + tid; + valid = remain_i < len; + value = valid ? in[remain_i] : T(); + f(value, remain_i, valid); + } + } +} + +template +struct alignas(128) Counter { + // We are processing the values in multiple passes, from most significant to + // least significant. In each pass, we keep the length of input (`len`) and + // the `k` of current pass, and update them at the end of the pass. + IdxT k; + IdxT len; + + // `previous_len` is the length of input in previous pass. Note that + // `previous_len` rather than `len` is used for the filtering step because + // filtering is indeed for previous pass (see comments before + // `radix_kernel`). + IdxT previous_len; + + // We determine the bits of the k_th value inside the mask processed by the + // pass. The already known bits are stored in `kth_value_bits`. It's used to + // discriminate a element is a result (written to `out`), a candidate for next + // pass (written to `out_buf`), or not useful (discarded). The bits that are + // not yet processed do not matter for this purpose. + typename cub::Traits::UnsignedBits kth_value_bits; + + // Record how many elements have passed filtering. It's used to determine the + // position in the `out_buf` where an element should be written. + alignas(128) IdxT filter_cnt; + + // For a row inside a batch, we may launch multiple thread blocks. This + // counter is used to determine if the current block is the last running + // block. If so, this block will execute scan() and choose_bucket(). + alignas(128) unsigned int finished_block_cnt; + + // Record how many elements have been written to the front of `out`. Elements + // less (if select_min==true) than the k-th value are written from front to + // back. + alignas(128) IdxT out_cnt; + + // Record how many elements have been written to the back of `out`. Elements + // equal to the k-th value are written from back to front. We need to keep + // count of them separately because the number of elements that <= the k-th + // value might exceed k. + alignas(128) IdxT out_back_cnt; +}; + +/** + * Fused filtering of the current pass and building histogram for the next pass + * (see steps 4 & 1 in `radix_kernel` description). + */ +template +__device__ void filter_and_histogram(const T *in_buf, const IdxT *in_idx_buf, T *out_buf, + IdxT *out_idx_buf, T *out, IdxT *out_idx, IdxT previous_len, + Counter *counter, IdxT *histogram, bool select_min, + int pass, bool early_stop) { + constexpr int num_buckets = calc_num_buckets(); + __shared__ IdxT histogram_smem[num_buckets]; + for (IdxT i = threadIdx.x; i < num_buckets; i += blockDim.x) { + histogram_smem[i] = 0; + } + __syncthreads(); + + const int start_bit = calc_start_bit(pass); + const unsigned mask = calc_mask(pass); + + if (pass == 0) { + // Passed to vectorized_process, this function executes in all blocks in + // parallel, i.e. the work is split along the input (both, in batches and + // chunks of a single row). Later, the histograms are merged using + // atomicAdd. + auto f = [select_min, start_bit, mask](T value, IdxT) { + int bucket = calc_bucket(value, start_bit, mask, select_min); + atomicAdd(histogram_smem + bucket, static_cast(1)); + }; + vectorized_process(static_cast(blockIdx.x) * blockDim.x + threadIdx.x, + static_cast(blockDim.x) * gridDim.x, in_buf, previous_len, f); + } else { + IdxT *p_filter_cnt = &counter->filter_cnt; + IdxT *p_out_cnt = &counter->out_cnt; + const auto kth_value_bits = counter->kth_value_bits; + const int previous_start_bit = calc_start_bit(pass - 1); + + // See the remark above on the distributed execution of `f` using + // vectorized_process. + auto f = [in_idx_buf, out_buf, out_idx_buf, out, out_idx, select_min, start_bit, mask, + previous_start_bit, kth_value_bits, p_filter_cnt, p_out_cnt, + early_stop](T value, IdxT i) { + const auto previous_bits = (twiddle_in(value, select_min) >> previous_start_bit) + << previous_start_bit; + if (previous_bits == kth_value_bits) { + if (early_stop) { + IdxT pos = atomicAdd(p_out_cnt, static_cast(1)); + if constexpr (store_out) { + out[pos] = value; + } + out_idx[pos] = in_idx_buf ? in_idx_buf[i] : i; + } else { + if (out_buf) { + IdxT pos = atomicAdd(p_filter_cnt, static_cast(1)); + out_buf[pos] = value; + out_idx_buf[pos] = in_idx_buf ? in_idx_buf[i] : i; + } + + int bucket = calc_bucket(value, start_bit, mask, select_min); + atomicAdd(histogram_smem + bucket, static_cast(1)); + } + } + // the condition `(out_buf || early_stop)` is a little tricky: + // If we skip writing to `out_buf` (when `out_buf` is nullptr), we should + // skip writing to `out` too. So we won't write the same value to `out` + // multiple times in different passes. And if we keep skipping the + // writing, values will be written in `last_filter_kernel()` at last. But + // when `early_stop` is true, we need to write to `out` since it's the + // last chance. + else if ((out_buf || early_stop) && previous_bits < kth_value_bits) { // NOLINT + IdxT pos = atomicAdd(p_out_cnt, static_cast(1)); + if constexpr (store_out) { + out[pos] = value; + } + out_idx[pos] = in_idx_buf ? in_idx_buf[i] : i; + } + }; + vectorized_process(static_cast(blockIdx.x) * blockDim.x + threadIdx.x, + static_cast(blockDim.x) * gridDim.x, in_buf, previous_len, f); + } + if (early_stop) { + return; + } + __syncthreads(); + + // merge histograms produced by individual blocks + for (int i = threadIdx.x; i < num_buckets; i += blockDim.x) { + if (histogram_smem[i] != 0) { + atomicAdd(histogram + i, histogram_smem[i]); + } + } +} + +/** + * Replace histogram with its own prefix sum + * (step 2 in `radix_kernel` description) + */ +template +__device__ void scan(volatile IdxT *histogram) { + constexpr int num_buckets = calc_num_buckets(); + if constexpr (num_buckets >= BlockSize) { + static_assert(num_buckets % BlockSize == 0); + constexpr int items_per_thread = num_buckets / BlockSize; + typedef cub::BlockLoad BlockLoad; + typedef cub::BlockStore + BlockStore; + typedef cub::BlockScan BlockScan; + + __shared__ union { + typename BlockLoad::TempStorage load; + typename BlockScan::TempStorage scan; + typename BlockStore::TempStorage store; + } temp_storage; + IdxT thread_data[items_per_thread]; // NOLINT(runtime/arrays) + + BlockLoad(temp_storage.load).Load(histogram, thread_data); + __syncthreads(); + + BlockScan(temp_storage.scan).InclusiveSum(thread_data, thread_data); + __syncthreads(); + + BlockStore(temp_storage.store).Store(histogram, thread_data); + } else { + typedef cub::BlockScan BlockScan; + __shared__ typename BlockScan::TempStorage temp_storage; + + IdxT thread_data = 0; + if (threadIdx.x < num_buckets) { + thread_data = histogram[threadIdx.x]; + } + + BlockScan(temp_storage).InclusiveSum(thread_data, thread_data); + __syncthreads(); + + if (threadIdx.x < num_buckets) { + histogram[threadIdx.x] = thread_data; + } + } +} + +/** + * Calculate in which bucket the k-th value will fall + * (steps 3 in `radix_kernel` description) + */ +template +__device__ void choose_bucket(Counter *counter, const IdxT *histogram, const IdxT k, + const int pass) { + constexpr int num_buckets = calc_num_buckets(); + for (int i = threadIdx.x; i < num_buckets; i += blockDim.x) { + IdxT prev = (i == 0) ? 0 : histogram[i - 1]; + IdxT cur = histogram[i]; + + // one and only one thread will satisfy this condition, so counter is + // written by only one thread + if (prev < k && cur >= k) { + counter->k = k - prev; // how many values still are there to find + counter->len = cur - prev; // number of values in next pass + typename cub::Traits::UnsignedBits bucket = i; + int start_bit = calc_start_bit(pass); + counter->kth_value_bits |= bucket << start_bit; + } + } +} + +template +__device__ void scan_warp_version(cg::thread_block_tile const &warp, + volatile IdxT *histogram, Counter *counter, const IdxT k, + const int pass) { + constexpr int num_buckets = calc_num_buckets(); + + __shared__ IdxT warp_histogram[num_buckets >> WARP_BITS]; + for (int i = threadIdx.x; i < num_buckets; i += BlockSize) { + IdxT data = histogram[i]; + IdxT warp_sum = cg::reduce(warp, data, cg::plus()); + + if (i % WARP_SIZE == 0) { + warp_histogram[i >> WARP_BITS] = warp_sum; + } + } + __syncthreads(); + + if (threadIdx.x < WARP_SIZE) { + IdxT value = warp_histogram[threadIdx.x * 2] + warp_histogram[threadIdx.x * 2 + 1]; + IdxT prefix = value; + for (int offset = 1; offset < WARP_SIZE; offset <<= 1) { + IdxT n = __shfl_up_sync(FULL_WARP_MASK, prefix, offset, WARP_SIZE); + if (threadIdx.x >= offset) prefix += n; + } + IdxT prefix_high = __shfl_sync(FULL_WARP_MASK, prefix, threadIdx.x - 1, WARP_SIZE); + if (threadIdx.x == 0) prefix_high = 0; + warp_histogram[threadIdx.x * 2] += prefix_high; + warp_histogram[threadIdx.x * 2 + 1] = value + prefix_high; + __syncwarp(); + + // Find the target warp bucket + IdxT target_warp = 2048; // invalid value + // bool is_one_in_warp=false; + for (int i = threadIdx.x; i < 64 && target_warp == 2048; i += WARP_SIZE) { + IdxT prev = (i == 0) ? 0 : warp_histogram[i - 1]; + IdxT cur = warp_histogram[i]; + bool is_selected = prev < k && cur >= k; + unsigned mask = __ballot_sync(FULL_WARP_MASK, is_selected); + if (__popc(mask) > 0) { + // target_warp = __ffs(mask) -1 + (i/WARP_SIZE)*WARP_SIZE; + target_warp = __ffs(mask) - 1 + ((i >> WARP_BITS) << WARP_BITS); + // is_one_in_warp= (target_warp==0? warp_histogram[0]: + // warp_histogram[target_warp]-warp_histogram[target_warp-1])==1?true:false; + } + } + + // Find the target bucket + // if(is_one_in_warp){ + // bool is_one=histogram[target_warp*WARP_SIZE+threadIdx.x]==1?1:0; + // unsigned mask = __ballot_sync(FULL_WARP_MASK, is_one); + // IdxT target_bucket=__ffs(mask)-1+target_warp*WARP_SIZE; + // IdxT prev=target_warp==0? 0: warp_histogram[target_warp-1]; + // IdxT cur=warp_histogram[target_warp]; + // if(threadIdx.x==0) { + // counter->k = k - prev; // how many values still are there + // to find counter->len = cur - prev; // number of values in next + // pass typename cub::Traits::UnsignedBits bucket = + // target_bucket; int start_bit = calc_start_bit(pass); counter->kth_value_bits |= bucket << + // start_bit; + // } + // }else{ + value = histogram[(target_warp << WARP_BITS) + threadIdx.x]; + for (int offset = 1; offset < WARP_SIZE; offset <<= 1) { + IdxT n = __shfl_up_sync(FULL_WARP_MASK, value, offset, WARP_SIZE); + if (threadIdx.x >= offset) value += n; + } + value += (target_warp == 0 ? 0 : warp_histogram[target_warp - 1]); + + for (int i = threadIdx.x; i < WARP_SIZE; i += WARP_SIZE) { + IdxT prev = __shfl_up_sync(FULL_WARP_MASK, value, 1, WARP_SIZE); + prev = (i == 0) ? (target_warp == 0 ? 0 : warp_histogram[target_warp - 1]) : prev; + IdxT cur = value; + if (prev < k && cur >= k) { + counter->k = k - prev; // how many values still are there to find + counter->len = cur - prev; // number of values in next pass + typename cub::Traits::UnsignedBits bucket = (target_warp << WARP_BITS) + i; + int start_bit = calc_start_bit(pass); + counter->kth_value_bits |= bucket << start_bit; + } + } + // } + } +} +// For one-block version, last_filter() could be called when pass < num_passes +// - 1. So `pass` could not be constexpr +template +__device__ void last_filter(const T *in_buf, const IdxT *in_idx_buf, T *out, IdxT *out_idx, + IdxT current_len, IdxT k, Counter *counter, + const bool select_min, const int pass) { + const auto kth_value_bits = counter->kth_value_bits; + const int start_bit = calc_start_bit(pass); + + // changed in choose_bucket(); need to reload + const IdxT needed_num_of_kth = counter->k; + IdxT *p_out_cnt = &counter->out_cnt; + IdxT *p_out_back_cnt = &counter->out_back_cnt; + for (IdxT i = threadIdx.x; i < current_len; i += blockDim.x) { + const T value = in_buf[i]; + const auto bits = (twiddle_in(value, select_min) >> start_bit) << start_bit; + if (bits < kth_value_bits) { + IdxT pos = atomicAdd(p_out_cnt, static_cast(1)); + if constexpr (store_out) { + out[pos] = value; + } + // For one-block version, `in_idx_buf` could be nullptr at pass 0. + // For non one-block version, if writing has been skipped, `in_idx_buf` + // could be nullptr if `in_buf` is `in` + out_idx[pos] = in_idx_buf ? in_idx_buf[i] : i; + } else if (bits == kth_value_bits) { + IdxT back_pos = atomicAdd(p_out_back_cnt, static_cast(1)); + if (back_pos < needed_num_of_kth) { + IdxT pos = k - 1 - back_pos; + if constexpr (store_out) { + out[pos] = value; + } + out_idx[pos] = in_idx_buf ? in_idx_buf[i] : i; + } + } + } +} + +template +__global__ void last_filter_kernel(const T *in, const IdxT *in_idx, const T *in_buf, + const IdxT *in_idx_buf, T *out, IdxT *out_idx, IdxT len, IdxT k, + Counter *counters, const bool select_min) { + const size_t batch_id = blockIdx.y; // size_t to avoid multiplication overflow + + Counter *counter = counters + batch_id; + IdxT previous_len = counter->previous_len; + if (previous_len == 0) { + return; + } + const IdxT buf_len = calc_buf_len(len); + if (previous_len > buf_len || in_buf == in) { + in_buf = in + batch_id * len; + in_idx_buf = in_idx ? (in_idx + batch_id * len) : nullptr; + previous_len = len; + } else { + in_buf += batch_id * buf_len; + in_idx_buf += batch_id * buf_len; + } + if constexpr (store_out) { + out += batch_id * k; + } + out_idx += batch_id * k; + + constexpr int pass = calc_num_passes() - 1; + constexpr int start_bit = calc_start_bit(pass); + + const auto kth_value_bits = counter->kth_value_bits; + const IdxT needed_num_of_kth = counter->k; + IdxT *p_out_cnt = &counter->out_cnt; + IdxT *p_out_back_cnt = &counter->out_back_cnt; + + auto f = [k, select_min, kth_value_bits, needed_num_of_kth, p_out_cnt, p_out_back_cnt, in_idx_buf, + out, out_idx](T value, IdxT i) { + const auto bits = (twiddle_in(value, select_min) >> start_bit) << start_bit; + if (bits < kth_value_bits) { + IdxT pos = atomicAdd(p_out_cnt, static_cast(1)); + if constexpr (store_out) { + out[pos] = value; + } + out_idx[pos] = in_idx_buf ? in_idx_buf[i] : i; + } else if (bits == kth_value_bits) { + IdxT back_pos = atomicAdd(p_out_back_cnt, static_cast(1)); + if (back_pos < needed_num_of_kth) { + IdxT pos = k - 1 - back_pos; + if constexpr (store_out) { + out[pos] = value; + } + out_idx[pos] = in_idx_buf ? in_idx_buf[i] : i; + } + } + }; + + vectorized_process(static_cast(blockIdx.x) * blockDim.x + threadIdx.x, + static_cast(blockDim.x) * gridDim.x, in_buf, previous_len, f); +} + +/** + * + * It is expected to call this kernel multiple times (passes), in each pass we + * process a radix, going from the most significant towards the least + * significant bits (MSD). + * + * Conceptually, each pass consists of 4 steps: + * + * 1. Calculate histogram + * First, transform bits into a digit, the value of which is in the range + * [0, 2^{BITS_PER_PASS}-1]. Then count the frequency of each digit value + * and the result is a histogram. That is, histogram[i] contains the count of + * inputs having value i. + * + * 2. Scan the histogram + * Inclusive prefix sum is computed for the histogram. After this step, + * histogram[i] contains the count of inputs having value <= i. + * + * 3. Find the bucket j of the histogram that the k-th value falls into + * + * 4. Filtering + * Input elements whose digit value +__device__ void radix_kernel_func(const T *in, const IdxT *in_idx, const T *in_buf, + const IdxT *in_idx_buf, T *out_buf, IdxT *out_idx_buf, T *out, + IdxT *out_idx, Counter *counter, IdxT *histogram, + const IdxT len, const IdxT k, const bool select_min, + const int pass) { + if (len <= k) { + if (pass == 0) { + for (int index = threadIdx.x; index < len; index += BlockSize) { + if constexpr (store_out) { + out[index] = in[index]; + } + out_idx[index] = in_idx ? in_idx[index] : index; + } + for (int index = threadIdx.x + len; index < k; index += BlockSize) { + if constexpr (store_out) { + out[index] = static_cast(-1.0f); + } + out_idx[index] = -1; + } + return; + } else { + return; + } + } + + IdxT current_k; + IdxT previous_len; + IdxT current_len; + if (pass == 0) { + current_k = k; + previous_len = len; + // Need to do this so setting counter->previous_len for the next pass is + // correct. This value is meaningless for pass 0, but it's fine because pass + // 0 won't be the last pass in this implementation so pass 0 won't hit the + // "if (pass == num_passes - 1)" branch. Maybe it's better to reload + // counter->previous_len and use it rather than current_len in last_filter() + current_len = len; + } else { + current_k = counter->k; + current_len = counter->len; + previous_len = counter->previous_len; + } + if (current_len == 0) { + return; + } + + // When k=len, early_stop will be true at pass 0. It means + // filter_and_histogram() should handle correctly the case that pass=0 and + // early_stop=true. However, this special case of k=len is handled in other + // way in select_k() so such case is not possible here. + const bool early_stop = (current_len == current_k); + const IdxT buf_len = calc_buf_len(len); + constexpr int num_buckets = calc_num_buckets(); + // "previous_len > buf_len" means previous pass skips writing buffer + if (pass == 0 || pass == 1 || previous_len > buf_len) { + in_buf = in; + in_idx_buf = in_idx ? in_idx : nullptr; + previous_len = len; + } + // "current_len > buf_len" means current pass will skip writing buffer + if (pass == 0 || current_len > buf_len) { + out_buf = nullptr; + out_idx_buf = nullptr; + } + + filter_and_histogram(in_buf, in_idx_buf, out_buf, out_idx_buf, + out, out_idx, previous_len, counter, + histogram, select_min, pass, early_stop); + __threadfence(); + + bool isLastBlock = false; + if (threadIdx.x == 0) { + unsigned int finished = atomicInc(&counter->finished_block_cnt, gridDim.x - 1); + isLastBlock = (finished == (gridDim.x - 1)); + } + + if (__syncthreads_or(isLastBlock)) { + if (early_stop) { + if (threadIdx.x == 0) { + // `last_filter_kernel()` requires setting previous_len + counter->previous_len = 0; + counter->len = 0; + } + return; + } + + constexpr int num_passes = calc_num_passes(); + + scan(histogram); + __syncthreads(); + choose_bucket(counter, histogram, current_k, pass); + __syncthreads(); + + // reset for next pass + if (pass != num_passes - 1) { + for (int i = threadIdx.x; i < num_buckets; i += blockDim.x) { + histogram[i] = 0; + } + } + if (threadIdx.x == 0) { + // `last_filter_kernel()` requires setting previous_len even in the last + // pass + counter->previous_len = current_len; + // not necessary for the last pass, but put it here anyway + counter->filter_cnt = 0; + } + + if constexpr (fused_last_filter) { + if (pass == num_passes - 1) { + last_filter( + out_buf ? out_buf : in_buf, out_idx_buf ? out_idx_buf : in_idx_buf, out, out_idx, + out_buf ? current_len : len, k, counter, select_min, pass); + } + } + } +} + +template +__global__ void radix_kernel(const T *in, const IdxT *in_idx, const T *in_buf, + const IdxT *in_idx_buf, T *out_buf, IdxT *out_idx_buf, T *out, + IdxT *out_idx, Counter *counters, IdxT *histograms, + const IdxT len, const IdxT k, const bool select_min, const int pass, + IdxT *lengths) { + const size_t batch_id = blockIdx.y; + auto counter = counters + batch_id; + constexpr int num_buckets = calc_num_buckets(); + auto histogram = histograms + batch_id * num_buckets; + + in += batch_id * len; + if (in_idx) { + in_idx += batch_id * len; + } + if constexpr (store_out) { + out += batch_id * k; + } + out_idx += batch_id * k; + + const IdxT buf_len = calc_buf_len(len); + in_buf += batch_id * buf_len; + in_idx_buf += batch_id * buf_len; + + out_buf += batch_id * buf_len; + out_idx_buf += batch_id * buf_len; + + IdxT actual_len = len; + if (lengths != nullptr) { + actual_len = lengths[batch_id]; + } + radix_kernel_func( + in, in_idx, in_buf, in_idx_buf, out_buf, out_idx_buf, out, out_idx, counter, histogram, + actual_len, k, select_min, pass); +} + +template +unsigned calc_grid_dim(int batch_size, IdxT len, int sm_cnt) { + static_assert(VECTORIZED_READ_SIZE / sizeof(T) >= 1); + + int active_blocks; + cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &active_blocks, radix_kernel, BlockSize, 0); + active_blocks *= sm_cnt; + + IdxT best_num_blocks = 0; + float best_tail_wave_penalty = 1.0f; + const IdxT max_num_blocks = ceildiv(len, VECTORIZED_READ_SIZE / sizeof(T) * BlockSize); + for (int num_waves = 1;; ++num_waves) { + IdxT num_blocks = std::min( + max_num_blocks, static_cast(std::max(num_waves * active_blocks / batch_size, 1))); + IdxT items_per_thread = ceildiv(len, num_blocks * BlockSize); + items_per_thread = alignTo(items_per_thread, VECTORIZED_READ_SIZE / sizeof(T)); + num_blocks = ceildiv(len, items_per_thread * BlockSize); + float actual_num_waves = static_cast(num_blocks) * batch_size / active_blocks; + float tail_wave_penalty = + (ceilf(actual_num_waves) - actual_num_waves) / ceilf(actual_num_waves); + + // 0.15 is determined experimentally. It also ensures breaking the loop + // early, e.g. when num_waves > 7, tail_wave_penalty will always <0.15 + if (tail_wave_penalty < 0.15) { + best_num_blocks = num_blocks; + break; + } else if (tail_wave_penalty < best_tail_wave_penalty) { + best_num_blocks = num_blocks; + best_tail_wave_penalty = tail_wave_penalty; + } + + if (num_blocks == max_num_blocks) { + break; + } + } + return best_num_blocks; +} + +template +__host__ __device__ void set_buf_pointers(const T *in, const IdxT *in_idx, T *buf1, IdxT *idx_buf1, + T *buf2, IdxT *idx_buf2, int pass, const T *&in_buf, + const IdxT *&in_idx_buf, T *&out_buf, + IdxT *&out_idx_buf) { + if (pass == 0) { + in_buf = in; + in_idx_buf = nullptr; + out_buf = nullptr; + out_idx_buf = nullptr; + } else if (pass == 1) { + in_buf = in; + in_idx_buf = in_idx; + out_buf = buf1; + out_idx_buf = idx_buf1; + } else if (pass % 2 == 0) { + in_buf = buf1; + in_idx_buf = idx_buf1; + out_buf = buf2; + out_idx_buf = idx_buf2; + } else { + in_buf = buf2; + in_idx_buf = idx_buf2; + out_buf = buf1; + out_idx_buf = idx_buf1; + } +} + +// The following a few functions are for the one-block version, which uses +// single thread block for each row of a batch. +template +__device__ void filter_and_histogram_for_one_block(const T *in_buf, const IdxT *in_idx_buf, + T *out_buf, IdxT *out_idx_buf, T *out, + IdxT *out_idx, Counter *counter, + IdxT *histogram, bool select_min, int pass) { + constexpr int num_buckets = calc_num_buckets(); + for (int i = threadIdx.x; i < num_buckets; i += blockDim.x) { + histogram[i] = 0; + } + IdxT *p_filter_cnt = &counter->filter_cnt; + if (threadIdx.x == 0) { + *p_filter_cnt = 0; + } + __syncthreads(); + + const int start_bit = calc_start_bit(pass); + const unsigned mask = calc_mask(pass); + const IdxT previous_len = counter->previous_len; + + if (pass == 0) { + if constexpr (is_vectorized) { + auto f = [histogram, select_min, start_bit, mask](T value, IdxT) { + int bucket = calc_bucket(value, start_bit, mask, select_min); + atomicAdd(histogram + bucket, static_cast(1)); + }; + vectorized_process(threadIdx.x, blockDim.x, in_buf, previous_len, f); + } else { + for (IdxT i = threadIdx.x; i < previous_len; i += blockDim.x) { + const T value = in_buf[i]; + int bucket = calc_bucket(value, start_bit, mask, select_min); + atomicAdd(histogram + bucket, static_cast(1)); + } + } + } else { + // not use vectorized_process here because it increases #registers a lot + IdxT *p_out_cnt = &counter->out_cnt; + const auto kth_value_bits = counter->kth_value_bits; + const int previous_start_bit = calc_start_bit(pass - 1); + + for (IdxT i = threadIdx.x; i < previous_len; i += blockDim.x) { + const T value = in_buf[i]; + const auto previous_bits = (twiddle_in(value, select_min) >> previous_start_bit) + << previous_start_bit; + if (previous_bits == kth_value_bits) { +#if CUDART_VERSION < 12000 + // Avoiding potential compiler bug in CUDA 11 + volatile +#endif + IdxT pos = atomicAdd(p_filter_cnt, static_cast(1)); + out_buf[pos] = value; + out_idx_buf[pos] = in_idx_buf ? in_idx_buf[i] : i; + + int bucket = calc_bucket(value, start_bit, mask, select_min); + atomicAdd(histogram + bucket, static_cast(1)); + } else if (previous_bits < kth_value_bits) { + IdxT pos = atomicAdd(p_out_cnt, static_cast(1)); + if constexpr (store_out) { + out[pos] = value; + } + out_idx[pos] = in_idx_buf ? in_idx_buf[i] : i; + } + } + } +} + +template +__device__ void radix_topk_one_block_func(const T *in, const IdxT *in_idx, const IdxT len, + const IdxT k, T *out, IdxT *out_idx, + const bool select_min, T *buf1, IdxT *idx_buf1, T *buf2, + IdxT *idx_buf2) { + if (len <= k) { + for (int index = threadIdx.x; index < len; index += BlockSize) { + if constexpr (store_out) { + out[index] = in[index]; + } + out_idx[index] = in_idx ? in_idx[index] : index; + } + for (int index = threadIdx.x + len; index < k; index += BlockSize) { + if constexpr (store_out) { + out[index] = static_cast(-1.0f); + } + out_idx[index] = -1; + } + return; + } + + constexpr int num_buckets = calc_num_buckets(); + __shared__ Counter counter; + __shared__ IdxT histogram[num_buckets]; + + if (threadIdx.x == 0) { + counter.k = k; + counter.len = len; + counter.previous_len = len; + counter.kth_value_bits = 0; + counter.out_cnt = 0; + counter.out_back_cnt = 0; + } + __syncthreads(); + + // const size_t batch_id = blockIdx.x; // size_t to avoid multiplication + // overflow + const T *in_buf = nullptr; + const IdxT *in_idx_buf = nullptr; + T *out_buf = nullptr; + IdxT *out_idx_buf = nullptr; + + constexpr int num_passes = calc_num_passes(); + auto block = cg::this_thread_block(); + auto warp = cg::tiled_partition(block); + for (int pass = 0; pass < num_passes; ++pass) { + set_buf_pointers(in, in_idx, buf1, idx_buf1, buf2, idx_buf2, pass, in_buf, in_idx_buf, out_buf, + out_idx_buf); + + IdxT current_len = counter.len; + IdxT current_k = counter.k; + + filter_and_histogram_for_one_block( + in_buf, in_idx_buf, out_buf, out_idx_buf, out, out_idx, &counter, histogram, select_min, + pass); + __syncthreads(); + + scan(histogram); + __syncthreads(); + + choose_bucket(&counter, histogram, current_k, pass); + // scan_warp_version( + // warp, histogram, &counter, current_k, pass); + if (threadIdx.x == 0) { + counter.previous_len = current_len; + } + __syncthreads(); + + if (counter.len == counter.k || pass == num_passes - 1) { + last_filter(pass == 0 ? in : out_buf, + pass == 0 ? in_idx : out_idx_buf, out, out_idx, + current_len, k, &counter, select_min, pass); + break; + } + } // end for pass +} // end kernel + +template +__global__ void radix_topk_one_block_kernel(const T *in, const IdxT *in_idx, const IdxT len, + const IdxT k, T *out, IdxT *out_idx, + const bool select_min, T *buf1, IdxT *idx_buf1, T *buf2, + IdxT *idx_buf2, IdxT *lengths) { + const size_t batch_id = blockIdx.x; // size_t to avoid multiplication overflow + IdxT actual_len = len; + if (lengths) { + actual_len = lengths[batch_id]; + } + + in += batch_id * len; + if (in_idx) { + in_idx += batch_id * len; + } + + out += batch_id * k; + out_idx += batch_id * k; + buf1 += batch_id * len; + idx_buf1 += batch_id * len; + buf2 += batch_id * len; + idx_buf2 += batch_id * len; + + radix_topk_one_block_func( + in, in_idx, actual_len, k, out, out_idx, select_min, buf1, idx_buf1, buf2, idx_buf2); +} // end kernel + +} // namespace topk + +/***************Runtime API****************/ + +template +void standalone_radix_topk_(void *buf, size_t &buf_size, const T *in, const IdxT *in_idx, + int batch_size, IdxT len, IdxT k, T *out, IdxT *out_idx, + bool select_min, bool fused_last_filter, unsigned grid_dim, + cudaStream_t stream, IdxT *lengths = nullptr) { + static_assert(topk::calc_num_passes() > 1); + constexpr int num_buckets = topk::calc_num_buckets(); + + topk::Counter *counters = nullptr; + IdxT *histograms = nullptr; + T *buf1 = nullptr; + IdxT *idx_buf1 = nullptr; + T *buf2 = nullptr; + IdxT *idx_buf2 = nullptr; + { + IdxT len_candidates = topk::calc_buf_len(len); + std::vector sizes = {sizeof(*counters) * batch_size, + sizeof(*histograms) * num_buckets * batch_size, + sizeof(*buf1) * len_candidates * batch_size, + sizeof(*idx_buf1) * len_candidates * batch_size, + sizeof(*buf2) * len_candidates * batch_size, + sizeof(*idx_buf2) * len_candidates * batch_size}; + size_t total_size = calc_aligned_size(sizes); + if (!buf) { + buf_size = total_size; + return; + } + + std::vector aligned_pointers = calc_aligned_pointers(buf, sizes); + counters = static_cast(aligned_pointers[0]); + histograms = static_cast(aligned_pointers[1]); + buf1 = static_cast(aligned_pointers[2]); + idx_buf1 = static_cast(aligned_pointers[3]); + buf2 = static_cast(aligned_pointers[4]); + idx_buf2 = static_cast(aligned_pointers[5]); + + cudaMemsetAsync( + buf, 0, static_cast(aligned_pointers[2]) - static_cast(aligned_pointers[0]), + stream); + } + + const T *in_buf = nullptr; + const IdxT *in_idx_buf = nullptr; + T *out_buf = nullptr; + IdxT *out_idx_buf = nullptr; + + dim3 blocks(grid_dim, batch_size); + + constexpr int num_passes = topk::calc_num_passes(); + + auto kernel = topk::radix_kernel; + + for (int pass = 0; pass < num_passes; ++pass) { + topk::set_buf_pointers(in, in_idx, buf1, idx_buf1, buf2, idx_buf2, pass, in_buf, in_idx_buf, + out_buf, out_idx_buf); + + if (fused_last_filter && pass == num_passes - 1 && out != nullptr) { + kernel = topk::radix_kernel; + } else if (fused_last_filter && pass == num_passes - 1 && out == nullptr) { + kernel = topk::radix_kernel; + } else if (out == nullptr) { + kernel = topk::radix_kernel; + } + + kernel<<>>(in, in_idx, in_buf, in_idx_buf, out_buf, out_idx_buf, + out, out_idx, counters, histograms, len, k, select_min, + pass, lengths); + } + + if (!fused_last_filter) { + if (out != nullptr) { + topk::last_filter_kernel<<>>( + in, in_idx, out_buf, out_idx_buf, out, out_idx, len, k, counters, select_min); + } else { + topk::last_filter_kernel<<>>( + in, in_idx, out_buf, out_idx_buf, out, out_idx, len, k, counters, select_min); + } + } +} + +template +void standalone_radix_topk_one_block_(void *buf, size_t &buf_size, const T *in, const IdxT *in_idx, + int batch_size, IdxT len, IdxT k, T *out, IdxT *out_idx, + bool select_min, cudaStream_t stream, + IdxT *lengths = nullptr) { + static_assert(topk::calc_num_passes() > 1); + + T *buf1 = nullptr; + IdxT *idx_buf1 = nullptr; + T *buf2 = nullptr; + IdxT *idx_buf2 = nullptr; + { + std::vector sizes = { + sizeof(*buf1) * len * batch_size, sizeof(*idx_buf1) * len * batch_size, + sizeof(*buf2) * len * batch_size, sizeof(*idx_buf2) * len * batch_size}; + size_t total_size = calc_aligned_size(sizes); + if (!buf) { + buf_size = total_size; + return; + } + + std::vector aligned_pointers = calc_aligned_pointers(buf, sizes); + buf1 = static_cast(aligned_pointers[0]); + idx_buf1 = static_cast(aligned_pointers[1]); + buf2 = static_cast(aligned_pointers[2]); + idx_buf2 = static_cast(aligned_pointers[3]); + } + + if (out != nullptr) { + topk::radix_topk_one_block_kernel + <<>>(in, in_idx, len, k, out, out_idx, select_min, buf1, + idx_buf1, buf2, idx_buf2, lengths); + } else { + topk::radix_topk_one_block_kernel + <<>>(in, in_idx, len, k, out, out_idx, select_min, buf1, + idx_buf1, buf2, idx_buf2, lengths); + } +} + +template +void standalone_topk(void *buf, size_t &buf_size, const T *in, int batch_size, idxT len, idxT k, + T *out, idxT *out_idx, bool greater, cudaStream_t stream = 0, + idxT *lengths = nullptr, bool is_prefill = false) { + constexpr int items_per_thread = 32; + constexpr int multi_block_dim = 256; + constexpr int single_block_dim = 1024; + constexpr bool fused_last_filter = false; + if (len <= single_block_dim * items_per_thread || is_prefill) { + standalone_radix_topk_one_block_( + buf, buf_size, in, static_cast(nullptr), batch_size, len, k, out, out_idx, !greater, + stream, lengths); + } else { + // Cache sm_cnt per device to avoid repeated host-side queries. + static int cached_dev = -1; + static int cached_sm_cnt = -1; + int sm_cnt; + { + int dev; + NVTE_CHECK_CUDA(cudaGetDevice(&dev)); + if (dev != cached_dev) { + NVTE_CHECK_CUDA( + cudaDeviceGetAttribute(&cached_sm_cnt, cudaDevAttrMultiProcessorCount, dev)); + cached_dev = dev; + } + sm_cnt = cached_sm_cnt; + } + unsigned grid_dim = topk::calc_grid_dim(batch_size, len, sm_cnt); + + if (grid_dim == 1) { + standalone_radix_topk_one_block_( + buf, buf_size, in, static_cast(nullptr), batch_size, len, k, out, out_idx, + !greater, stream, lengths); + } else { + standalone_radix_topk_( + buf, buf_size, in, static_cast(nullptr), batch_size, len, k, out, out_idx, + !greater, fused_last_filter, grid_dim, stream, lengths); + } + } +} +} // namespace nv diff --git a/transformer_engine/common/util/topk.cu b/transformer_engine/common/util/topk.cu new file mode 100644 index 000000000..21018a494 --- /dev/null +++ b/transformer_engine/common/util/topk.cu @@ -0,0 +1,60 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include + +#include "../common.h" +#include "standalone_topk.cuh" + +void nvte_topk(cudaStream_t stream, const NVTETensor keys_in, const NVTETensor lengths_in, + NVTETensor keys_out, NVTETensor indices_out, NVTETensor workspace, int batch_size, + int seq_len, int k) { + NVTE_API_CALL(nvte_topk); + using namespace transformer_engine; + + Tensor *workspace_tensor = convertNVTETensor(workspace); + + if (workspace_tensor->data.numel() == 0) { + size_t workspace_bytes = 0; + nv::standalone_topk(nullptr, workspace_bytes, nullptr, batch_size, seq_len, k, + nullptr, nullptr, /*greater=*/true, /*stream=*/nullptr, + /*lengths=*/nullptr, /*is_prefill=*/false); + workspace_tensor->data.shape = {workspace_bytes}; + workspace_tensor->data.dtype = DType::kByte; + return; + } + + const Tensor *keys_in_tensor = convertNVTETensorCheck(keys_in); + const Tensor *lengths_tensor = convertNVTETensorCheck(lengths_in); + Tensor *keys_out_tensor = convertNVTETensor(keys_out); + Tensor *indices_tensor = convertNVTETensor(indices_out); + + void *d_workspace = workspace_tensor->data.dptr; + size_t workspace_bytes = workspace_tensor->data.numel(); + const int *d_lengths = reinterpret_cast(lengths_tensor->data.dptr); + int *d_indices = reinterpret_cast(indices_tensor->data.dptr); + + auto dtype = keys_in_tensor->data.dtype; + +#define DISPATCH_TOPK(T) \ + do { \ + const T *d_in = reinterpret_cast(keys_in_tensor->data.dptr); \ + T *d_out = reinterpret_cast(keys_out_tensor->data.dptr); \ + nv::standalone_topk(d_workspace, workspace_bytes, d_in, batch_size, seq_len, k, d_out, \ + d_indices, /*greater=*/true, stream, const_cast(d_lengths), \ + /*is_prefill=*/false); \ + } while (0) + + if (dtype == DType::kBFloat16) { + DISPATCH_TOPK(__nv_bfloat16); + } else if (dtype == DType::kFloat32) { + DISPATCH_TOPK(float); + } else { + NVTE_ERROR("nvte_topk: unsupported key dtype (supported: float32, bfloat16)"); + } + +#undef DISPATCH_TOPK +} diff --git a/transformer_engine/jax/cpp_extensions/__init__.py b/transformer_engine/jax/cpp_extensions/__init__.py index d203fcea9..fe1f93dc7 100644 --- a/transformer_engine/jax/cpp_extensions/__init__.py +++ b/transformer_engine/jax/cpp_extensions/__init__.py @@ -10,3 +10,4 @@ from .softmax import * from .gemm import * from .router import * +from .topk import * diff --git a/transformer_engine/jax/cpp_extensions/topk.py b/transformer_engine/jax/cpp_extensions/topk.py new file mode 100644 index 000000000..b8e3a92f3 --- /dev/null +++ b/transformer_engine/jax/cpp_extensions/topk.py @@ -0,0 +1,136 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""TopK custom op""" + +import functools +from typing import Tuple + +import jax +import jax.numpy as jnp +from jax import dtypes, ffi + +from .base import BasePrimitive, register_primitive +from .misc import te_dtype_to_jax_dtype + +__all__ = ["topk"] + + +@functools.lru_cache(maxsize=512) +def get_topk_workspace_sizes(batch_size: int, seq_len: int, k: int): + """Query the workspace shape and dtype required for TopK. + + The result is memoised per (batch_size, seq_len, k) tuple so that repeated + JIT compilations with the same shapes incur only one host-side CUDA call. + """ + import transformer_engine_jax as _te_jax + + (wkspace_info,) = _te_jax.get_topk_workspace_sizes(batch_size, seq_len, k) + return wkspace_info + + +class TopKPrimitive(BasePrimitive): + """ + TopK Primitive + + Selects the top-k entries (by value) from each row of a 2-D input using the + AIR radix-selection algorithm. Returns both the top-k key values and their + column indices within each row. + """ + + name = "te_topk_ffi" + multiple_results = True + impl_static_args = (2,) # k_value + inner_primitive = None + outer_primitive = None + + @staticmethod + def abstract( + in_keys_aval, + in_lengths_aval, + *, + k_value, + ): + keys_dtype = dtypes.canonicalize_dtype(in_keys_aval.dtype) + assert keys_dtype in [ + jnp.float32, + jnp.bfloat16, + ], f"topk: unsupported key dtype {keys_dtype}; supported: float32, bfloat16" + assert in_keys_aval.ndim == 2, "topk: keys input must be 2D (batch_size, seq_len)" + assert dtypes.canonicalize_dtype(in_lengths_aval.dtype) == jnp.int32 + + batch_size, seq_len = in_keys_aval.shape + wkspace_info = get_topk_workspace_sizes(batch_size, seq_len, k_value) + + out_shape = (batch_size, k_value) + out_keys_aval = jax.core.ShapedArray(shape=out_shape, dtype=keys_dtype) + out_indices_aval = jax.core.ShapedArray(shape=out_shape, dtype=jnp.int32) + workspace_aval = jax.core.ShapedArray( + shape=wkspace_info[0], dtype=te_dtype_to_jax_dtype(wkspace_info[1]) + ) + return (out_keys_aval, out_indices_aval, workspace_aval) + + @staticmethod + def outer_abstract(*args, **kwargs): + out_keys_aval, out_indices_aval, _workspace_aval = TopKPrimitive.abstract(*args, **kwargs) + return (out_keys_aval, out_indices_aval) + + @staticmethod + def lowering(ctx, in_keys, in_lengths, k_value): + return ffi.ffi_lowering(TopKPrimitive.name)( + ctx, + in_keys, + in_lengths, + k_value=k_value, + ) + + @staticmethod + def impl(in_keys, in_lengths, k_value): + assert TopKPrimitive.inner_primitive is not None + out_keys, out_indices, _workspace = TopKPrimitive.inner_primitive.bind( + in_keys, + in_lengths, + k_value=k_value, + ) + return (out_keys, out_indices) + + +register_primitive(TopKPrimitive) + + +def topk( + x: jnp.ndarray, + k_value: int, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Select the top-k largest entries from each row using the AIR radix algorithm. + + Args: + x: Input array of shape ``(batch_size, seq_len)`` or ``(seq_len,)``. + Supported dtypes: ``float32``, ``bfloat16``. + k_value: Number of top entries to select per row. + + Returns: + A tuple ``(values, indices)`` where both arrays have shape + ``(batch_size, k_value)`` (or ``(k_value,)`` for 1-D input). The + outputs are *unordered*: use ``jax.lax.sort_key_val`` if a sorted result + is required. ``indices`` are the column positions within the original row. + """ + squeezed = x.ndim == 1 + if squeezed: + x = x[jnp.newaxis, :] # (1, seq_len) + + assert x.ndim == 2, f"topk expected 2D input tensor 'x' but {x.shape=}" + batch_size, seq_len = x.shape + lengths = jnp.full((batch_size,), seq_len, dtype=jnp.int32) + + out_keys, out_indices = TopKPrimitive.outer_primitive.bind( + x, + lengths, + k_value=k_value, + ) + + if squeezed: + out_keys = out_keys[0] # (k_value,) + out_indices = out_indices[0] # (k_value,) + + return out_keys, out_indices diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 3ba0e7e9b..2ecfedc8a 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -199,6 +199,10 @@ XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedTopkWithScoreFunctionBackwardHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedMoEAuxLossForwardHandler); XLA_FFI_DECLARE_HANDLER_SYMBOL(FusedMoEAuxLossBackwardHandler); +// TopK +XLA_FFI_DECLARE_HANDLER_SYMBOL(TopkHandler); +pybind11::tuple GetTopkWorkspaceSizes(int batch_size, int seq_len, int k); + } // namespace jax } // namespace transformer_engine diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index e3bc12240..b00264394 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -101,6 +101,9 @@ pybind11::dict Registrations() { dict["te_fused_moe_aux_loss_forward_ffi"] = EncapsulateFFI(FusedMoEAuxLossForwardHandler); dict["te_fused_moe_aux_loss_backward_ffi"] = EncapsulateFFI(FusedMoEAuxLossBackwardHandler); + // TopK + dict["te_topk_ffi"] = EncapsulateFFI(TopkHandler); + return dict; } @@ -118,6 +121,7 @@ PYBIND11_MODULE(transformer_engine_jax, m) { m.def("get_norm_bwd_workspace_sizes", &GetNormBackwardWorkspaceSizes); m.def("get_fused_attn_fwd_workspace_sizes", &GetFusedAttnForwardWorkspaceSizes); m.def("get_fused_attn_bwd_workspace_sizes", &GetFusedAttnBackwardWorkspaceSizes); + m.def("get_topk_workspace_sizes", &GetTopkWorkspaceSizes); m.def("nvte_get_qkv_format", &nvte_get_qkv_format); m.def("is_non_nt_fp8_gemm_supported", &nvte_is_non_tn_fp8_gemm_supported); m.def("initialize_cgemm_communicator", &InitializeCgemmCommunicator); diff --git a/transformer_engine/jax/csrc/extensions/topk.cpp b/transformer_engine/jax/csrc/extensions/topk.cpp new file mode 100644 index 000000000..450ff08b3 --- /dev/null +++ b/transformer_engine/jax/csrc/extensions/topk.cpp @@ -0,0 +1,104 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "transformer_engine/topk.h" + +#include "../extensions.h" +#include "xla/ffi/api/c_api.h" + +namespace transformer_engine { +namespace jax { + +// --------------------------------------------------------------------------- +// JAX FFI handler +// --------------------------------------------------------------------------- + +Error_Type TopkFFI(cudaStream_t stream, Buffer_Type keys_in_buf, Buffer_Type lengths_buf, + Result_Type keys_out_buf, Result_Type indices_out_buf, Result_Type workspace_buf, + int64_t k_value) { + auto keys_in_dtype = convert_ffi_datatype_to_te_dtype(keys_in_buf.element_type()); + auto keys_out_dtype = convert_ffi_datatype_to_te_dtype(keys_out_buf->element_type()); + auto idx_out_dtype = convert_ffi_datatype_to_te_dtype(indices_out_buf->element_type()); + NVTE_CHECK(keys_in_dtype == keys_out_dtype, "TopkFFI: input and output key dtypes must match"); + NVTE_CHECK(idx_out_dtype == DType::kInt32, "TopkFFI: index output must be int32"); + + auto keys_in_shape = keys_in_buf.dimensions(); + NVTE_CHECK(keys_in_shape.size() == 2, "TopkFFI: keys input must be 2D (batch_size, seq_len)"); + + int batch_size = static_cast(keys_in_shape[0]); + int seq_len = static_cast(keys_in_shape[1]); + int k = static_cast(k_value); + + // Validate key dtype (float32 and bfloat16 only). + switch (keys_in_dtype) { + case DType::kFloat32: + case DType::kBFloat16: + break; + default: + NVTE_ERROR("TopkFFI: unsupported key dtype (float32 and bfloat16 only)"); + } + + auto workbuf_bytes = product(workspace_buf->dimensions()); + + // Build flat TensorWrappers over the full (batch_size * seq_len) / (batch_size * k) buffers. + auto flat_in_shape = + std::vector{static_cast(batch_size) * static_cast(seq_len)}; + auto flat_out_shape = + std::vector{static_cast(batch_size) * static_cast(k)}; + auto len_shape = std::vector{static_cast(batch_size)}; + auto ws_shape = std::vector{workbuf_bytes}; + + auto keys_in_tensor = TensorWrapper(keys_in_buf.untyped_data(), flat_in_shape, keys_in_dtype); + auto lengths_tensor = TensorWrapper(lengths_buf.untyped_data(), len_shape, DType::kInt32); + auto keys_out_tensor = + TensorWrapper(keys_out_buf->untyped_data(), flat_out_shape, keys_out_dtype); + auto idx_out_tensor = + TensorWrapper(indices_out_buf->untyped_data(), flat_out_shape, DType::kInt32); + auto workspace_tensor = TensorWrapper(workspace_buf->untyped_data(), ws_shape, DType::kByte); + + nvte_topk(stream, keys_in_tensor.data(), lengths_tensor.data(), keys_out_tensor.data(), + idx_out_tensor.data(), workspace_tensor.data(), batch_size, seq_len, k); + + return ffi_with_cuda_error_check(); +} + +XLA_FFI_DEFINE_HANDLER_SYMBOL(TopkHandler, TopkFFI, + FFI::Bind() + .Ctx() // stream + .Arg() // keys_in + .Arg() // lengths + .Ret() // keys_out + .Ret() // indices_out + .Ret() // workspace + .Attr("k_value"), + FFI_CudaGraph_Traits); + +// --------------------------------------------------------------------------- +// Workspace-size query exposed to Python +// --------------------------------------------------------------------------- + +pybind11::tuple GetTopkWorkspaceSizes(int batch_size, int seq_len, int k) { + auto flat_in_shape = + std::vector{static_cast(batch_size) * static_cast(seq_len)}; + auto flat_out_shape = + std::vector{static_cast(batch_size) * static_cast(k)}; + auto len_shape = std::vector{static_cast(batch_size)}; + + auto keys_in_tensor = TensorWrapper(nullptr, flat_in_shape, DType::kFloat32); + auto lengths_tensor = TensorWrapper(nullptr, len_shape, DType::kInt32); + auto keys_out_tensor = TensorWrapper(nullptr, flat_out_shape, DType::kFloat32); + auto idx_out_tensor = TensorWrapper(nullptr, flat_out_shape, DType::kInt32); + TensorWrapper workspace_tensor; + + nvte_topk(nullptr, keys_in_tensor.data(), lengths_tensor.data(), keys_out_tensor.data(), + idx_out_tensor.data(), workspace_tensor.data(), batch_size, seq_len, k); + + auto work_shape = MakeShapeVector(workspace_tensor.shape()); + return pybind11::make_tuple(std::make_pair(work_shape, workspace_tensor.dtype())); +} + +} // namespace jax +} // namespace transformer_engine From 5d947a03775797875c90ce1c3cf249d0d3dd33cb Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Thu, 23 Apr 2026 20:44:35 -0700 Subject: [PATCH 016/180] Fix the race in the dbias computation in MXFP8 quantization and grouped quantization kernel (#2921) Fix the race in the dbias computation Signed-off-by: Przemek Tredak --- transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh | 2 ++ transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh | 2 ++ 2 files changed, 4 insertions(+) diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index ce827d24e..aa697d4bf 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -713,6 +713,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) group_quantize_mxfp8_kernel if constexpr (COLWISE_SCALING) { thread_partial_dbias = partial_dbias_colwise; } else { + ptx::cp_async_bulk_wait_group_read<0>(); + __syncthreads(); float *partial_dbias_rowwise = reinterpret_cast(dshmem); constexpr size_t DBIAS_BUFF_WIDTH = THREADS_X * (SCALE_DIM_X + 1); diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index f36b07108..a0ae7dde8 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -498,6 +498,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) if constexpr (COLWISE_SCALING) { thread_partial_dbias = partial_dbias_colwise; } else { + ptx::cp_async_bulk_wait_group_read<0>(); + __syncthreads(); // Reusing dshmem (in_sh) as dbias buffer [HEIGHT x WIDTH] // HEIGHT = THREADS_Y // WIDTH = THREADS_X * (SCALE_DIM_X + 1) From 9ad2e7bc9dde40cc4546eb6879475754f142bc59 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Fri, 24 Apr 2026 16:07:09 -0700 Subject: [PATCH 017/180] Remove uncessary ctype being passed to GroupedGEMMQuant kernel (#2922) * remove ctype to eliminate memory usage from the cudnn kernel Signed-off-by: Varun Thumbe * Remove c_dtype from fusible ops test Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Varun Thumbe Signed-off-by: Kirthi Shankar Sivamani Co-authored-by: Kirthi Shankar Sivamani --- tests/pytorch/test_fusible_ops.py | 1 - transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py | 1 - transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py | 1 - 3 files changed, 3 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 0f40e9218..c73f56056 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -4658,7 +4658,6 @@ def test_grouped_gemm_quant_cute_matches_mxfp8_quantized() -> None: norm_const_tensor=None, prob_tensor=inputs["prob_tensor"], acc_dtype=torch.float32, - c_dtype=torch.bfloat16, d_dtype=torch.bfloat16, cd_major="n", sf_vec_size=32, diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index aca49e986..29273a5b4 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -687,7 +687,6 @@ def fuser_backward( "norm_const_tensor": None, "prob_tensor": torch.ones((out_shape[0], 1, 1), dtype=torch.float32, device=device), "acc_dtype": torch.float32, - "c_dtype": dtype, "d_dtype": dtype, "cd_major": "n", "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index 90c4204f0..cad31e2c5 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -436,7 +436,6 @@ def fuser_forward( "norm_const_tensor": None, "prob_tensor": fc2_scales_tensor, "acc_dtype": torch.float32, - "c_dtype": dtype, "d_dtype": dtype, "cd_major": "n", "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, From f2e31dbb604cac5d045384e455dac09b37687868 Mon Sep 17 00:00:00 2001 From: Muu Date: Tue, 28 Apr 2026 04:25:17 +0800 Subject: [PATCH 018/180] fix: TransformerEngineBaseModule quantizers init values type (#2927) Signed-off-by: Muu --- transformer_engine/pytorch/module/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index ebfb98b2d..e6bedee0c 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -776,7 +776,7 @@ def __init__(self, name: Optional[str] = None) -> None: self.fp8_meta["fp8_checkpoint"] = False self.fp8_meta["fp8_group"] = None self.fp8_meta_tensors_initialized = False - self.quantizers = {"scaling_fwd": {}, "scaling_bwd": {}} + self.quantizers = {"scaling_fwd": [], "scaling_bwd": []} self.tp_group = None self.tp_size = 1 self.sequence_parallel = False From 82ace626294c7da1cbe8d0d8e6f4b5af627635c8 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Mon, 27 Apr 2026 18:55:30 -0700 Subject: [PATCH 019/180] [Common] Fix "0" literal for compilation (#2934) --- transformer_engine/common/fused_attn/flash_attn.cu | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/common/fused_attn/flash_attn.cu b/transformer_engine/common/fused_attn/flash_attn.cu index 5037be828..38bf09f81 100644 --- a/transformer_engine/common/fused_attn/flash_attn.cu +++ b/transformer_engine/common/fused_attn/flash_attn.cu @@ -311,7 +311,7 @@ __device__ __forceinline__ void permute_vec_loop(const T *__restrict__ in, T *__ const size_t s_local = w / pad_elems; const size_t s_i = s_begin + s_local; const size_t d_off = D + (w % pad_elems); - out[out_base + s_i * D_out + d_off] = static_cast(0); + out[out_base + s_i * D_out + d_off] = static_cast(0.f); } } } From df0025b646c26f9059bc3c43bd5fb39863c00289 Mon Sep 17 00:00:00 2001 From: Jacket <44538064+kainzhong@users.noreply.github.com> Date: Tue, 28 Apr 2026 10:11:07 -0700 Subject: [PATCH 020/180] [Common, PyTorch] Add triton mHC kernels & pytorch APIs (#2790) * [Common, PyTorch] Add triton mHC kernels & pytorch operators Signed-off-by: Kaining Zhong * fix Signed-off-by: Kaining Zhong * nit Signed-off-by: Kaining Zhong * make linter happy Signed-off-by: Kaining Zhong * nit Signed-off-by: Kaining Zhong * ah OK Signed-off-by: Kaining Zhong * new configs to improve perf Signed-off-by: Kaining Zhong * add APIs to docs Signed-off-by: Kaining Zhong * fix typos, check deterministic, refactor Signed-off-by: Kaining Zhong * fix Signed-off-by: Kaining Zhong * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * reset rng for all tests Signed-off-by: Kaining Zhong * add docstring Signed-off-by: Kaining Zhong * fix api doc Signed-off-by: Kaining Zhong * whoops Signed-off-by: Kaining Zhong * grad_x doesn't have to zero Signed-off-by: Kaining Zhong * nit Signed-off-by: Kaining Zhong * nit Signed-off-by: Kaining Zhong * force pytorch to not use bf16 for reduction Signed-off-by: Kaining Zhong * use TE's general_gemm instead Signed-off-by: Kaining Zhong * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Looks like this is how to make TE use fp32 acc Signed-off-by: Kaining Zhong * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kaining Zhong Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/api/pytorch.rst | 10 + qa/L0_pytorch_lint/test.sh | 0 qa/L0_pytorch_unittest/test.sh | 2 + tests/pytorch/test_mhc.py | 497 +++++ transformer_engine/common/triton/mhc.py | 1693 +++++++++++++++++ transformer_engine/pytorch/triton/__init__.py | 1 + transformer_engine/pytorch/triton/mhc.py | 999 ++++++++++ 7 files changed, 3202 insertions(+) mode change 100644 => 100755 qa/L0_pytorch_lint/test.sh create mode 100644 tests/pytorch/test_mhc.py create mode 100644 transformer_engine/common/triton/mhc.py create mode 100644 transformer_engine/pytorch/triton/mhc.py diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index 3217d29c3..db8649800 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -229,6 +229,16 @@ Operation fuser .. autoapiclass:: transformer_engine.pytorch.ops.SwiGLU +.. autoapifunction:: transformer_engine.pytorch.triton.mhc.mhc_fused_sinkhorn + +.. autoapifunction:: transformer_engine.pytorch.triton.mhc.mhc_fused_scale + +.. autoapifunction:: transformer_engine.pytorch.triton.mhc.mhc_fused_aggregate + +.. autoapifunction:: transformer_engine.pytorch.triton.mhc.mhc_fused_expand_combine + +.. autoapifunction:: transformer_engine.pytorch.triton.mhc.mhc_fused_projection + Deprecated functions -------------------- diff --git a/qa/L0_pytorch_lint/test.sh b/qa/L0_pytorch_lint/test.sh old mode 100644 new mode 100755 diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 377c9ddb0..a8f8cf875 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -58,6 +58,8 @@ fi python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_checkpoint.xml $TE_PATH/tests/pytorch/test_checkpoint.py || test_fail "test_checkpoint.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_router.xml $TE_PATH/tests/pytorch/test_fused_router.py || test_fail "test_fused_router.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_partial_cast.xml $TE_PATH/tests/pytorch/test_partial_cast.py || test_fail "test_partial_cast.py" +# Disable autotuning to make unittests faster. In addition, disable TF32 path to fully align with the pytorch reference implementation's precision +NVTE_DISABLE_TRITON_AUTOTUNING=1 NVIDIA_TF32_OVERRIDE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_mhc.xml $TE_PATH/tests/pytorch/test_mhc.py || test_fail "test_mhc.py" if [ "$RET" -ne 0 ]; then echo "Error in the following test cases:$FAILED_CASES" diff --git a/tests/pytorch/test_mhc.py b/tests/pytorch/test_mhc.py new file mode 100644 index 000000000..541ce9a8c --- /dev/null +++ b/tests/pytorch/test_mhc.py @@ -0,0 +1,497 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +from dataclasses import dataclass +import pytest +import torch +import torch.nn.functional as F + +from utils import reset_rng_states +from transformer_engine.pytorch.triton.mhc import ( + mhc_fused_sinkhorn, + mhc_fused_scale, + mhc_fused_aggregate, + mhc_fused_expand_combine, + mhc_fused_projection, +) + +# Disable TF32 for matmul to ensure consistency between the fused and reference implementations +torch.backends.cuda.matmul.allow_tf32 = False + + +def mhc_projection_ref(x, phi): + """ + Reference operator for mHC's projection building operation. + + x: (M, nC) where M = s * b + phi: (2n + n^2, nC), which consists of the following matrices + - phi_pre: (n, nC) + - phi_post: (n, nC) + - phi_res: (n^2, nC) + n: number of Hyper Connection streams + C: hidden dimension per stream + """ + x_dtype = x.dtype + x = x.to(torch.float32) + phi = phi.to(torch.float32) + + Hs = x @ phi.T # (M, 2n + n^2) + + x_fp32 = x.to(torch.float32) # Use fp32 for better numerical stability in variance calculation + ms = (x_fp32 * x_fp32).mean(dim=1) + + return Hs.to(x_dtype), ms + + +def mhc_scale_ref(H, alpha, beta, ms, n): + """ + Reference operator for mHC's H matrices scaling operation + + :param: H: (M, 2n + n^2), the unprocessed H matrices where M = s * b + :param: alpha: (3,), three scalar parameters + :param: beta: (1, 2n + n^2), bias term + :param: r: (M,), the denominator for RMSNorm + :param: n: int, the width of Hyper-Connection + + :return Hs: (M, 2n + n^2), the processed H matrices + """ + + input_dtype = H.dtype + H = H.to(torch.float32) + alpha = alpha.to(torch.float32) + beta = beta.to(torch.float32) + eps = torch.finfo(torch.float32).eps + rms = torch.sqrt(ms + eps) # (M,) + rms = rms.to(torch.float32) + + H_pre = H[:, :n] # (M, n) + H_post = H[:, n : 2 * n] # (M, n) + H_res = H[:, 2 * n :] # (M, n^2) + + beta_pre = beta[0, :n] + beta_post = beta[0, n : 2 * n] + beta_res = beta[0, 2 * n : 2 * n + n * n] + + alpha_pre, alpha_post, alpha_res = alpha[0], alpha[1], alpha[2] + + H_pre = H_pre * alpha_pre + H_post = H_post * alpha_post + H_res = H_res * alpha_res + + H_pre = H_pre / rms[:, None] + H_post = H_post / rms[:, None] + H_res = H_res / rms[:, None] + + H_pre = H_pre + beta_pre + H_post = H_post + beta_post + H_res = H_res + beta_res + + H_pre = F.sigmoid(H_pre) + H_post = 2 * F.sigmoid(H_post) + + return H_pre.to(input_dtype), H_post.to(input_dtype), H_res.to(input_dtype) + + +def mhc_sinkhorn_ref(H_res, n=4, iterations=20): + """ + Reference operator for mHC's Sinkhorn-Knopp algorithm to convert a matrix into a doubly stochastic matrix. + Calculated in log space for numerical stability. + + :param H_res: a tensor of shape (s, b, n, n) + :return: a tensor of shape (s, b, n, n) + """ + s, b = H_res.shape[:2] + device = H_res.device + dtype = H_res.dtype + + H_res_f = H_res.to( + torch.float32 + ).clone() # Use float32 for better numerical stability during Sinkhorn iterations + + log_mu = torch.zeros(s, b, n, device=device, dtype=torch.float32) + log_nu = torch.zeros(s, b, n, device=device, dtype=torch.float32) + + f = torch.zeros(s, b, n, device=device, dtype=torch.float32) + g = torch.zeros(s, b, n, device=device, dtype=torch.float32) + + for _ in range(iterations): + # Update f: logsumexp over the column dimension (3) + f = log_mu - torch.logsumexp(H_res_f + g.unsqueeze(2), dim=3) + # Update g: logsumexp over the row dimension (2) + g = log_nu - torch.logsumexp(H_res_f + f.unsqueeze(3), dim=2) + + log_P = f.unsqueeze(3) + H_res_f + g.unsqueeze(2) + H_res_out = torch.exp(log_P).to(dtype) # Convert back to original dtype + + return H_res_out + + +def mhc_aggregate_ref(x, H_pre, n): + """ + Reference operator for applying mHC's aggregation transformation + + x: (s, b, C, n) + H_pre: (s, b, n) + """ + H_pre = H_pre.contiguous() + + s, b, C, n = x.shape + H_pre = H_pre.view(s, b, n, 1) + + out = (x @ H_pre).view(s, b, C) + + return out + + +def mhc_expand_combine_ref(f, bias, H_post, x, H_res, n): + """ + Reference operator for applying mHC's expansion and combination transformation + + f: (s, b, C) + bias: (C,) or None + H_post: (s, b, n) + x: (s, b, C, n) + H_res: (s, b, n, n) + """ + + s, b, C, n = x.shape + + # My triton kernels use FMA and MMA instructions with fp32 accumulator for bf16 test cases + # which has better numerical stability than this pytorch implementation + # To match the kernel's accuracy we need to cast to fp32 here to match kernels' result + input_dtype = f.dtype + f = f.to(torch.float32) + bias = bias.to(torch.float32) if bias is not None else None + H_post = H_post.to(torch.float32) + x = x.to(torch.float32) + H_res = H_res.to(torch.float32) + + if bias is not None: + f = f + bias[None, None, :] + + f = f.view(s, b, C, 1) + H_post = H_post.view(s, b, 1, n) + + out = f @ H_post + x @ H_res # (s, b, C, n) + + return out.to(input_dtype) + + +@dataclass +class MHCConfig: + s: int = 2048 # Sequence length + b: int = 32 # Batch size + C: int = 1024 # Hidden dimension + n: int = 4 # Number of Hyper Connection streams + + allow_n = [ + 4, + ] + + def __init__(self, b, s, C, n=4): + assert n in self.allow_n, f"n must be one of {self.allow_n}" + self.b = b + self.s = s + self.C = C + self.n = n + + @staticmethod + def desc(cfg): + return f"b{cfg.b}_s{cfg.s}_C{cfg.C}_n{cfg.n}" + + +mhc_configs = [ + MHCConfig(8, 32, 32), + MHCConfig(8, 128, 16 * 64), + MHCConfig( + 4, + 128, + 16 * 64, + ), + MHCConfig(2, 2048, 24 * 128), + MHCConfig( + 1, + 2048, + 24 * 128, + ), + MHCConfig( + 13, + 1, + 16 * 128, + ), + MHCConfig( + 7, + 1, + 16 * 256, + ), + MHCConfig( + 8, + 1, + 16 * 192, + ), + MHCConfig( + 8, + 128, + 5129, + ), + MHCConfig( + 8, + 512, + 8000, + ), + MHCConfig( + 4, + 1024, + 8192, + ), + MHCConfig( + 2, + 4096, + 8192, + ), + MHCConfig( + 8, + 128, + 16384, + ), +] + + +def get_tols(dtype): + if dtype == torch.bfloat16: + tols = dict(atol=2.5e-2, rtol=2.5e-2) + else: + tols = dict(atol=5e-3, rtol=5e-3) + return tols + + +@pytest.mark.parametrize("cfg", mhc_configs, ids=MHCConfig.desc) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) +def test_mhc_projection(cfg: MHCConfig, dtype): + reset_rng_states() + + s, b, C, n = cfg.s, cfg.b, cfg.C, cfg.n + nC = n * C + N = 2 * n + n * n + + tols = get_tols(dtype) + use_tf32 = False + + x = torch.randn(s * b, nC, device="cuda", requires_grad=True, dtype=dtype) + phi = torch.randn(N, nC, dtype=dtype, requires_grad=True, device="cuda") + + x_ref = x.detach().clone().requires_grad_(True) + phi_ref = phi.detach().clone().requires_grad_(True) + + ref_out_Hs, ref_out_ms = mhc_projection_ref(x_ref, phi_ref) + fused_out_Hs_padded, fused_out_ms = mhc_fused_projection(x, phi, use_tf32) + fused_out_Hs = fused_out_Hs_padded[:, :N] + + torch.testing.assert_close(fused_out_Hs, ref_out_Hs, **tols) + torch.testing.assert_close(fused_out_ms, ref_out_ms, **tols) + (ref_out_Hs.sum() + ref_out_ms.sum()).backward() + (fused_out_Hs.sum() + fused_out_ms.sum()).backward() + + torch.testing.assert_close(x.grad, x_ref.grad, **tols) + torch.testing.assert_close(phi.grad, phi_ref.grad, **tols) + + +@pytest.mark.parametrize("cfg", mhc_configs, ids=MHCConfig.desc) +@pytest.mark.parametrize("dtype", [torch.float32], ids=["fp32"]) +def test_mhc_scale(cfg: MHCConfig, dtype): + reset_rng_states() + + s, b, C, n = cfg.s, cfg.b, cfg.C, cfg.n + N = 2 * n + n * n + + tols = get_tols(dtype) + + H_padded = torch.randn(s * b, 32, device="cuda", requires_grad=True, dtype=dtype) + H = H_padded[:, :N] + alpha = torch.randn(3, device="cuda", requires_grad=True, dtype=dtype) + beta = torch.randn(1, 2 * n + n * n, device="cuda", requires_grad=True, dtype=dtype) + ms_raw = torch.randn(s * b, device="cuda", dtype=dtype).abs() + 1.0 + ms = ms_raw.detach().clone().requires_grad_(True) + + H_ref = H.detach().clone().requires_grad_(True) + alpha_ref = alpha.detach().clone().requires_grad_(True) + beta_ref = beta.detach().clone().requires_grad_(True) + ms_ref = ms.detach().clone().requires_grad_(True) + + ref_out = mhc_scale_ref(H_ref[:, :N], alpha_ref, beta_ref, ms_ref, n) + fused_out = mhc_fused_scale(H_padded, alpha, beta, ms, n) + + for i in range(3): + torch.testing.assert_close(fused_out[i], ref_out[i], **tols) + + torch.cat([ref_out[i] for i in range(3)], dim=-1).sum().backward() + torch.cat([fused_out[i] for i in range(3)], dim=-1).sum().backward() + + torch.testing.assert_close(H_padded.grad[:, :N], H_ref.grad, **tols) + torch.testing.assert_close(alpha.grad, alpha_ref.grad, **tols) + torch.testing.assert_close(beta.grad, beta_ref.grad, **tols) + torch.testing.assert_close(ms.grad, ms_ref.grad, **tols) + + +@pytest.mark.parametrize("cfg", mhc_configs, ids=MHCConfig.desc) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) +def test_mhc_combined(cfg: MHCConfig, dtype): + reset_rng_states() + + s, b, C, n = cfg.s, cfg.b, cfg.C, cfg.n + N = 2 * n + n * n + nC = n * C + + tols = get_tols(dtype) + use_tf32 = False + + x = torch.randn(s * b, nC, device="cuda", requires_grad=True, dtype=dtype) + phi = torch.randn(N, nC, dtype=dtype, requires_grad=True, device="cuda") + + alpha = torch.randn(3, device="cuda", requires_grad=True, dtype=dtype) + beta = torch.randn(1, 2 * n + n * n, device="cuda", requires_grad=True, dtype=dtype) + + x_ref = x.detach().clone().requires_grad_(True) + phi_ref = phi.detach().clone().requires_grad_(True) + + alpha_ref = alpha.detach().clone().requires_grad_(True) + beta_ref = beta.detach().clone().requires_grad_(True) + + ref_out_H, ref_out_r = mhc_projection_ref(x_ref, phi_ref) + fused_out_H_padded, fused_out_r = mhc_fused_projection(x, phi, use_tf32) + + ref_H_pre, ref_H_post, ref_H_res = mhc_scale_ref( + ref_out_H[:, :N], alpha_ref, beta_ref, ref_out_r, n + ) + fused_H_pre, fused_H_post, fused_H_res = mhc_fused_scale( + fused_out_H_padded, alpha, beta, fused_out_r, n + ) + + def mhc_combined(x_ref, phi_ref, alpha_ref, beta_ref): + dtype = x_ref.dtype + x_ref = x_ref.to(torch.float32) + phi_ref = phi_ref.to(torch.float32) + alpha_ref = alpha_ref.to(torch.float32) + beta_ref = beta_ref.to(torch.float32) + + # Check if after spliting RMSNorm to two steps in projection and scaling, + # theresult is close to applying RMSNorm in the correct order + x_rmsnorm = F.rms_norm(x_ref, normalized_shape=(nC,)) + H = x_rmsnorm @ phi_ref.T + H_pre = H[:, :n] + H_post = H[:, n : 2 * n] + H_res = H[:, 2 * n :] + + out_pre = H_pre * alpha_ref[0] + beta_ref[:, :n] + out_post = H_post * alpha_ref[1] + beta_ref[:, n : 2 * n] + out_res = H_res * alpha_ref[2] + beta_ref[:, 2 * n :] + + out_pre = out_pre.sigmoid() + out_post = 2 * out_post.sigmoid() + out_res = out_res + + return out_pre.to(dtype), out_post.to(dtype), out_res.to(dtype) + + combined_H_pre, combined_H_post, combined_H_res = mhc_combined( + x_ref, phi_ref, alpha_ref, beta_ref + ) + + torch.testing.assert_close(combined_H_pre, ref_H_pre, **tols) + torch.testing.assert_close(combined_H_post, ref_H_post, **tols) + torch.testing.assert_close(combined_H_res, ref_H_res, **tols) + + torch.testing.assert_close(combined_H_pre, fused_H_pre, **tols) + torch.testing.assert_close(combined_H_post, fused_H_post, **tols) + torch.testing.assert_close(combined_H_res, fused_H_res, **tols) + + +@pytest.mark.parametrize("cfg", mhc_configs, ids=MHCConfig.desc) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) +@pytest.mark.parametrize("recompute", [False, True], ids=["no_recompute", "recompute"]) +def test_mhc_sinkhorn(cfg: MHCConfig, dtype, recompute): + reset_rng_states() + + s, b, C, n = cfg.s, cfg.b, cfg.C, cfg.n + + tols = get_tols(dtype) + + x = torch.randn(s, b, n, n, device="cuda", requires_grad=True, dtype=dtype) + x_ref = x.detach().clone().requires_grad_(True) + + ref_out = mhc_sinkhorn_ref(x_ref, n) + fused_out = mhc_fused_sinkhorn(x, n, recompute) + + torch.testing.assert_close(fused_out, ref_out, **tols) + + ref_out.sum().backward() + fused_out.sum().backward() + + torch.testing.assert_close(x.grad, x_ref.grad, **tols) + + +@pytest.mark.parametrize("cfg", mhc_configs, ids=MHCConfig.desc) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) +def test_mhc_aggregate(cfg: MHCConfig, dtype): + reset_rng_states() + + s, b, C, n = cfg.s, cfg.b, cfg.C, cfg.n + + tols = get_tols(dtype) + + x = torch.randn(s, b, C, n, device="cuda", requires_grad=True, dtype=dtype) + H_pre = torch.randn(s, b, n, device="cuda", requires_grad=True, dtype=dtype) + + x_ref = x.detach().clone().requires_grad_(True) + H_pre_ref = H_pre.detach().clone().requires_grad_(True) + + ref_out = mhc_aggregate_ref(x_ref, H_pre_ref, n) + fused_out = mhc_fused_aggregate(x, H_pre, n, False) + + torch.testing.assert_close(fused_out, ref_out, **tols) + + ref_out.sum().backward() + fused_out.sum().backward() + + torch.testing.assert_close(x.grad, x_ref.grad, **tols) + torch.testing.assert_close(H_pre.grad, H_pre_ref.grad, **tols) + + +@pytest.mark.parametrize("cfg", mhc_configs, ids=MHCConfig.desc) +@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=["fp32", "bf16"]) +@pytest.mark.parametrize("with_bias", [True, False], ids=["with_bias", "no_bias"]) +def test_mhc_expand_combine(cfg: MHCConfig, dtype, with_bias): + reset_rng_states() + + s, b, C, n = cfg.s, cfg.b, cfg.C, cfg.n + + tols = get_tols(dtype) + + f = torch.randn(s, b, C, device="cuda", requires_grad=True, dtype=dtype) + bias = None + if with_bias: + bias = torch.randn(C, device="cuda", requires_grad=True, dtype=dtype) + H_post = torch.randn(s, b, n, device="cuda", requires_grad=True, dtype=dtype) + x = torch.randn(s, b, C, n, device="cuda", requires_grad=True, dtype=dtype) + H_res = torch.randn(s, b, n, n, device="cuda", requires_grad=True, dtype=dtype) + + f_ref = f.detach().clone().requires_grad_(True) + bias_ref = None if bias is None else bias.detach().clone().requires_grad_(True) + H_post_ref = H_post.detach().clone().requires_grad_(True) + x_ref = x.detach().clone().requires_grad_(True) + H_res_ref = H_res.detach().clone().requires_grad_(True) + + ref_out = mhc_expand_combine_ref(f_ref, bias_ref, H_post_ref, x_ref, H_res_ref, n) + fused_out = mhc_fused_expand_combine(f, bias, H_post, x, H_res, n, False) + + torch.testing.assert_close(fused_out, ref_out, **tols) + + ref_out.sum().backward() + fused_out.sum().backward() + + torch.testing.assert_close(f.grad, f_ref.grad, **tols) + torch.testing.assert_close(H_post.grad, H_post_ref.grad, **tols) + torch.testing.assert_close(x.grad, x_ref.grad, **tols) + torch.testing.assert_close(H_res.grad, H_res_ref.grad, **tols) + if bias is not None: + torch.testing.assert_close(bias.grad, bias_ref.grad, **tols) diff --git a/transformer_engine/common/triton/mhc.py b/transformer_engine/common/triton/mhc.py new file mode 100644 index 000000000..965bb437f --- /dev/null +++ b/transformer_engine/common/triton/mhc.py @@ -0,0 +1,1693 @@ +# pylint: disable=missing-function-docstring + +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""triton kernels for mHC (manifold Hyper-Connection) operations""" + +import itertools +import os + +import triton +import triton.language as tl + + +def projection_config_fwd(): + block_m = [64, 128] + block_k = [1024] + step_k = [32, 64] + warps = [4] + stages = [3, 4] + + configs = [] + for m, bk, sk, w, s in itertools.product(block_m, block_k, step_k, warps, stages): + configs.append( + triton.Config( + {"BLOCK_SIZE_M": m, "BLOCK_SIZE_K": bk, "STEP_SIZE_K": sk}, + num_warps=w, + num_stages=s, + ) + ) + if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": + configs = configs[:1] + return configs + + +def projection_config_bwd(): + block_m = [32, 128] + block_k = [128] + warps = [2] + stages = [2, 3, 4] + + configs = [] + for m, bk, w, s in itertools.product(block_m, block_k, warps, stages): + configs.append( + triton.Config({"BLOCK_SIZE_M": m, "BLOCK_SIZE_K": bk}, num_warps=w, num_stages=s) + ) + if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": + configs = configs[:1] + return configs + + +@triton.autotune(configs=projection_config_fwd(), key=["M", "K"], reset_to_zero=["h_ptr", "ms_ptr"]) +@triton.jit +def _mhc_projection_fwd_fused( + x_ptr, # (M, K) + phi_ptr, # (N, K) + h_ptr, # (M, 32) + ms_ptr, # (M,) + M, + N, + K, + stride_xm, + stride_xk: tl.constexpr, + stride_phin, + stride_phik: tl.constexpr, + stride_hm: tl.constexpr, + stride_hn: tl.constexpr, + stride_ms: tl.constexpr, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + STEP_SIZE_K: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + precision: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_k = tl.program_id(axis=1) + + tl.assume(pid_m >= 0) + tl.assume(pid_k >= 0) + tl.assume(stride_xm > 0) + tl.assume(stride_xk == 1) + tl.assume(stride_phin == K) + tl.assume(stride_phik == 1) + tl.assume(stride_hm == 32) + tl.assume(stride_hn == 1) + tl.assume(stride_ms == 1) + + tl.assume(BLOCK_SIZE_M % 32 == 0) + tl.assume(BLOCK_SIZE_K % 32 == 0) + tl.assume(BLOCK_SIZE_N == 32) + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n_full = tl.arange(0, BLOCK_SIZE_N) + mask_m = offs_m < M + + h_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + ms_acc = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32) + + k_base = pid_k * BLOCK_SIZE_K + for k_start in range(0, tl.cdiv(BLOCK_SIZE_K, STEP_SIZE_K)): + k_offs = k_base + k_start * STEP_SIZE_K + tl.arange(0, STEP_SIZE_K) + mask_k = k_offs < K + x_ptrs = x_ptr + offs_m[:, None] * stride_xm + k_offs[None, :] * stride_xk + x = tl.load( + x_ptrs, mask=mask_m[:, None] & mask_k[None, :], other=0.0 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_K) + phi_ptrs = phi_ptr + offs_n_full[:, None] * stride_phin + k_offs[None, :] * stride_phik + phi = tl.load( + phi_ptrs, + mask=(offs_n_full[:, None] < N) & mask_k[None, :], + other=0.0, + cache_modifier=".ca", + ) # (BLOCK_SIZE_N, BLOCK_SIZE_K) + ms_acc += tl.sum(x * x, axis=1) + h_acc = tl.dot( + x, tl.trans(phi, (1, 0)), h_acc, input_precision=precision, out_dtype=tl.float32 + ) + + h_ptrs = h_ptr + offs_m[:, None] * stride_hm + offs_n_full[None, :] * stride_hn + tl.atomic_add(h_ptrs, h_acc, mask=mask_m[:, None], sem="relaxed") + + offs_ms = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + masks_ms = offs_ms < M + offs_ms %= M + ms_ptrs = ms_ptr + offs_ms * stride_ms + ms = ms_acc / tl.cast(K, tl.float32) + tl.atomic_add(ms_ptrs, ms, mask=masks_ms, sem="relaxed") + + +@triton.autotune( + configs=projection_config_bwd(), + key=["M", "K"], +) +@triton.jit +def _mhc_projection_bwd_fused( + x_ptr, + grad_x_ptr, # (M, K) + phi_ptr, # (N, K) + grad_h_ptr, # (M, N) + grad_ms_ptr, # (M,) + M, + N, + K, + stride_xm, + stride_xk: tl.constexpr, + stride_grad_xm, + stride_grad_xk: tl.constexpr, + stride_phin, + stride_phik: tl.constexpr, + stride_grad_phin, + stride_grad_phik: tl.constexpr, + stride_grad_hm: tl.constexpr, + stride_grad_hn: tl.constexpr, + stride_grad_ms: tl.constexpr, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + precision: tl.constexpr, +): + pid_m = tl.program_id(axis=0) + pid_k = tl.program_id(axis=1) + + tl.assume(pid_m >= 0) + tl.assume(pid_k >= 0) + tl.assume(stride_xm > 0) + tl.assume(stride_xk == 1) + tl.assume(stride_grad_hm == 32) + tl.assume(stride_grad_hn == 1) + tl.assume(stride_phin == K) + tl.assume(stride_phik == 1) + tl.assume(stride_grad_phin == K) + tl.assume(stride_grad_phik == 1) + tl.assume(stride_grad_ms == 1) + + tl.assume(BLOCK_SIZE_M % 32 == 0) + tl.assume(BLOCK_SIZE_K % 32 == 0) + tl.assume(BLOCK_SIZE_N == 32) + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_k = pid_k * BLOCK_SIZE_K + tl.arange(0, BLOCK_SIZE_K) + offs_n_full = tl.arange(0, BLOCK_SIZE_N) + mask_m = offs_m < M + mask_k = offs_k < K + + x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_k[None, :] * stride_xk + x = tl.load( + x_ptrs, mask=mask_m[:, None] & mask_k[None, :], other=0.0 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_K) + + grad_h_ptrs = ( + grad_h_ptr + offs_m[:, None] * stride_grad_hm + offs_n_full[None, :] * stride_grad_hn + ) + grad_h = tl.load( + grad_h_ptrs, mask=mask_m[:, None] & (offs_n_full[None, :] < N), other=0.0 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_N) + + phi_ptrs = phi_ptr + offs_n_full[:, None] * stride_phin + offs_k[None, :] * stride_phik + offs_ms = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + grad_ms_ptrs = grad_ms_ptr + offs_ms * stride_grad_ms + + phi = tl.load( + phi_ptrs, mask=(offs_n_full[:, None] < N) & mask_k[None, :], other=0.0 + ) # (BLOCK_SIZE_N, BLOCK_SIZE_K) + grad_ms = tl.load( + grad_ms_ptrs, mask=offs_ms < M, other=0.0, cache_modifier=".ca" + ) # (BLOCK_SIZE_M,) + + grad_x = x * (grad_ms * 2 / tl.cast(K, tl.float32))[:, None] + grad_x = tl.dot( + grad_h, phi, acc=grad_x, input_precision=precision, out_dtype=tl.float32 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_K) + grad_x_ptrs = grad_x_ptr + offs_m[:, None] * stride_grad_xm + offs_k[None, :] * stride_grad_xk + grad_x = grad_x.to(x.dtype) + tl.store(grad_x_ptrs, grad_x, mask=mask_m[:, None] & mask_k[None, :]) + + +def scale_config(): + block_m = [128] + warps = [4] + stages = [1, 2, 4] + + configs = [] + for m, w, s in itertools.product(block_m, warps, stages): + configs.append(triton.Config({"BLOCK_SIZE_M": m}, num_warps=w, num_stages=s)) + + if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": + configs = configs[:1] + return configs + + +@triton.autotune( + configs=scale_config(), + key=["M"], +) +@triton.jit +def _mhc_scale_fwd_fused( + h_ptr, # (M, 2n + n^2), which is padded to (M, 32) in the last dimension + a_ptr, # (3,) + b_ptr, # (2n + n^2) + ms_ptr, # (M,) + out_ptr, # (M, 2n + n^2), which is padded to (M, 32) in the last dimension + M, + n, + stride_hm, + stride_hn, + stride_a, + stride_b, + stride_ms, + stride_out_m, + stride_out_n, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + eps: tl.constexpr, +): + pid = tl.program_id(0) + + tl.assume(M > 0) + tl.assume(n == 4) + tl.assume(stride_hm == 32) + tl.assume(stride_hn == 1) + tl.assume(stride_out_m == 32) + tl.assume(stride_out_n == 1) + tl.assume(stride_a == 1) + tl.assume(stride_b == 1) + tl.assume(stride_ms == 1) + tl.assume(BLOCK_SIZE_N == 32) + + N = 2 * n + n * n + + offs_m = pid * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + cols = tl.arange(0, BLOCK_SIZE_N) + mask_m = offs_m < M + + # Expand a to BLOCK_SIZE_N length + offs_a = tl.zeros_like(cols) + offs_a = tl.where((cols >= n) & (cols < 2 * n), 1, offs_a) + offs_a = tl.where((cols >= 2 * n) & (cols < 2 * n + n * n), 2, offs_a) + # Pick a[0] from a for the first 4 columns, a[1] for the next 4 columns, and a[2] for the rest of the columns + a = tl.load( + a_ptr + offs_a * stride_a, mask=offs_a < 3, other=0.0 + ) # a[2*n + n*n:] is filled with garbage + a = tl.where(cols < N, a, 0.0) # Mask out the garbage values in a + + b = tl.load(b_ptr + cols * stride_b, mask=cols < N, other=0.0) # (BLOCK_SIZE_N,) + ms = tl.load(ms_ptr + offs_m * stride_ms, mask=mask_m, other=0.0) # (BLOCK_SIZE_M,) + # In projection kernel we use split-K so we only have the accumulated ms, + # and now we need to take sqrt on the accumulated ms to obtain the RMSNorm denominator. + rms = tl.sqrt(ms + eps) + + h = tl.load( + h_ptr + offs_m[:, None] * stride_hm + cols[None, :] * stride_hn, + mask=mask_m[:, None], + other=0.0, + ) # (BLOCK_SIZE_M, BLOCK_SIZE_N) + + h = a[None, :] * h + h = tl.fma( + h, 1.0 / rms[:, None], b[None, :] + ) # (BLOCK_SIZE_M, BLOCK_SIZE_N), where the first 2n columns are H_pre and H_post, and the rest are H_res + h_sigmoid_pre = tl.sigmoid(h) + h_sigmoid_post = 2 * h_sigmoid_pre + + # Use this mask to select h[:, :2n] + h = tl.where(cols[None, :] < n, h_sigmoid_pre, h) + h = tl.where((cols[None, :] >= n) & (cols[None, :] < 2 * n), h_sigmoid_post, h) + + tl.store( + out_ptr + offs_m[:, None] * stride_out_m + cols[None, :] * stride_out_n, + h, + mask=mask_m[:, None], + ) + + +@triton.autotune( + configs=scale_config(), + key=["M"], + reset_to_zero=["grad_a_ptr", "grad_b_ptr"], +) +@triton.jit +def _mhc_scale_bwd_fused( + grad_out_ptr, + out_ptr, # (M, 2n + n^2), which is padded to (M, 32) in the last dimension + grad_h_ptr, + h_ptr, # (M, 2n + n^2), which is padded to (M, 32) in the last dimension + grad_a_ptr, + a_ptr, # (3,) + grad_b_ptr, # (2n + n^2,) + grad_ms_ptr, + ms_ptr, # (M,) + M, + n, + stride_grad_out_m, + stride_grad_out_n, + stride_out_m, + stride_out_n, + stride_grad_hm, + stride_grad_hn, + stride_hm, + stride_hn, + stride_grad_a, + stride_a, + stride_grad_b, + stride_grad_ms, + stride_ms, + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, + eps: tl.constexpr, +): + pid = tl.program_id(0) + + tl.assume(M > 0) + tl.assume(n == 4) + tl.assume(stride_grad_out_m == 32) + tl.assume(stride_grad_out_n == 1) + tl.assume(stride_out_m == 32) + tl.assume(stride_out_n == 1) + tl.assume(stride_grad_hm == 32) + tl.assume(stride_grad_hn == 1) + tl.assume(stride_hm == 32) + tl.assume(stride_hn == 1) + tl.assume(stride_grad_a == 1) + tl.assume(stride_a == 1) + tl.assume(stride_grad_b == 1) + tl.assume(stride_grad_ms == 1) + tl.assume(stride_ms == 1) + tl.assume(BLOCK_SIZE_N == 32) + + N = 2 * n + n * n + + offs_m = pid * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + cols = tl.arange(0, BLOCK_SIZE_N) + mask_m = offs_m < M + mask_n = cols < N + + # Expand a to BLOCK_SIZE_N length + offs_a = tl.zeros_like(cols) + offs_a = tl.where((cols >= n) & (cols < 2 * n), 1, offs_a) + offs_a = tl.where((cols >= 2 * n) & (cols < 2 * n + n * n), 2, offs_a) + # Pick a[0] from a for the first 4 columns, a[1] for the next 4 columns, and a[2] for the rest of the columns + a = tl.load( + a_ptr + offs_a * stride_a, mask=offs_a < 3, other=0.0 + ) # a[2*n + n*n:] is filled with garbage + a = tl.where(cols < N, a, 0.0) # Mask out the garbage values in a + + ms_offsets = offs_m + ms_mask = mask_m + ms = tl.load(ms_ptr + ms_offsets * stride_ms, mask=ms_mask, other=1.0) # (BLOCK_SIZE_M,) + rms = tl.sqrt(ms + eps) + + grad_out = tl.load( + grad_out_ptr + offs_m[:, None] * stride_grad_out_m + cols[None, :] * stride_grad_out_n, + mask=mask_m[:, None] & mask_n[None, :], + other=0.0, + ) # (BLOCK_SIZE_M, BLOCK_SIZE_N) + out = tl.load( + out_ptr + offs_m[:, None] * stride_out_m + cols[None, :] * stride_out_n, + mask=mask_m[:, None] & mask_n[None, :], + other=0.0, + ) # (BLOCK_SIZE_M, BLOCK_SIZE_N) + h = tl.load( + h_ptr + offs_m[:, None] * stride_hm + cols[None, :] * stride_hn, + mask=mask_m[:, None] & mask_n[None, :], + other=0.0, + ) # (BLOCK_SIZE_M, BLOCK_SIZE_N) + + # Gradiient of H before H_pre and H_post go through sigmoid + grad_out_out = grad_out * out + grad_h_pre = grad_out_out * (1 - out) + grad_h_post = grad_out_out * 0.5 * (2 - out) + grad_h = grad_out + grad_h = tl.where(cols[None, :] < n, grad_h_pre, grad_h) + grad_h = tl.where((cols[None, :] >= n) & (cols[None, :] < 2 * n), grad_h_post, grad_h) + + grad_a = tl.sum(h * grad_h / rms[:, None], axis=0).to(a.dtype) + # Write grad_a[0:4].sum to grad_a_ptr[0], grad_a[4:8].sum to grad_a_ptr[1], and grad_a[8:24].sum to grad_a_ptr[2] + tl.atomic_add(grad_a_ptr, tl.where(cols[None, :] < n, grad_a, 0.0).sum(), sem="relaxed") + tl.atomic_add( + grad_a_ptr + stride_grad_a, + tl.where((cols[None, :] >= n) & (cols[None, :] < 2 * n), grad_a, 0.0).sum(), + sem="relaxed", + ) + tl.atomic_add( + grad_a_ptr + 2 * stride_grad_a, + tl.where((cols[None, :] >= 2 * n) & (cols[None, :] < 2 * n + n * n), grad_a, 0.0).sum(), + sem="relaxed", + ) + + grad_b = tl.sum(grad_h, axis=0).to(a.dtype) + tl.atomic_add(grad_b_ptr + cols * stride_grad_b, grad_b, mask=cols < N, sem="relaxed") + + grad_rms = (tl.sum((-grad_h * h * a[None, :]), axis=1) / (rms * rms)).to(rms.dtype) + grad_ms = grad_rms / (2 * rms) + tl.store(grad_ms_ptr + ms_offsets * stride_grad_ms, grad_ms, mask=ms_mask) + + grad_h = a[None, :] * grad_h / rms[:, None] + tl.store( + grad_h_ptr + offs_m[:, None] * stride_grad_hm + cols[None, :] * stride_grad_hn, + grad_h, + mask=mask_m[:, None] & mask_n[None, :], + ) + + +def sinkhorn_config(): + block = [256, 1024] + warps = [2, 8] + stages = [2, 4] + configs = [] + for b, w, s in itertools.product(block, warps, stages): + configs.append(triton.Config({"BLOCK_SIZE": b}, num_warps=w, num_stages=s)) + if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": + configs = configs[:1] + return configs + + +@triton.autotune( + configs=sinkhorn_config(), + key=["M"], +) +@triton.jit +def _mhc_sinkhorn_fwd_fused_recompute( + x_ptr, # (M, n*n) + output_ptr, # (M, n*n) + stride_xm, + stride_xn, + stride_out_m, + stride_out_n, + M, + n: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + iters, +): + pid = tl.program_id(0) + + tl.static_assert(BLOCK_SIZE % (n * n) == 0, "BLOCK_SIZE must be divisible by n*n") + tl.assume(M > 0 and iters > 0) + tl.assume(n == 4) + + BATCH_SIZE: tl.constexpr = BLOCK_SIZE // (n * n) + + offs_batch = pid * BATCH_SIZE + tl.arange(0, BATCH_SIZE) + offs_nn = tl.arange(0, n * n) + mask_batch = offs_batch < M + + x_ptrs = x_ptr + offs_batch[:, None] * stride_xm + offs_nn[None, :] * stride_xn + x = tl.load(x_ptrs, mask=mask_batch[:, None], other=0.0) # (BATCH_SIZE, n*n) + x = tl.reshape(x, (BATCH_SIZE, n, n)) # (BATCH_SIZE, n, n) + + log_mu = tl.zeros((BATCH_SIZE, n), dtype=x.dtype) # (BATCH_SIZE, n) + log_nu = tl.zeros((BATCH_SIZE, n), dtype=x.dtype) # (BATCH_SIZE, n) + + f = tl.zeros((BATCH_SIZE, n), dtype=x.dtype) # (BATCH_SIZE, n) + g = tl.zeros((BATCH_SIZE, n), dtype=x.dtype) # (BATCH_SIZE, n) + + for _ in range(iters): + # Update f: logsumexp over the column dimension (1) + f = x + g[:, None, :] # Broadcast g to (BATCH_SIZE, n, n) + f_max = tl.max(f, axis=2) + f = tl.log(tl.sum(tl.exp(f - f_max[:, :, None]), axis=2)) # logsumexp over columns + f = log_mu - f - f_max + + # Update g: logsumexp over the row dimension (2) + g = x + f[:, :, None] # Broadcast f to (BATCH_SIZE, n, n) + g_max = tl.max(g, axis=1) + g = tl.log(tl.sum(tl.exp(g - g_max[:, None, :]), axis=1)) # logsumexp over rows + g = log_nu - g - g_max + + log_P = f[:, :, None] + x + g[:, None, :] + log_P = tl.reshape( + log_P, + ( + BATCH_SIZE, + n * n, + ), + ) + P = tl.exp(log_P) + + output_ptrs = output_ptr + offs_batch[:, None] * stride_out_m + offs_nn[None, :] * stride_out_n + tl.store(output_ptrs, P, mask=mask_batch[:, None]) + + +@triton.autotune( + configs=sinkhorn_config(), + key=["M"], +) +@triton.jit +def _mhc_sinkhorn_bwd_fused_recompute( + grad_out_ptr, + output_ptr, + grad_x_ptr, + x_ptr, + hist_f_ptr, + hist_g_ptr, + stride_grad_out_m, + stride_grad_out_n, + stride_out_m, + stride_out_n, + stride_grad_xm, + stride_grad_xn, + stride_xm, + stride_xn, + M, + n: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + iters, +): + pid = tl.program_id(0) + + tl.static_assert(BLOCK_SIZE % (n * n) == 0, "BLOCK_SIZE must be divisible by n*n") + tl.assume(M > 0 and iters > 0) + tl.assume(n == 4) + + BATCH_SIZE: tl.constexpr = BLOCK_SIZE // (n * n) # Assume there's no remainder for simplicity + + offs_batch = pid * BATCH_SIZE + tl.arange(0, BATCH_SIZE) + offs_nn = tl.arange(0, n * n) + offs_n_hist = tl.arange(0, n) + mask_batch = offs_batch < M + + x_ptrs = x_ptr + offs_batch[:, None] * stride_xm + offs_nn[None, :] * stride_xn + x = tl.load(x_ptrs, mask=mask_batch[:, None], other=0.0) # (BATCH_SIZE, n*n) + x = tl.reshape(x, (BATCH_SIZE, n, n)) # (BATCH_SIZE, n, n) + + P_ptrs = output_ptr + offs_batch[:, None] * stride_out_m + offs_nn[None, :] * stride_out_n + P = tl.load(P_ptrs, mask=mask_batch[:, None], other=0.0) # (BATCH_SIZE, n*n) + P = tl.reshape(P, (BATCH_SIZE, n, n)) + + grad_out_ptrs = ( + grad_out_ptr + + offs_batch[:, None] * stride_grad_out_m + + offs_nn[None, :] * stride_grad_out_n + ) + grad_out = tl.load(grad_out_ptrs, mask=mask_batch[:, None], other=0.0) # (BATCH_SIZE, n*n) + grad_out = tl.reshape(grad_out, (BATCH_SIZE, n, n)) # (BATCH_SIZE, n, n) + + sbn = M * n + + # Recompute the full history of f and g + log_mu = tl.zeros((BATCH_SIZE, n), dtype=x.dtype) # (BATCH_SIZE, n) + log_nu = tl.zeros((BATCH_SIZE, n), dtype=x.dtype) # (BATCH_SIZE, n) + + f = tl.zeros((BATCH_SIZE, n), dtype=x.dtype) # (BATCH_SIZE, n) + g = tl.zeros((BATCH_SIZE, n), dtype=x.dtype) # (BATCH_SIZE, n) + + f_hist_ptrs = hist_f_ptr + offs_batch[:, None] * n + offs_n_hist[None, :] + g_hist_ptrs = hist_g_ptr + offs_batch[:, None] * n + offs_n_hist[None, :] + tl.store(f_hist_ptrs, f, mask=mask_batch[:, None]) + tl.store(g_hist_ptrs, g, mask=mask_batch[:, None]) + + for iter_idx in range(iters): + # Update f: logsumexp over the column dimension (1) + f = x + g[:, None, :] # Broadcast g to (BATCH_SIZE, n, n) + f_max = tl.max(f, axis=2) + f = tl.log(tl.sum(tl.exp(f - f_max[:, :, None]), axis=2)) # logsumexp over columns + f = log_mu - f - f_max + + f_hist_ptrs = ( + hist_f_ptr + (iter_idx + 1) * sbn + offs_batch[:, None] * n + offs_n_hist[None, :] + ) + tl.store(f_hist_ptrs, f, mask=mask_batch[:, None]) + + # Update g: logsumexp over the row dimension (2) + g = x + f[:, :, None] # Broadcast f to (BATCH_SIZE, n, n) + g_max = tl.max(g, axis=1) + g = tl.log(tl.sum(tl.exp(g - g_max[:, None, :]), axis=1)) # logsumexp over rows + g = log_nu - g - g_max + + g_hist_ptrs = ( + hist_g_ptr + (iter_idx + 1) * sbn + offs_batch[:, None] * n + offs_n_hist[None, :] + ) + tl.store(g_hist_ptrs, g, mask=mask_batch[:, None]) + + # Backward pass + grad_log_P = grad_out * P # (BATCH_SIZE, n, n) + zeros = tl.zeros_like(grad_log_P) + grad_g = tl.sum(grad_log_P, axis=1) # (BATCH_SIZE, n) + grad_x = grad_log_P + + g_hist_ptrs = hist_g_ptr + iters * sbn + offs_batch[:, None] * n + offs_n_hist[None, :] + g = tl.load(g_hist_ptrs, mask=mask_batch[:, None], other=0.0) + g = tl.reshape(g, (BATCH_SIZE, n)) + + for iter_idx in range(iters, 0, -1): + f_hist_ptrs = hist_f_ptr + iter_idx * sbn + offs_batch[:, None] * n + offs_n_hist[None, :] + f = tl.load(f_hist_ptrs, mask=mask_batch[:, None], other=0.0) + f = tl.reshape(f, (BATCH_SIZE, n)) + + g_hist_ptrs = ( + hist_g_ptr + (iter_idx - 1) * sbn + offs_batch[:, None] * n + offs_n_hist[None, :] + ) + g_next = tl.load(g_hist_ptrs, mask=mask_batch[:, None], other=0.0) + g_next = tl.reshape(g_next, (BATCH_SIZE, n)) + + term_g = -grad_g[:, None, :] * tl.exp(f[:, :, None] + x + g[:, None, :]) + grad_f = tl.sum(term_g + grad_log_P, axis=2) # (BATCH_SIZE, n) + # Only the last iteration's f will contribute to gradients with both grad_g1 and grad_log_P + grad_log_P = zeros # Zero out grad_log_P for next iterations + + g = g_next + + term_f = -grad_f[:, :, None] * tl.exp(f[:, :, None] + x + g[:, None, :]) + grad_g = tl.sum(term_f, axis=1) # (BATCH_SIZE, n) + + grad_x += term_f + term_g + + grad_x_ptrs = ( + grad_x_ptr + offs_batch[:, None] * stride_grad_xm + offs_nn[None, :] * stride_grad_xn + ) + tl.store( + grad_x_ptrs, + tl.reshape( + grad_x, + ( + BATCH_SIZE, + n * n, + ), + ), + mask=mask_batch[:, None], + ) + + +@triton.autotune( + configs=sinkhorn_config(), + key=["M"], +) +@triton.jit +def _mhc_sinkhorn_fwd_fused( + x_ptr, # (M, n*n) + output_ptr, # (M, n*n) + hist_f_ptr, # (iters+1, M, n) + hist_g_ptr, # (iters+1, M, n) + stride_xm, + stride_xn, + stride_out_m, + stride_out_n, + M, + n: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + iters, +): + pid = tl.program_id(0) + + tl.static_assert(BLOCK_SIZE % (n * n) == 0, "BLOCK_SIZE must be divisible by n*n") + tl.assume(M > 0 and iters > 0) + tl.assume(n == 4) + + BATCH_SIZE: tl.constexpr = BLOCK_SIZE // (n * n) # Assume there's no remainder for simplicity + + offs_batch = pid * BATCH_SIZE + tl.arange(0, BATCH_SIZE) + offs_nn = tl.arange(0, n * n) + offs_n_hist = tl.arange(0, n) + mask_batch = offs_batch < M + + x_ptrs = x_ptr + offs_batch[:, None] * stride_xm + offs_nn[None, :] * stride_xn + x = tl.load(x_ptrs, mask=mask_batch[:, None], other=0.0) # (BATCH_SIZE, n*n) + x = tl.reshape(x, (BATCH_SIZE, n, n)) # (BATCH_SIZE, n, n) + + log_mu = tl.zeros((BATCH_SIZE, n), dtype=x.dtype) # (BATCH_SIZE, n) + log_nu = tl.zeros((BATCH_SIZE, n), dtype=x.dtype) # (BATCH_SIZE, n) + + f = tl.zeros((BATCH_SIZE, n), dtype=x.dtype) # (BATCH_SIZE, n) + g = tl.zeros((BATCH_SIZE, n), dtype=x.dtype) # (BATCH_SIZE, n) + + sbn = M * n + + # Store the initial f and g to history + f_hist_ptrs = hist_f_ptr + offs_batch[:, None] * n + offs_n_hist[None, :] + g_hist_ptrs = hist_g_ptr + offs_batch[:, None] * n + offs_n_hist[None, :] + tl.store(f_hist_ptrs, f, mask=mask_batch[:, None]) + tl.store(g_hist_ptrs, g, mask=mask_batch[:, None]) + + for iter_idx in range(iters): + # Update f: logsumexp over the column dimension (1) + f = x + g[:, None, :] # Broadcast g to (BATCH_SIZE, n, n) + f_max = tl.max(f, axis=2) + f = tl.log(tl.sum(tl.exp(f - f_max[:, :, None]), axis=2)) # logsumexp over columns + f = log_mu - f - f_max + + f_hist_ptrs = ( + hist_f_ptr + (iter_idx + 1) * sbn + offs_batch[:, None] * n + offs_n_hist[None, :] + ) + tl.store(f_hist_ptrs, f, mask=mask_batch[:, None]) + + # Update g: logsumexp over the row dimension (2) + g = x + f[:, :, None] # Broadcast f to (BATCH_SIZE, n, n) + g_max = tl.max(g, axis=1) + g = tl.log(tl.sum(tl.exp(g - g_max[:, None, :]), axis=1)) # logsumexp over rows + g = log_nu - g - g_max + + g_hist_ptrs = ( + hist_g_ptr + (iter_idx + 1) * sbn + offs_batch[:, None] * n + offs_n_hist[None, :] + ) + tl.store(g_hist_ptrs, g, mask=mask_batch[:, None]) + + log_P = f[:, :, None] + x + g[:, None, :] + log_P = tl.reshape( + log_P, + ( + BATCH_SIZE, + n * n, + ), + ) + P = tl.exp(log_P) + + output_ptrs = output_ptr + offs_batch[:, None] * stride_out_m + offs_nn[None, :] * stride_out_n + tl.store(output_ptrs, P, mask=mask_batch[:, None]) + + +@triton.autotune( + configs=sinkhorn_config(), + key=["M"], +) +@triton.jit +def _mhc_sinkhorn_bwd_fused( + grad_out_ptr, # (M, n*n) + output_ptr, # (M, n*n) + grad_x_ptr, # (M, n*n) + x_ptr, # (M, n*n) + hist_f_ptr, # (iters+1, M, n) + hist_g_ptr, # (iters+1, M, n) + stride_grad_out_m, + stride_grad_out_n, + stride_out_m, + stride_out_n, + stride_grad_xm, + stride_grad_xn, + stride_xm, + stride_xn, + M, + n: tl.constexpr, + BLOCK_SIZE: tl.constexpr, + iters, +): + pid = tl.program_id(0) + + tl.static_assert(BLOCK_SIZE % (n * n) == 0, "BLOCK_SIZE must be divisible by n*n") + tl.assume(M > 0 and iters > 0) + tl.assume(n == 4) + + BATCH_SIZE: tl.constexpr = BLOCK_SIZE // (n * n) # Assume there's no remainder for simplicity + + offs_batch = pid * BATCH_SIZE + tl.arange(0, BATCH_SIZE) + offs_nn = tl.arange(0, n * n) + offs_n_hist = tl.arange(0, n) + mask_batch = offs_batch < M + + x_ptrs = x_ptr + offs_batch[:, None] * stride_xm + offs_nn[None, :] * stride_xn + x = tl.load(x_ptrs, mask=mask_batch[:, None], other=0.0) # (BATCH_SIZE, n*n) + x = tl.reshape(x, (BATCH_SIZE, n, n)) # (BATCH_SIZE, n, n) + + P_ptrs = output_ptr + offs_batch[:, None] * stride_out_m + offs_nn[None, :] * stride_out_n + P = tl.load(P_ptrs, mask=mask_batch[:, None], other=0.0) # (BATCH_SIZE, n*n) + P = tl.reshape(P, (BATCH_SIZE, n, n)) + + grad_out_ptrs = ( + grad_out_ptr + + offs_batch[:, None] * stride_grad_out_m + + offs_nn[None, :] * stride_grad_out_n + ) + grad_out = tl.load(grad_out_ptrs, mask=mask_batch[:, None], other=0.0) # (BATCH_SIZE, n*n) + grad_out = tl.reshape(grad_out, (BATCH_SIZE, n, n)) # (BATCH_SIZE, n, n) + + sbn = M * n + + # Backward pass + grad_log_P = grad_out * P # (BATCH_SIZE, n, n) + zeros = tl.zeros_like(grad_log_P) + grad_g = tl.sum(grad_log_P, axis=1) # (BATCH_SIZE, n) + grad_x = grad_log_P + + g_hist_ptrs = hist_g_ptr + iters * sbn + offs_batch[:, None] * n + offs_n_hist[None, :] + g = tl.load(g_hist_ptrs, mask=mask_batch[:, None], other=0.0) + g = tl.reshape(g, (BATCH_SIZE, n)) + + for iter_idx in range(iters, 0, -1): + f_hist_ptrs = hist_f_ptr + iter_idx * sbn + offs_batch[:, None] * n + offs_n_hist[None, :] + f = tl.load(f_hist_ptrs, mask=mask_batch[:, None], other=0.0) + f = tl.reshape(f, (BATCH_SIZE, n)) + + g_hist_ptrs = ( + hist_g_ptr + (iter_idx - 1) * sbn + offs_batch[:, None] * n + offs_n_hist[None, :] + ) + g_next = tl.load(g_hist_ptrs, mask=mask_batch[:, None], other=0.0) + g_next = tl.reshape(g_next, (BATCH_SIZE, n)) + + term_g = -grad_g[:, None, :] * tl.exp(f[:, :, None] + x + g[:, None, :]) + grad_f = tl.sum(term_g + grad_log_P, axis=2) # (BATCH_SIZE, n) + # Only the last iteration's f will contribute to gradients with both grad_g1 and grad_log_P + grad_log_P = zeros # Zero out grad_log_P for next iterations + + g = g_next + + term_f = -grad_f[:, :, None] * tl.exp(f[:, :, None] + x + g[:, None, :]) + grad_g = tl.sum(term_f, axis=1) # (BATCH_SIZE, n) + + grad_x += term_f + term_g + + grad_x_ptrs = ( + grad_x_ptr + offs_batch[:, None] * stride_grad_xm + offs_nn[None, :] * stride_grad_xn + ) + tl.store( + grad_x_ptrs, + tl.reshape( + grad_x, + ( + BATCH_SIZE, + n * n, + ), + ), + mask=mask_batch[:, None], + ) + + +def aggregate_config(): + block_m = [1, 2, 4] + block_c = [64, 128, 256] + warps = [1, 2, 4] + stages = [1, 2, 3, 4] + + configs = [] + for m, c, w, s in itertools.product(block_m, block_c, warps, stages): + configs.append( + triton.Config({"BLOCK_SIZE_M": m, "BLOCK_SIZE_C": c}, num_warps=w, num_stages=s) + ) + if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": + configs = configs[:1] + return configs + + +@triton.autotune( + configs=aggregate_config(), + key=["M", "C"], +) +@triton.jit +def _mhc_aggregate_fwd( + x_ptr, # # (M, C, n) + H_pre_ptr, # (M, n) + output_ptr, # (M, C) + M, + C, + n: tl.constexpr, + stride_xm, + stride_xCn, + stride_output_m, + stride_output_c, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_C: tl.constexpr, +): + """ + output = x @ H_pre: (M, C, n) @ (M, n, 1) = (M, C, 1) + """ + pid_m = tl.program_id(1) + pid_c = tl.program_id(0) + + tl.static_assert(n == 4) + tl.assume(M > 0) + tl.assume(C > 0) + tl.assume(n == 4) + tl.assume(stride_xm > 0 and stride_xCn == 1) + tl.assume(stride_output_m > 0 and stride_output_c == 1) + + tl.assume(BLOCK_SIZE_C % 32 == 0) + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_c = pid_c * BLOCK_SIZE_C + tl.arange(0, BLOCK_SIZE_C) + offs_cn = pid_c * BLOCK_SIZE_C * n + tl.arange(0, BLOCK_SIZE_C * n) + mask_m = offs_m < M + mask_c = offs_c < C + mask_cn = offs_cn < C * n + + offs_H_pre = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) + H_pre = tl.load( + H_pre_ptr + offs_H_pre, mask=offs_H_pre < M * n, other=0.0, cache_modifier=".ca" + ) # (BLOCK_SIZE_M * n) + H_pre = H_pre.reshape(BLOCK_SIZE_M, 2, 2) + H_pre01, H_pre23 = tl.split(H_pre) + H_pre0, H_pre1 = tl.split(H_pre01) + H_pre2, H_pre3 = tl.split(H_pre23) # (BLOCK_SIZE_M, 1) + + x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_cn[None, :] * stride_xCn + x = tl.load( + x_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C * n) + + x = tl.reshape(x, (BLOCK_SIZE_M, BLOCK_SIZE_C, 2, 2)) + x01, x23 = tl.split(x) + x0, x1 = tl.split(x01) + x2, x3 = tl.split(x23) # (BLOCK_SIZE_M, BLOCK_SIZE_C) + + # x @ H_pre: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, 1) + # triton doesn't support dot prod with inner dimension < 16, so we need to manually unroll the computation for n=4: + # x @ H_pre = x[:, :, 0] * H_pre[:, 0] + # + x[:, :, 1] * H_pre[:, 1] + # + x[:, :, 2] * H_pre[:, 2] + # + x[:, :, 3] * H_pre[:, 3] + out_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_C), dtype=tl.float32) + out_acc = tl.fma(x0, H_pre0[:, None], out_acc) + out_acc = tl.fma(x1, H_pre1[:, None], out_acc) + out_acc = tl.fma(x2, H_pre2[:, None], out_acc) + out_acc = tl.fma(x3, H_pre3[:, None], out_acc) + + out = out_acc.to(x.dtype) + + output_ptrs = output_ptr + offs_m[:, None] * stride_output_m + offs_c[None, :] * stride_output_c + tl.store(output_ptrs, out, mask=mask_m[:, None] & mask_c[None, :]) + + +@triton.autotune(configs=aggregate_config(), key=["M", "C"], reset_to_zero=["grad_H_pre_ptr"]) +@triton.jit +def _mhc_aggregate_bwd( + grad_output_ptr, # (M, C) + H_pre_ptr, # (M, n) + grad_H_pre_ptr, # (M, n) + x_ptr, # (M, C, n) + grad_x_ptr, # # (M, C, n) + M, + C, + n: tl.constexpr, + stride_grad_output_m, + stride_grad_output_c, + stride_xm, + stride_xCn, + stride_grad_xm, + stride_grad_xCn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_C: tl.constexpr, + precision: tl.constexpr, +): + """ + Forward: + out = x @ H_pre: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, 1) = (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) + Backward: + grad_H_pre = x.T @ grad_output: (BLOCK_SIZE_M, n, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) = (BLOCK_SIZE_M, n, 1) + grad_H_pre.T = grad_output.T @ x: (BLOCK_SIZE_M, 1, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, 1, n) + which is easier to compute since transposing grad_H_pre and grad_output is just view change + grad_x = grad_output @ H_pre.T: (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + """ + pid_m = tl.program_id(1) + pid_c = tl.program_id(0) + + tl.static_assert(n == 4) + tl.assume(M > 0) + tl.assume(C > 0) + tl.assume(n == 4) + tl.assume(stride_xm > 0 and stride_xCn == 1) + tl.assume(stride_grad_xm > 0 and stride_grad_xCn == 1) + tl.assume(stride_grad_output_m > 0 and stride_grad_output_c == 1) + + tl.assume(BLOCK_SIZE_C % 32 == 0) + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_c = pid_c * BLOCK_SIZE_C + tl.arange(0, BLOCK_SIZE_C) + offs_cn = pid_c * BLOCK_SIZE_C * n + tl.arange(0, BLOCK_SIZE_C * n) + mask_m = offs_m < M + mask_c = offs_c < C + mask_cn = offs_cn < C * n + + grad_output_ptrs = ( + grad_output_ptr + + offs_m[:, None] * stride_grad_output_m + + offs_c[None, :] * stride_grad_output_c + ) + grad_output = tl.load( + grad_output_ptrs, mask=mask_m[:, None] & mask_c[None, :], other=0.0 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C) + + x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_cn[None, :] * stride_xCn + x = tl.load( + x_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C * n) + + grad_H_pre = tl.dot( + tl.reshape(grad_output, (BLOCK_SIZE_M, 1, BLOCK_SIZE_C)), + tl.reshape(x, (BLOCK_SIZE_M, BLOCK_SIZE_C, n)), + input_precision=precision, + out_dtype=tl.float32, + ) + grad_H_pre = tl.reshape(grad_H_pre, (BLOCK_SIZE_M * n,)) # (BLOCK_SIZE_M * n) + offs_grad_H_pre = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) + grad_H_pre_ptrs = grad_H_pre_ptr + offs_grad_H_pre + tl.atomic_add(grad_H_pre_ptrs, grad_H_pre, mask=offs_grad_H_pre < M * n, sem="relaxed") + + H_pre_offs = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) + H_pre = tl.load( + H_pre_ptr + H_pre_offs, mask=H_pre_offs < M * n, other=0.0, cache_modifier=".ca" + ) # (BLOCK_SIZE_M * n) + H_pre = tl.reshape(H_pre, (BLOCK_SIZE_M, n)) # (BLOCK_SIZE_M, n) + + # grad_x = grad_output @ H_pre.T: (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + grad_x = grad_output[:, :, None] * H_pre[:, None, :] # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + grad_x = tl.reshape(grad_x, (BLOCK_SIZE_M, BLOCK_SIZE_C * n)) + + grad_x_ptrs = grad_x_ptr + offs_m[:, None] * stride_grad_xm + offs_cn[None, :] * stride_grad_xCn + tl.store( + grad_x_ptrs, + grad_x, + mask=mask_m[:, None] & mask_cn[None, :], + ) + + +def expand_combine_config(): + block_m = [1, 2, 4] + block_c = [128, 256] + warps = [1, 2] + stages = [1, 2, 3, 4] + + configs = [] + for m, c, w, s in itertools.product(block_m, block_c, warps, stages): + configs.append( + triton.Config({"BLOCK_SIZE_M": m, "BLOCK_SIZE_C": c}, num_warps=w, num_stages=s) + ) + if os.environ.get("NVTE_DISABLE_TRITON_AUTOTUNING", "0") == "1": + configs = configs[:1] + return configs + + +@triton.autotune( + configs=expand_combine_config(), + key=["M", "C"], +) +@triton.jit +def _mhc_expand_combine_fwd( + f_ptr, # (M, C) + H_post_ptr, # (M, n) + x_ptr, # (M, C, n) + H_res_ptr, # (M, n, n) + output_ptr, # # (M, C, n) + M, + C, + n: tl.constexpr, + stride_fm, + stride_fc, + stride_xm, + stride_xCn, + stride_output_m, + stride_output_Cn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_C: tl.constexpr, +): + """ + output = f @ H_post: (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + + x @ H_res: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + """ + pid_m = tl.program_id(1) + pid_c = tl.program_id(0) + + tl.static_assert(n == 4) + tl.assume(M > 0) + tl.assume(C > 0) + tl.assume(n == 4) + tl.assume(stride_fm > 0 and stride_fc == 1) + tl.assume(stride_xm > 0 and stride_xCn == 1) + tl.assume(stride_output_m > 0 and stride_output_Cn == 1) + + tl.assume(BLOCK_SIZE_C % 32 == 0) + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_c = pid_c * BLOCK_SIZE_C + tl.arange(0, BLOCK_SIZE_C) + offs_cn = pid_c * BLOCK_SIZE_C * n + tl.arange(0, BLOCK_SIZE_C * n) + mask_m = offs_m < M + mask_c = offs_c < C + mask_cn = offs_cn < C * n + + f_ptrs = f_ptr + offs_m[:, None] * stride_fm + offs_c[None, :] * stride_fc + f = tl.load(f_ptrs, mask=mask_m[:, None] & mask_c[None, :], other=0.0) + + offs_H_post = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) + H_post = tl.load( + H_post_ptr + offs_H_post, mask=offs_H_post < M * n, other=0.0, cache_modifier=".ca" + ) + H_post = tl.reshape(H_post, (BLOCK_SIZE_M, n)) # (BLOCK_SIZE_M, n) + + # Residual connection path: res_out = f @ H_post: + # (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, n) = (BLOCK_SIZE_M, n, BLOCK_SIZE_C) + # Due to broadcasting, it's equivalent to a multiplicaiton + out_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_C, n), dtype=tl.float32) + out_acc = tl.fma(f[:, :, None], H_post[:, None, :], out_acc) + + H_res_offs = pid_m * BLOCK_SIZE_M * n * n + tl.arange(0, BLOCK_SIZE_M * n * n) + H_res = tl.load( + H_res_ptr + H_res_offs, mask=H_res_offs < M * n * n, other=0.0, cache_modifier=".ca" + ) + H_res = tl.reshape(H_res, (BLOCK_SIZE_M, n, n)) # (BLOCK_SIZE_M, n, n) + + x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_cn[None, :] * stride_xCn + x = tl.load( + x_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + + # Manifold connection path: manifold_out = H_res @ x: + # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + # triton doesn't support dot prod with inner dimension < 16, so we need to manually unroll the computation for n=4: + # x @ H_res = x[:, :, 0] @ H_res[:, 0, :] + # + x[:, :, 1] @ H_res[:, 1, :] + # + x[:, :, 2] @ H_res[:, 2, :] + # + x[:, :, 3] @ H_res[:, 3, :] + + x_reshape = tl.reshape(x, (BLOCK_SIZE_M, BLOCK_SIZE_C, 2, 2)) + x01, x23 = tl.split( + x_reshape + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, 2), (BLOCK_SIZE_M, BLOCK_SIZE_C, 2) + x0, x1 = tl.split(x01) # (BLOCK_SIZE_M, BLOCK_SIZE_C), (BLOCK_SIZE_M, BLOCK_SIZE_C) + x2, x3 = tl.split(x23) # (BLOCK_SIZE_M, BLOCK_SIZE_C), (BLOCK_SIZE_M, BLOCK_SIZE_C) + + H_resT = tl.reshape(tl.trans(H_res, (0, 2, 1)), (BLOCK_SIZE_M, n, 2, 2)) + H_res01, H_res23 = tl.split(H_resT) # (BLOCK_SIZE_M, n, 2), (BLOCK_SIZE_M, n, 2) + H_res0, H_res1 = tl.split(H_res01) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) + H_res2, H_res3 = tl.split(H_res23) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) + + out_acc = tl.fma(x0[:, :, None], H_res0[:, None, :], out_acc) + out_acc = tl.fma(x1[:, :, None], H_res1[:, None, :], out_acc) + out_acc = tl.fma(x2[:, :, None], H_res2[:, None, :], out_acc) + out_acc = tl.fma(x3[:, :, None], H_res3[:, None, :], out_acc) + + out = out_acc.to(x.dtype) + out = tl.reshape(out, (BLOCK_SIZE_M, BLOCK_SIZE_C * n)) # (BLOCK_SIZE_M, BLOCK_SIZE_C*n) + + output_ptrs = ( + output_ptr + offs_m[:, None] * stride_output_m + offs_cn[None, :] * stride_output_Cn + ) + tl.store(output_ptrs, out, mask=mask_m[:, None] & mask_cn[None, :]) + + +@triton.autotune( + configs=expand_combine_config(), + key=["M", "C"], + reset_to_zero=["grad_H_post_ptr", "grad_H_res_ptr"], +) +@triton.jit +def _mhc_expand_combine_bwd( + grad_output_ptr, # (M, C, n) + f_ptr, # (M, C) + H_post_ptr, # (M, n) + x_ptr, # (M, C, n) + H_res_ptr, # (M, n, n) + grad_H_post_ptr, # (M, n) + grad_f_ptr, # (M, C) + grad_H_res_ptr, # (M, n, n) + grad_x_ptr, # (M, C, n) + M, + C, + n: tl.constexpr, + stride_grad_output_m, + stride_grad_output_Cn, + stride_fm, + stride_fc, + stride_xm, + stride_xCn, + stride_grad_fm, + stride_grad_fc, + stride_grad_xm, + stride_grad_xCn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_C: tl.constexpr, + precision: tl.constexpr, +): + """ + Each block + It reads + - (BLOCK_SIZE_M, BLOCK_SIZE_C) of f, which is the output of the attention / FFN module + - (BLOCK_SIZE_M, n) of H_post, which is applied for the transformation of the attention / FFN output + - (BLOCK_SIZE_M, BLOCK_SIZE_C, n) of x, which is the skip connection's input + - (BLOCK_SIZE_M, n*n) of H_res, which is applied for the transformation of the skip connection + and writes + - (BLOCK_SIZE_M, n) of grad_H_post + - (BLOCK_SIZE_M, BLOCK_SIZE_C) of grad_f + - (BLOCK_SIZE_M, n, n) of grad_H_res + - (BLOCK_SIZE_M, BLOCK_SIZE_C, n) of grad_x + + Forward: + out = f @ H_post + x @ H_res + Backward: + GEMM: + grad_H_post = f.T @ grad_output: (BLOCK_SIZE_M, 1, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, 1, n) + grad_H_res = x.T @ grad_output: (BLOCK_SIZE_M, n, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, n, n) + Not GEMM: + grad_f = grad_output @ H_post.T: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, 1) = (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) + grad_x = grad_output @ H_res.T: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + """ + + pid_m = tl.program_id(1) + pid_c = tl.program_id(0) + + tl.static_assert(n == 4) + tl.assume(M > 0) + tl.assume(C > 0) + tl.assume(n == 4) + tl.assume(stride_fm > 0 and stride_fc == 1) + tl.assume(stride_xm > 0 and stride_xCn == 1) + tl.assume(stride_grad_output_m > 0 and stride_grad_output_Cn == 1) + tl.assume(stride_grad_fm > 0 and stride_grad_fc == 1) + tl.assume(stride_grad_xm > 0 and stride_grad_xCn == 1) + + tl.assume(BLOCK_SIZE_C % 32 == 0) + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_c = pid_c * BLOCK_SIZE_C + tl.arange(0, BLOCK_SIZE_C) + offs_cn = pid_c * BLOCK_SIZE_C * n + tl.arange(0, BLOCK_SIZE_C * n) + mask_m = offs_m < M + mask_c = offs_c < C + mask_cn = offs_cn < C * n + + f_ptrs = f_ptr + offs_m[:, None] * stride_fm + offs_c[None, :] * stride_fc + f = tl.load(f_ptrs, mask=mask_m[:, None] & mask_c[None, :], other=0.0) + + H_post_offs = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) + H_post = tl.load(H_post_ptr + H_post_offs, mask=H_post_offs < M * n, other=0.0) + H_post = tl.reshape(H_post, (BLOCK_SIZE_M, n)) # (BLOCK_SIZE_M, n) + + H_res_offs = pid_m * BLOCK_SIZE_M * n * n + tl.arange(0, BLOCK_SIZE_M * n * n) + H_res = tl.load( + H_res_ptr + H_res_offs, mask=H_res_offs < M * n * n, other=0.0 + ) # (BLOCK_SIZE_M, n, n) + H_res = tl.reshape(H_res, (BLOCK_SIZE_M, n, n)) # (BLOCK_SIZE_M, n, n) + + grad_out_ptrs = ( + grad_output_ptr + + offs_m[:, None] * stride_grad_output_m + + offs_cn[None, :] * stride_grad_output_Cn + ) + grad_out = tl.load( + grad_out_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C * n) + grad_out = tl.reshape( + grad_out, (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + + # grad_H_post = f.T @ grad_output # (BLOCK_SIZE_M, 1, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, 1, n) + grad_H_post = tl.dot( + tl.reshape(f, (BLOCK_SIZE_M, 1, BLOCK_SIZE_C)), + tl.reshape(grad_out, (BLOCK_SIZE_M, BLOCK_SIZE_C, n)), + input_precision=precision, + out_dtype=tl.float32, + ) # (BLOCK_SIZE_M, 1, n) + grad_H_post = tl.reshape(grad_H_post, (BLOCK_SIZE_M * n,)) # (BLOCK_SIZE_M * n) + offs_grad_H_post = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) + grad_H_post_ptrs = grad_H_post_ptr + offs_grad_H_post + tl.atomic_add(grad_H_post_ptrs, grad_H_post, mask=offs_grad_H_post < M * n, sem="relaxed") + + x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_cn[None, :] * stride_xCn + x = tl.load( + x_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C*n) + x = tl.reshape(x, (BLOCK_SIZE_M, BLOCK_SIZE_C, n)) # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + + # grad_H_res = x.T @ grad_output: (BLOCK_SIZE_M, n, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, n, n) + grad_H_res = tl.dot( + tl.trans(x, (0, 2, 1)), grad_out, input_precision=precision, out_dtype=tl.float32 + ) # (BLOCK_SIZE_M, n, n) + grad_H_res = tl.reshape(grad_H_res, (BLOCK_SIZE_M * n * n,)) # (BLOCK_SIZE_M * n * n) + offs_grad_H_res = pid_m * BLOCK_SIZE_M * n * n + tl.arange(0, BLOCK_SIZE_M * n * n) + grad_H_res_ptrs = grad_H_res_ptr + offs_grad_H_res + tl.atomic_add( + grad_H_res_ptrs, grad_H_res.to(tl.float32), mask=offs_grad_H_res < M * n * n, sem="relaxed" + ) + + grad_out_reshape = tl.reshape( + grad_out, (BLOCK_SIZE_M, BLOCK_SIZE_C, 2, 2) + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, 2, 2) + grad_out01, grad_out23 = tl.split( + grad_out_reshape + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, 2), (BLOCK_SIZE_M, BLOCK_SIZE_C, 2) + grad_out0, grad_out1 = tl.split( + grad_out01 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C), (BLOCK_SIZE_M, BLOCK_SIZE_C) + grad_out2, grad_out3 = tl.split( + grad_out23 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C), (BLOCK_SIZE_M, BLOCK_SIZE_C) + + # grad_f = grad_output @ H_post.T: (BLOCK_SIZE_M, 1, n) @ (BLOCK_SIZE_M, n, BLOCK_SIZE_C) = (BLOCK_SIZE_M, 1, BLOCK_SIZE_C) + # Triton doesn't support dot prod with inner dimension < 16, so we need to hack this: + # grad_f = grad_out[:, :, 0] @ H_post.T[:, 0, :] (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, 1) + # + grad_out[:, :, 1] @ H_post.T[:, 1, :] + # + grad_out[:, :, 2] @ H_post.T[:, 2, :] + # + grad_out[:, :, 3] @ H_post.T[:, 3, :] + # where H_post.T[:, i, :] = H_post[:, :, i] + H_post = tl.reshape(H_post, (BLOCK_SIZE_M, 2, 2)) + H_post01, H_post23 = tl.split(H_post) # (BLOCK_SIZE_M, 2), (BLOCK_SIZE_M, 2) + H_post0, H_post1 = tl.split(H_post01) # (BLOCK_SIZE_M,), (BLOCK_SIZE_M,) + H_post2, H_post3 = tl.split(H_post23) # (BLOCK_SIZE_M,), (BLOCK_SIZE_M,) + + grad_f_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_C), dtype=tl.float32) + # (BLOCK_SIZE_M, BLOCK_SIZE_C) * (BLOCK_SIZE_M, 1) -> (BLOCK_SIZE_M, BLOCK_SIZE_C) + grad_f_acc = tl.fma(grad_out0, H_post0[:, None], grad_f_acc) + grad_f_acc = tl.fma(grad_out1, H_post1[:, None], grad_f_acc) + grad_f_acc = tl.fma(grad_out2, H_post2[:, None], grad_f_acc) + grad_f_acc = tl.fma(grad_out3, H_post3[:, None], grad_f_acc) + grad_f = grad_f_acc.to(f.dtype) + + grad_f_ptrs = grad_f_ptr + offs_m[:, None] * stride_grad_fm + offs_c[None, :] * stride_grad_fc + tl.store(grad_f_ptrs, grad_f, mask=mask_m[:, None] & mask_c[None, :]) + + # grad_x = grad_output @ H_res.T: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, n, BLOCK_SIZE_C) + # The inner dim is n=4 which is too small for triton, so we will manually unroll the matmul + # grad_x = grad_out[:, :, 0] @ H_res.T[:, 0, :] + # + grad_out[:, :, 1] @ H_res.T[:, 1, :] + # + grad_out[:, :, 2] @ H_res.T[:, 2, :] + # + grad_out[:, :, 3] @ H_res.T[:, 3, :] + # where H_res.T[:, i, :] = H_res[:, :, i] + # Due to broadcasting, it's equivalent to multiplying each H_res[:, i, :].T with grad_out[:, i, :] + + H_res_reshape = tl.reshape(H_res, (BLOCK_SIZE_M, n, 2, 2)) # (BLOCK_SIZE_M, n, 2, 2) + H_res01, H_res23 = tl.split(H_res_reshape) # (BLOCK_SIZE_M, n, 2), (BLOCK_SIZE_M, n, 2) + H_res0, H_res1 = tl.split(H_res01) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) + H_res2, H_res3 = tl.split(H_res23) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) + + grad_x_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_C, n), dtype=tl.float32) + grad_x_acc = tl.fma(grad_out0[:, :, None], H_res0[:, None, :], grad_x_acc) + grad_x_acc = tl.fma(grad_out1[:, :, None], H_res1[:, None, :], grad_x_acc) + grad_x_acc = tl.fma(grad_out2[:, :, None], H_res2[:, None, :], grad_x_acc) + grad_x_acc = tl.fma(grad_out3[:, :, None], H_res3[:, None, :], grad_x_acc) + + grad_x = grad_x_acc.to(x.dtype) + grad_x = tl.reshape(grad_x, (BLOCK_SIZE_M, BLOCK_SIZE_C * n)) # (BLOCK_SIZE_M, BLOCK_SIZE_C*n) + + grad_x_ptrs = grad_x_ptr + offs_m[:, None] * stride_grad_xm + offs_cn[None, :] * stride_grad_xCn + tl.store(grad_x_ptrs, grad_x, mask=mask_m[:, None] & mask_cn[None, :]) + + +@triton.autotune( + configs=expand_combine_config(), + key=["M", "C"], +) +@triton.jit +def _mhc_expand_combine_with_bias_fwd( + f_ptr, # (M, C) + bias_ptr, # (C,) + H_post_ptr, # (M, n) + x_ptr, # (M, C, n) + H_res_ptr, # (M, n, n) + output_ptr, # # (M, C, n) + M, + C, + n: tl.constexpr, + stride_fm, + stride_fc, + stride_bias, + stride_xm, + stride_xCn, + stride_output_m, + stride_output_Cn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_C: tl.constexpr, +): + """ + output = (f + bias[None, :, None]) @ H_post: (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + + x @ H_res: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + """ + pid_m = tl.program_id(1) + pid_c = tl.program_id(0) + + tl.static_assert(n == 4) + tl.assume(M > 0) + tl.assume(C > 0) + tl.assume(n == 4) + tl.assume(stride_fm > 0 and stride_fc == 1) + tl.assume(stride_bias == 1) + tl.assume(stride_xm > 0 and stride_xCn == 1) + tl.assume(stride_output_m > 0 and stride_output_Cn == 1) + + tl.assume(BLOCK_SIZE_C % 32 == 0) + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_c = pid_c * BLOCK_SIZE_C + tl.arange(0, BLOCK_SIZE_C) + offs_cn = pid_c * BLOCK_SIZE_C * n + tl.arange(0, BLOCK_SIZE_C * n) + mask_m = offs_m < M + mask_c = offs_c < C + mask_cn = offs_cn < C * n + + f_ptrs = f_ptr + offs_m[:, None] * stride_fm + offs_c[None, :] * stride_fc + f = tl.load(f_ptrs, mask=mask_m[:, None] & mask_c[None, :], other=0.0) + bias = tl.load(bias_ptr + offs_c * stride_bias, mask=mask_c, other=0.0) # (BLOCK_SIZE_C,) + + offs_H_post = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) + H_post = tl.load( + H_post_ptr + offs_H_post, mask=offs_H_post < M * n, other=0.0, cache_modifier=".ca" + ) + H_post = tl.reshape(H_post, (BLOCK_SIZE_M, n)) # (BLOCK_SIZE_M, n) + + # Residual connection path: res_out = f @ H_post + bias @ H_post: + # (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, n) = (BLOCK_SIZE_M, n, BLOCK_SIZE_C) + # Due to broadcasting, it's equivalent to a multiplicaiton + out_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_C, n), dtype=tl.float32) + out_acc = tl.fma(bias[None, :, None], H_post[:, None, :], out_acc) + out_acc = tl.fma(f[:, :, None], H_post[:, None, :], out_acc) + + H_res_offs = pid_m * BLOCK_SIZE_M * n * n + tl.arange(0, BLOCK_SIZE_M * n * n) + H_res = tl.load( + H_res_ptr + H_res_offs, mask=H_res_offs < M * n * n, other=0.0, cache_modifier=".ca" + ) + H_res = tl.reshape(H_res, (BLOCK_SIZE_M, n, n)) # (BLOCK_SIZE_M, n, n) + + x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_cn[None, :] * stride_xCn + x = tl.load( + x_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + + # Manifold connection path: manifold_out = H_res @ x: + # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + # triton doesn't support dot prod with inner dimension < 16, so we need to manually unroll the computation for n=4: + # x @ H_res = x[:, :, 0] @ H_res[:, 0, :] + # + x[:, :, 1] @ H_res[:, 1, :] + # + x[:, :, 2] @ H_res[:, 2, :] + # + x[:, :, 3] @ H_res[:, 3, :] + + x_reshape = tl.reshape(x, (BLOCK_SIZE_M, BLOCK_SIZE_C, 2, 2)) + x01, x23 = tl.split( + x_reshape + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, 2), (BLOCK_SIZE_M, BLOCK_SIZE_C, 2) + x0, x1 = tl.split(x01) # (BLOCK_SIZE_M, BLOCK_SIZE_C), (BLOCK_SIZE_M, BLOCK_SIZE_C) + x2, x3 = tl.split(x23) # (BLOCK_SIZE_M, BLOCK_SIZE_C), (BLOCK_SIZE_M, BLOCK_SIZE_C) + + H_resT = tl.reshape(tl.trans(H_res, (0, 2, 1)), (BLOCK_SIZE_M, n, 2, 2)) + H_res01, H_res23 = tl.split(H_resT) # (BLOCK_SIZE_M, n, 2), (BLOCK_SIZE_M, n, 2) + H_res0, H_res1 = tl.split(H_res01) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) + H_res2, H_res3 = tl.split(H_res23) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) + + out_acc = tl.fma(x0[:, :, None], H_res0[:, None, :], out_acc) + out_acc = tl.fma(x1[:, :, None], H_res1[:, None, :], out_acc) + out_acc = tl.fma(x2[:, :, None], H_res2[:, None, :], out_acc) + out_acc = tl.fma(x3[:, :, None], H_res3[:, None, :], out_acc) + + out = out_acc.to(x.dtype) + out = tl.reshape(out, (BLOCK_SIZE_M, BLOCK_SIZE_C * n)) # (BLOCK_SIZE_M, BLOCK_SIZE_C*n) + + output_ptrs = ( + output_ptr + offs_m[:, None] * stride_output_m + offs_cn[None, :] * stride_output_Cn + ) + tl.store(output_ptrs, out, mask=mask_m[:, None] & mask_cn[None, :]) + + +@triton.autotune( + configs=expand_combine_config(), + key=["M", "C"], + reset_to_zero=["grad_H_post_ptr", "grad_H_res_ptr", "grad_bias_ptr"], +) +@triton.jit +def _mhc_expand_combine_with_bias_bwd( + grad_output_ptr, # (M, C, n) + f_ptr, # (M, C) + bias_ptr, # (C,) + H_post_ptr, # (M, n) + x_ptr, # (M, C, n) + H_res_ptr, # (M, n, n) + grad_H_post_ptr, # (M, n) + grad_f_ptr, # (M, C) + grad_bias_ptr, # (C,) + grad_H_res_ptr, # (M, n, n) + grad_x_ptr, # (M, C, n) + M, + C, + n: tl.constexpr, + stride_grad_output_m, + stride_grad_output_Cn, + stride_fm, + stride_fc, + stride_bias, + stride_xm, + stride_xCn, + stride_grad_fm, + stride_grad_fc, + stride_grad_bias, + stride_grad_xm, + stride_grad_xCn, + # Meta-parameters + BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_C: tl.constexpr, + precision: tl.constexpr, +): + """ + Each block + It reads + - (BLOCK_SIZE_M, BLOCK_SIZE_C) of f, which is the output of the attention / FFN module + - (BLOCK_SIZE_M, n) of H_post, which is applied for the transformation of the attention / FFN output + - (BLOCK_SIZE_M, BLOCK_SIZE_C, n) of x, which is the skip connection's input + - (BLOCK_SIZE_M, n*n) of H_res, which is applied for the transformation of the skip connection + and writes + - (BLOCK_SIZE_M, n) of grad_H_post + - (BLOCK_SIZE_M, BLOCK_SIZE_C) of grad_f + - (BLOCK_SIZE_M, n, n) of grad_H_res + - (BLOCK_SIZE_M, BLOCK_SIZE_C, n) of grad_x + + Forward: + out = f @ H_post + x @ H_res + Backward: + GEMM: + grad_H_post = f.T @ grad_output: (BLOCK_SIZE_M, 1, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, 1, n) + grad_H_res = x.T @ grad_output: (BLOCK_SIZE_M, n, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, n, n) + Not GEMM: + grad_f = grad_output @ H_post.T: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, 1) = (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) + grad_x = grad_output @ H_res.T: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + """ + + pid_m = tl.program_id(1) + pid_c = tl.program_id(0) + + tl.static_assert(n == 4) + tl.assume(M > 0) + tl.assume(C > 0) + tl.assume(n == 4) + tl.assume(stride_fm > 0 and stride_fc == 1) + tl.assume(stride_bias == 1) + tl.assume(stride_xm > 0 and stride_xCn == 1) + tl.assume(stride_grad_output_m > 0 and stride_grad_output_Cn == 1) + tl.assume(stride_grad_fm > 0 and stride_grad_fc == 1) + tl.assume(stride_grad_bias == 1) + tl.assume(stride_grad_xm > 0 and stride_grad_xCn == 1) + + tl.assume(BLOCK_SIZE_C % 32 == 0) + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_c = pid_c * BLOCK_SIZE_C + tl.arange(0, BLOCK_SIZE_C) + offs_cn = pid_c * BLOCK_SIZE_C * n + tl.arange(0, BLOCK_SIZE_C * n) + mask_m = offs_m < M + mask_c = offs_c < C + mask_cn = offs_cn < C * n + + f_ptrs = f_ptr + offs_m[:, None] * stride_fm + offs_c[None, :] * stride_fc + f = tl.load(f_ptrs, mask=mask_m[:, None] & mask_c[None, :], other=0.0) + + bias = tl.load(bias_ptr + offs_c * stride_bias, mask=mask_c, other=0.0) # (BLOCK_SIZE_C,) + + H_post_offs = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) + H_post = tl.load(H_post_ptr + H_post_offs, mask=H_post_offs < M * n, other=0.0) + H_post = tl.reshape(H_post, (BLOCK_SIZE_M, n)) # (BLOCK_SIZE_M, n) + + H_res_offs = pid_m * BLOCK_SIZE_M * n * n + tl.arange(0, BLOCK_SIZE_M * n * n) + H_res = tl.load( + H_res_ptr + H_res_offs, mask=H_res_offs < M * n * n, other=0.0 + ) # (BLOCK_SIZE_M, n, n) + H_res = tl.reshape(H_res, (BLOCK_SIZE_M, n, n)) # (BLOCK_SIZE_M, n, n) + + grad_out_ptrs = ( + grad_output_ptr + + offs_m[:, None] * stride_grad_output_m + + offs_cn[None, :] * stride_grad_output_Cn + ) + grad_out = tl.load( + grad_out_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C * n) + grad_out = tl.reshape( + grad_out, (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + + # grad_H_post = f.T @ grad_output # (BLOCK_SIZE_M, 1, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, 1, n) + grad_H_post = tl.dot( + tl.reshape(f, (BLOCK_SIZE_M, 1, BLOCK_SIZE_C)), + tl.reshape(grad_out, (BLOCK_SIZE_M, BLOCK_SIZE_C, n)), + input_precision=precision, + out_dtype=tl.float32, + ) # (BLOCK_SIZE_M, 1, n) + grad_H_post = tl.dot( + tl.broadcast_to(bias[None, None, :], (BLOCK_SIZE_M, 1, BLOCK_SIZE_C)), + tl.reshape(grad_out, (BLOCK_SIZE_M, BLOCK_SIZE_C, n)), + acc=grad_H_post, + input_precision=precision, + out_dtype=tl.float32, + ) # (BLOCK_SIZE_M, 1, n) + grad_H_post = tl.reshape(grad_H_post, (BLOCK_SIZE_M * n,)) # (BLOCK_SIZE_M * n) + offs_grad_H_post = pid_m * BLOCK_SIZE_M * n + tl.arange(0, BLOCK_SIZE_M * n) + grad_H_post_ptrs = grad_H_post_ptr + offs_grad_H_post + tl.atomic_add(grad_H_post_ptrs, grad_H_post, mask=offs_grad_H_post < M * n, sem="relaxed") + + x_ptrs = x_ptr + offs_m[:, None] * stride_xm + offs_cn[None, :] * stride_xCn + x = tl.load( + x_ptrs, mask=mask_m[:, None] & mask_cn[None, :], other=0.0 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C*n) + x = tl.reshape(x, (BLOCK_SIZE_M, BLOCK_SIZE_C, n)) # (BLOCK_SIZE_M, BLOCK_SIZE_C, n) + + # grad_H_res = x.T @ grad_output: (BLOCK_SIZE_M, n, BLOCK_SIZE_C) @ (BLOCK_SIZE_M, BLOCK_SIZE_C, n) = (BLOCK_SIZE_M, n, n) + grad_H_res = tl.dot( + tl.trans(x, (0, 2, 1)), grad_out, input_precision=precision, out_dtype=tl.float32 + ) # (BLOCK_SIZE_M, n, n) + grad_H_res = tl.reshape(grad_H_res, (BLOCK_SIZE_M * n * n,)) # (BLOCK_SIZE_M * n * n) + offs_grad_H_res = pid_m * BLOCK_SIZE_M * n * n + tl.arange(0, BLOCK_SIZE_M * n * n) + grad_H_res_ptrs = grad_H_res_ptr + offs_grad_H_res + tl.atomic_add( + grad_H_res_ptrs, grad_H_res.to(tl.float32), mask=offs_grad_H_res < M * n * n, sem="relaxed" + ) + + grad_out_reshape = tl.reshape( + grad_out, (BLOCK_SIZE_M, BLOCK_SIZE_C, 2, 2) + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, 2, 2) + grad_out01, grad_out23 = tl.split( + grad_out_reshape + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C, 2), (BLOCK_SIZE_M, BLOCK_SIZE_C, 2) + grad_out0, grad_out1 = tl.split( + grad_out01 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C), (BLOCK_SIZE_M, BLOCK_SIZE_C) + grad_out2, grad_out3 = tl.split( + grad_out23 + ) # (BLOCK_SIZE_M, BLOCK_SIZE_C), (BLOCK_SIZE_M, BLOCK_SIZE_C) + + # grad_f = grad_output @ H_post.T: (BLOCK_SIZE_M, 1, n) @ (BLOCK_SIZE_M, n, BLOCK_SIZE_C) = (BLOCK_SIZE_M, 1, BLOCK_SIZE_C) + # Triton doesn't support dot prod with inner dimension < 16, so we need to hack this: + # = grad_out[:, :, 0] @ H_post.T[:, 0, :] (BLOCK_SIZE_M, BLOCK_SIZE_C, 1) @ (BLOCK_SIZE_M, 1, 1) + # + grad_out[:, :, 1] @ H_post.T[:, 1, :] + # + grad_out[:, :, 2] @ H_post.T[:, 2, :] + # + grad_out[:, :, 3] @ H_post.T[:, 3, :] + # where H_post.T[:, i, :] = H_post[:, :, i] + H_post = tl.reshape(H_post, (BLOCK_SIZE_M, 2, 2)) + H_post01, H_post23 = tl.split(H_post) # (BLOCK_SIZE_M, 2), (BLOCK_SIZE_M, 2) + H_post0, H_post1 = tl.split(H_post01) # (BLOCK_SIZE_M,), (BLOCK_SIZE_M,) + H_post2, H_post3 = tl.split(H_post23) # (BLOCK_SIZE_M,), (BLOCK_SIZE_M,) + + grad_f_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_C), dtype=tl.float32) + # (BLOCK_SIZE_M, BLOCK_SIZE_C) * (BLOCK_SIZE_M, 1) -> (BLOCK_SIZE_M, BLOCK_SIZE_C) + grad_f_acc = tl.fma(grad_out0, H_post0[:, None], grad_f_acc) + grad_f_acc = tl.fma(grad_out1, H_post1[:, None], grad_f_acc) + grad_f_acc = tl.fma(grad_out2, H_post2[:, None], grad_f_acc) + grad_f_acc = tl.fma(grad_out3, H_post3[:, None], grad_f_acc) + grad_f = grad_f_acc.to(f.dtype) + + grad_f_ptrs = grad_f_ptr + offs_m[:, None] * stride_grad_fm + offs_c[None, :] * stride_grad_fc + tl.store(grad_f_ptrs, grad_f, mask=mask_m[:, None] & mask_c[None, :]) + + grad_bias = tl.sum(grad_f_acc, axis=0) # (BLOCK_SIZE_C,) + grad_bias_ptrs = grad_bias_ptr + offs_c * stride_grad_bias + tl.atomic_add(grad_bias_ptrs, grad_bias, mask=mask_c, sem="relaxed") + + # grad_x = grad_output @ H_res.T: (BLOCK_SIZE_M, BLOCK_SIZE_C, n) @ (BLOCK_SIZE_M, n, n) = (BLOCK_SIZE_M, n, BLOCK_SIZE_C) + # The inner dim is n=4 which is too small for triton, so we will manually unroll the matmul + # grad_x = grad_out[:, :, 0] @ H_res.T[:, 0, :] + # + grad_out[:, :, 1] @ H_res.T[:, 1, :] + # + grad_out[:, :, 2] @ H_res.T[:, 2, :] + # + grad_out[:, :, 3] @ H_res.T[:, 3, :] + # where H_res.T[:, i, :] = H_res[:, :, i] + # Due to broadcasting, it's equivalent to multiplying each H_res[:, i, :].T with grad_out[:, i, :] + + H_res_reshape = tl.reshape(H_res, (BLOCK_SIZE_M, n, 2, 2)) # (BLOCK_SIZE_M, n, 2, 2) + H_res01, H_res23 = tl.split(H_res_reshape) # (BLOCK_SIZE_M, n, 2), (BLOCK_SIZE_M, n, 2) + H_res0, H_res1 = tl.split(H_res01) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) + H_res2, H_res3 = tl.split(H_res23) # (BLOCK_SIZE_M, n), (BLOCK_SIZE_M, n) + + grad_x_acc = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_C, n), dtype=tl.float32) + grad_x_acc = tl.fma(grad_out0[:, :, None], H_res0[:, None, :], grad_x_acc) + grad_x_acc = tl.fma(grad_out1[:, :, None], H_res1[:, None, :], grad_x_acc) + grad_x_acc = tl.fma(grad_out2[:, :, None], H_res2[:, None, :], grad_x_acc) + grad_x_acc = tl.fma(grad_out3[:, :, None], H_res3[:, None, :], grad_x_acc) + + grad_x = grad_x_acc.to(x.dtype) + grad_x = tl.reshape(grad_x, (BLOCK_SIZE_M, BLOCK_SIZE_C * n)) # (BLOCK_SIZE_M, BLOCK_SIZE_C*n) + + grad_x_ptrs = grad_x_ptr + offs_m[:, None] * stride_grad_xm + offs_cn[None, :] * stride_grad_xCn + tl.store(grad_x_ptrs, grad_x, mask=mask_m[:, None] & mask_cn[None, :]) diff --git a/transformer_engine/pytorch/triton/__init__.py b/transformer_engine/pytorch/triton/__init__.py index d86cededd..6d3141253 100644 --- a/transformer_engine/pytorch/triton/__init__.py +++ b/transformer_engine/pytorch/triton/__init__.py @@ -3,3 +3,4 @@ # See LICENSE for license information. """PyTorch wrappers for Triton kernels.""" +from transformer_engine.pytorch.triton import mhc diff --git a/transformer_engine/pytorch/triton/mhc.py b/transformer_engine/pytorch/triton/mhc.py new file mode 100644 index 000000000..987216e32 --- /dev/null +++ b/transformer_engine/pytorch/triton/mhc.py @@ -0,0 +1,999 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""PyTorch wrapper functions for mHC (manifold Hyper-Connection) Triton kernels.""" + +import os +import torch +import triton + +from transformer_engine.common.triton.mhc import ( + _mhc_scale_fwd_fused, + _mhc_scale_bwd_fused, + _mhc_expand_combine_with_bias_fwd, + _mhc_expand_combine_with_bias_bwd, + _mhc_expand_combine_fwd, + _mhc_expand_combine_bwd, + _mhc_aggregate_fwd, + _mhc_aggregate_bwd, + _mhc_projection_fwd_fused, + _mhc_projection_bwd_fused, + _mhc_sinkhorn_fwd_fused, + _mhc_sinkhorn_fwd_fused_recompute, + _mhc_sinkhorn_bwd_fused, + _mhc_sinkhorn_bwd_fused_recompute, +) +from transformer_engine.pytorch.cpp_extensions.gemm import general_gemm + + +def check_deterministic(operator: str): + """ + Checks if the non-deterministic algorithm is allowed for the given operator. If not, raises an assertion error with instructions on how to allow it. + Since atomic add is used in this mHC implementation, it breaks the determinism guarantee due to non-associativity of floating point addition. + """ + allow_nondeterministic = os.environ.get("NVTE_ALLOW_NONDETERMINISTIC_ALGO", "1") == "1" + assert allow_nondeterministic, ( + f"[{operator}]: This operation uses atomic add which violates determinism. Set" + " NVTE_ALLOW_NONDETERMINISTIC_ALGO=1 to allow this non-deterministic behavior." + ) + + +def mhc_fused_sinkhorn( + H_res: torch.Tensor, n: int = 4, recompute_hist: bool = True, iters: int = 20 +): + """ + Sinkhorn operation to compute the final H_res matrix (see eq. 19, section 4.3.1 of the DeepSeek mHC paper): + + The Sinkhorn operation conducts an iterative normalization process that alternately rescales rows and columns to sum to 1. + This kernel performs this operation in the log space for numerical stability. + + Parameters + ---------- + H_res : torch.Tensor + input H_res matrix of shape (s, b, n, n) that needs to be normalized into a doubly stochastic matrix. + n : int + number of hyper connections, where only n=4 is supported in the current implementation + recompute_hist : bool + whether to recompute the intermediate history in the backward pass to save memory + iters : int + number of Sinkhorn iterations, according to the DeepSeek paper 20 is enough for convergence + + Returns + ------- + out : torch.Tensor + out of shape (s, b, n, n), which is the final H_res after Sinkhorn normalization + """ + assert n == 4, "Only n=4 is supported in this implementation" + out = mHCSinkhornOp.apply(H_res, n, recompute_hist, iters) + return out + + +def mhc_fused_scale( + H: torch.Tensor, alpha: torch.Tensor, beta: torch.Tensor, ms: torch.Tensor, n: int +): + """ + Fused scale operation to compute the scaled H matrices (see eq. 16-18, section 4.3.1 of the DeepSeek mHC paper): + + H_pre = H[:, 0:n] * alpha[0] / sqrt(ms) + beta[0:n] + H_post = H[:, n:2n] * alpha[1] / sqrt(ms) + beta[n:2n] + H_res = H[:, 2n:2n+n*n] * alpha[2] / sqrt(ms) + beta[2n:2n+n*n] + + H_pre = sigmoid(H_pre) + H_post = 2*sigmoid(H_post) + + Parameters + ---------- + H : torch.Tensor + input H matrix of shape (M, 32), where M=s*b, and only the first N elements in the last dimension are valid + alpha : torch.Tensor + scaling factor for H, of shape (3,), where + alpha[0] is applied to H[:, 0:n] for H_pre + alpha[1] is applied to H[:, n:2n] for H_post + alpha[2] is applied to H[:, 2n:2n+n*n] for H_res + beta : torch.Tensor + bias term for H, of shape (1, 2*n+n*n), where + beta[0, 0:n] is applied to H[:, 0:n] for H_pre + beta[0, n:2n] is applied to H[:, n:2n] for H_post + beta[0, 2n:2n+n*n] is applied to H[:, 2n:2n+n*n] for H_res + ms : torch.Tensor + mean square for each row of H from the projection kernel, of shape (M,), used for RMSNorm scaling + n : int + number of hyper connections, where only n=4 is supported in the current implementation + + Returns + ------- + h_pre : torch.Tensor + Scaled H_pre of shape (M, n), which aggregates (s, b, C, n) input of a Hyper Connection block into (s, b, n) as the input of attention / MLP + h_post : torch.Tensor + Scaled H_post of shape (M, n), which expands the output of attention / MLP of shape (s, b, n) back to (s, b, C, n) for the residual connection + h_res : torch.Tensor + Scaled H_res of shape (M, n*n), which mixes the n streams of the (s, b, C, n) input of a Hyper Connection block + + """ + assert n == 4, "Only n=4 is supported in this implementation" + check_deterministic("mhc_fused_scale") + out = mHCScaleFusedOp.apply(H, alpha, beta, ms, n) + h_pre = out[..., :n] + h_post = out[..., n : 2 * n] + h_res = out[..., 2 * n : n * n + 2 * n] + return h_pre, h_post, h_res + + +def mhc_fused_aggregate(x: torch.Tensor, H_pre: torch.Tensor, n: int, use_tf32: bool = True): + """ + Aggregate operation to merge n activation streams into one (see section 4.3.1 of the DeepSeek mHC paper): + out = x @ H_pre: (s, b, C, n) @ (s, b, n, 1) -> (s, b, C, 1) -> (s, b, C) after squeezing the last dimension + + Parameters + ---------- + x : torch.Tensor + input activation tensor of shape (s, b, C, n), + where s is the sequence length, b is the batch size, C is the hidden dimension per hyper connection, and n is the number of hyper connections. Note that C is equal to the original hidden dimension divided by n. + H_pre: torch.Tensor + input H_pre matrix of shape (s, b, n) + n: int + number of hyper connections, where only n=4 is supported in the current implementation + use_tf32: bool + whether to use TF32 precision for matmul operations. If False, it will use ieee for better precision. + This is mainly used by our unittests since TF32 precision will introduce some errors and cause tests to fail + + Returns + ------- + out: torch.Tensor + output activation tensor of shape (s, b, C), which is the aggregated output after merging n hyper connections + """ + assert n == 4, "Only n=4 is supported in this implementation" + check_deterministic("mhc_fused_aggregate") + out = mHCAggregateOp.apply(x, H_pre, n, use_tf32) + return out + + +def mhc_fused_expand_combine( + f: torch.Tensor, + bias: torch.Tensor, + H_post: torch.Tensor, + x: torch.Tensor, + H_res: torch.Tensor, + n: int, + use_tf32: bool = True, +): + """ + Expand and combine operation for merging n hyper connections (see section 4.3.1 of the DeepSeek mHC paper): + + out = (f [+ bias]) @ H_post + x @ H_res: (s, b, C, 1) @ (s, b, 1, n) + (s, b, C, n) @ (s, b, n, n) -> (s, b, C, n) + + Parameters + ---------- + f : torch.Tensor + input activation tensor of shape (s, b, C), which is the output from the attention / FFN sub-layer in a transformer block + bias : torch.Tensor or None + optional bias tensor of shape (C,) from the last linear layer, where f + bias is fused in this kernel for better performance + H_post : torch.Tensor + input H_post matrix of shape (s, b, n) + x : torch.Tensor + input activation tensor of shape (s, b, C, n), which is the hyper connection input before the aggregation operation + H_res : torch.Tensor + input H_res matrix of shape (s, b, n, n) + n : int + number of hyper connections + use_tf32 : bool + whether to use TF32 precision for matmul operations. If False, it will use ieee for better precision. + This is mainly used by our unittests since TF32 precision will introduce some errors and cause tests to fail + + Returns + ------- + out : torch.Tensor + out of shape (s, b, C, n), which is the expanded and combined output after merging n hyper connections + """ + assert n == 4, "Only n=4 is supported in this implementation" + check_deterministic("mhc_fused_expand_combine") + out = mHCExpandCombineOp.apply( + f, + bias, + H_post, + x, + H_res, + n, + use_tf32, + ) + return out + + +def mhc_fused_projection(x: torch.Tensor, phi: torch.Tensor, use_tf32: bool = True): + """ + Fused projection operation to compute H matrices and mean square for RMSNorm (see eq. 14-15, section 4.3.1 of the DeepSeek mHC paper): + + H = x @ phi^T: (M, K) @ (K, N) -> (M, N), which is padded to (M, 32) for better memory access pattern in the next kernels. + ms = mean(x^2, dim=-1): (M,) + + Note: the current implementation only supports n=4 + + Parameters + ---------- + x : torch.Tensor + input tensor of shape (M, K), where M=s*b is the batch size and K=nC is the hidden dimension after expansion. + phi : torch.Tensor + projection matrix of shape (N, K), where N=2n+n*n (=24 for n=4) + use_tf32 : bool + whether to use TF32 precision for matmul operations. If False, it will use ieee for better precision. + This is mainly used by our unittests since TF32 precision will introduce some errors and cause tests to fail. + + Returns + ------- + H : torch.Tensor + Projected matrix of shape (M, 32), where only the first N elements in the last dimension are valid. + ms : torch.Tensor + Mean square of shape (M,), which is used for RMSNorm in the next kernel. + """ + assert ( + phi.shape[0] == 24 + ), "Currently only n=4 is supported, which means phi should have 24 in its first dimension" + check_deterministic("mhc_fused_projection") + H, ms = mHCProjectionOp.apply(x, phi, use_tf32) + return H, ms + + +class mHCProjectionOp(torch.autograd.Function): + """ + PyTorch operator for the fused projection operation in mHC, whose wrapper API is mhc_fused_projection. + """ + + @staticmethod + def forward(ctx, x, phi, use_tf32=True): + """ + The forward pass of the fused projection operation. Computes H = x @ phi^T and the mean + square ms = mean(x^2, dim=-1) for RMSNorm in a single fused kernel. + + Parameters: + ctx : The context object. + x (tensor): The input tensor of shape (M, K), where M=s*b is the flattened batch dimension and K=nC is the hidden dimension after expansion. + phi (tensor): The projection matrix of shape (N, K), where N=2n+n*n (=24 for n=4). + use_tf32 (bool): Whether to use TF32 precision for matmul operations. If False, uses IEEE for better precision. + + Returns: + tuple: A tuple of (H, ms) where H is the projected matrix of shape (M, 32) padded for memory alignment (only the first N elements are valid), and ms is the mean square of shape (M,) in FP32. + """ + x = x.contiguous() + phi = phi.contiguous() + + ctx.use_tf32 = use_tf32 + ctx.dtype = x.dtype + + M, K = x.shape + device = x.device + + N = phi.shape[0] + + # Pad H to (s, b, 32) for better memory access pattern in the kernel, but only the first N elements in the last dimension are valid + H = torch.zeros((M, 32), device=device, dtype=torch.float32) + ms = torch.zeros( + (M,), device=device, dtype=torch.float32 + ) # Mean square for x, used to compute RMSNorm in the next kernel + + # pylint: disable=unnecessary-lambda-assignment + grid = lambda META: ( + triton.cdiv(M, META["BLOCK_SIZE_M"]), + triton.cdiv(K, META["BLOCK_SIZE_K"]), + ) + + _mhc_projection_fwd_fused[grid]( + x_ptr=x, # (M, K) + phi_ptr=phi, # (N, K) + h_ptr=H, # (M, 32) + ms_ptr=ms, # (M,) + M=M, + N=N, + K=K, + stride_xm=K, + stride_xk=1, + stride_phin=K, + stride_phik=1, + stride_hm=32, + stride_hn=1, + stride_ms=1, + BLOCK_SIZE_N=32, + precision="tf32" if use_tf32 else "ieee", + ) + + ctx.save_for_backward(x, phi, ms) + ctx.phi_dtype = phi.dtype + + return H.to(ctx.dtype), ms # Keep ms in fp32 + + @staticmethod + def backward(ctx, grad_H, grad_ms): + """ + The backward pass of the fused projection operation. Computes gradients for x and phi. + + grad_phi = grad_H^T @ x, truncated to the first N rows. + grad_x = grad_H @ phi + 2 * x * grad_ms / K, where the second term is the gradient contribution from + the mean square computation fused in the forward pass. + + Parameters: + ctx : The context object with saved tensors. + grad_H (tensor): The gradient of the loss with respect to H, of shape (M, 32). + grad_ms (tensor): The gradient of the loss with respect to the mean square, of shape (M,). + + Returns: + tuple: A tuple with the gradients (grad_x, grad_phi, None). + """ + x, phi, ms = ctx.saved_tensors + M, K = x.shape + device = x.device + + N = phi.shape[0] + + grad_H = grad_H.contiguous().view(M, -1) + grad_ms = grad_ms.contiguous().view( + M, + ) + ms = ms.contiguous().view( + M, + ) + + grad_x = torch.empty((M, K), device=device, dtype=x.dtype) + + grad_x = torch.empty((M, K), device=device, dtype=x.dtype) + grad_phi = general_gemm(x, grad_H, out_dtype=torch.float32, layout="NT")[0][:N, :].to( + phi.dtype + ) # (2n + n^2, M) @ (M, nC) = (2n + n^2, nC); grad_H's last dim is padded to 32 + + # pylint: disable=unnecessary-lambda-assignment + grid = lambda META: ( + triton.cdiv(M, META["BLOCK_SIZE_M"]), + triton.cdiv(K, META["BLOCK_SIZE_K"]), + ) + + _mhc_projection_bwd_fused[grid]( + x_ptr=x, + grad_x_ptr=grad_x, # (M, K) + phi_ptr=phi, # (N, K) + grad_h_ptr=grad_H, # (M, 32) + grad_ms_ptr=grad_ms, # (M,) + M=M, + N=N, + K=K, + stride_xm=K, + stride_xk=1, + stride_grad_xm=K, + stride_grad_xk=1, + stride_phin=K, + stride_phik=1, + stride_grad_phin=K, + stride_grad_phik=1, + stride_grad_hm=32, + stride_grad_hn=1, + stride_grad_ms=1, + BLOCK_SIZE_N=32, + precision="tf32" if ctx.use_tf32 else "ieee", + ) + + return grad_x.to(ctx.dtype), grad_phi.to(ctx.dtype), None + + +class mHCScaleFusedOp(torch.autograd.Function): + """ + PyTorch operator for the fused scale operation in mHC, whose wrapper API is mhc_fused_scale. + """ + + @staticmethod + def forward(ctx, H, alpha, beta, ms, n): + """ + The forward pass of the fused scale operation. Applies RMSNorm scaling, bias, and activation + functions to produce H_pre, H_post, and H_res: + + H_pre = sigmoid(H[:, 0:n] * alpha[0] / sqrt(ms) + beta[0:n]) + H_post = 2 * sigmoid(H[:, n:2n] * alpha[1] / sqrt(ms) + beta[n:2n]) + H_res = H[:, 2n:2n+n*n] * alpha[2] / sqrt(ms) + beta[2n:2n+n*n] + + Parameters: + ctx : The context object. + H (tensor): The input H matrix of shape (M, 32), where only the first N=2n+n*n elements are valid. + alpha (tensor): The scaling factors of shape (3,), one for each of H_pre, H_post, H_res. + beta (tensor): The bias terms of shape (1, 2n+n*n). + ms (tensor): The mean square from the projection kernel, of shape (M,), used for RMSNorm scaling. + n (int): The number of hyper connections (only n=4 is supported). + + Returns: + tensor: The scaled output of shape (M, 32), where only the first N elements are valid. + """ + + ctx.dtype = H.dtype + H = H.to(torch.float32) + alpha = alpha.to(torch.float32) + beta = beta.to(torch.float32) + ms = ms.to(torch.float32) + + M, _ = H.shape + + H = H.contiguous() + beta = beta.contiguous() + ms = ms.contiguous() + + out = torch.empty( + (M, 32), device=H.device, dtype=H.dtype + ) # Pad the output to 32 in the last dimension + + # pylint: disable=unnecessary-lambda-assignment + grid = lambda META: (triton.cdiv(M, META["BLOCK_SIZE_M"]),) + + _mhc_scale_fwd_fused[grid]( + h_ptr=H, # (M, N), which is padded to (M, 32) + b_ptr=beta, # (N,) + a_ptr=alpha, # (N,) + ms_ptr=ms, # (M,) + out_ptr=out, # (M, N), which is padded to (M, 32) + M=M, + n=n, + stride_hm=32, + stride_hn=1, + stride_a=1, + stride_b=1, + stride_ms=1, + stride_out_m=32, + stride_out_n=1, # strides for out, which is padded to 32 in the last dimension + BLOCK_SIZE_N=32, + eps=torch.finfo(ms.dtype).eps, + ) + + ctx.save_for_backward(H, alpha, ms, out) + ctx.n = n + + return out.to(ctx.dtype) # Cast back to the original dtype of H + + @staticmethod + def backward(ctx, grad_out): + """ + The backward pass of the fused scale operation. Computes gradients for H, alpha, beta, and ms + by backpropagating through the sigmoid activations, RMSNorm scaling, and bias additions. + + Parameters: + ctx : The context object with saved tensors. + grad_out (tensor): The gradient of the loss with respect to the output, of shape (M, 32). + + Returns: + tuple: A tuple with the gradients (grad_H, grad_alpha, grad_beta, grad_ms, None). + """ + H, alpha, ms, out = ctx.saved_tensors + n = ctx.n + + grad_out = grad_out.contiguous() + grad_out = grad_out.to(torch.float32) + + M, _ = grad_out.shape + N = 2 * n + n * n + + grad_h = torch.zeros( + (M, 32), device=grad_out.device, dtype=grad_out.dtype + ) # Pad the grad_h to 32 in the last dimension + grad_alpha = torch.zeros((3,), device=grad_out.device, dtype=grad_out.dtype) + grad_beta_padded = torch.zeros((1, 32), device=grad_out.device, dtype=grad_out.dtype) + grad_beta = grad_beta_padded[ + :, :N + ] # Use only the first N elements for grad_beta, the rest are just padding + grad_ms = torch.zeros((M,), device=grad_out.device, dtype=grad_out.dtype) + + # pylint: disable=unnecessary-lambda-assignment + grid = lambda META: (triton.cdiv(M, META["BLOCK_SIZE_M"]),) + + _mhc_scale_bwd_fused[grid]( + grad_out_ptr=grad_out, + out_ptr=out, + grad_h_ptr=grad_h, + h_ptr=H, + grad_a_ptr=grad_alpha, + a_ptr=alpha, + grad_b_ptr=grad_beta, + grad_ms_ptr=grad_ms, + ms_ptr=ms, + M=M, + n=n, + stride_grad_out_m=32, + stride_grad_out_n=1, + stride_out_m=32, + stride_out_n=1, + stride_grad_hm=32, + stride_grad_hn=1, + stride_hm=32, + stride_hn=1, + stride_grad_a=1, + stride_a=1, + stride_grad_b=1, + stride_grad_ms=1, + stride_ms=1, + BLOCK_SIZE_N=32, + eps=torch.finfo(ms.dtype).eps, + ) + + return ( + grad_h.to(ctx.dtype), + grad_alpha.to(ctx.dtype), + grad_beta.to(ctx.dtype), + grad_ms.to(ctx.dtype), + None, + ) + + +class mHCSinkhornOp(torch.autograd.Function): + """ + PyTorch operator for the Sinkhorn operation in mHC, whose wrapper API is mhc_fused_sinkhorn. + """ + + @staticmethod + def forward(ctx, H_res, n=4, recompute_hist=True, iters=20): + """ + The forward pass of the Sinkhorn operation. Performs iterative row-column normalization + in log space to convert H_res into a doubly stochastic matrix. Each iteration alternately + rescales rows and columns to sum to 1: + + f = log_mu - logsumexp(H_res + g, dim=cols) + g = log_nu - logsumexp(H_res + f, dim=rows) + output = exp(f + H_res + g) + + Parameters: + ctx : The context object. + H_res (tensor): The input H_res matrix of shape (s, b, n, n). + n (int): The number of hyper connections (only n=4 is supported). + recompute_hist (bool): Whether to recompute the intermediate f/g history in the backward pass to save memory. If False, stores history buffers of shape (iters+1, s, b, n). + iters (int): The number of Sinkhorn iterations (20 is enough for convergence per the DeepSeek paper). + + Returns: + tensor: The doubly stochastic matrix of shape (s, b, n, n). + """ + + s, b, _, _ = H_res.shape + + ctx.dtype = H_res.dtype + H_res = H_res.to(torch.float32) + + H_res = H_res.contiguous().view(s * b, n * n) + + hist_f, hist_g = None, None + if not recompute_hist: + # History buffers: (iters+1, s, b, n) + hist_f = torch.empty((iters + 1, s, b, n), device=H_res.device, dtype=H_res.dtype) + hist_g = torch.empty((iters + 1, s, b, n), device=H_res.device, dtype=H_res.dtype) + H_res_out = torch.empty_like(H_res) # (s*b, n*n) + + # pylint: disable=unnecessary-lambda-assignment + grid = lambda META: (triton.cdiv(s * b * n * n, META["BLOCK_SIZE"]),) + + if recompute_hist: + _mhc_sinkhorn_fwd_fused_recompute[grid]( + x_ptr=H_res, + output_ptr=H_res_out, + stride_xm=n * n, + stride_xn=1, + stride_out_m=n * n, + stride_out_n=1, + M=s * b, + n=n, + iters=iters, + ) + else: + _mhc_sinkhorn_fwd_fused[grid]( + x_ptr=H_res, + output_ptr=H_res_out, + hist_f_ptr=hist_f, + hist_g_ptr=hist_g, + stride_xm=n * n, + stride_xn=1, + stride_out_m=n * n, + stride_out_n=1, + M=s * b, + n=n, + iters=iters, + ) + + if recompute_hist: + ctx.save_for_backward(H_res, H_res_out) + else: + ctx.save_for_backward(H_res, H_res_out, hist_f, hist_g) + ctx.recompute_hist = recompute_hist + ctx.iters = iters + ctx.n = n + + H_res_out = H_res_out.view(s, b, n, n) + return H_res_out.to(ctx.dtype) # Cast back to the original dtype of H + + @staticmethod + def backward(ctx, grad_out): + """ + The backward pass of the Sinkhorn operation. Backpropagates through the iterative + normalization by reversing through the f/g update steps. If recompute_hist is True, + the forward pass history is recomputed to save memory. + + Parameters: + ctx : The context object with saved tensors. + grad_out (tensor): The gradient of the loss with respect to the output, of shape (s, b, n, n). + + Returns: + tuple: A tuple with the gradients (grad_H_res, None, None, None). + """ + + s, b, n, _ = grad_out.shape + M = s * b + + hist_f, hist_g = None, None + recompute_hist = ctx.recompute_hist + iters = ctx.iters + if recompute_hist: + H_res, H_res_out = ctx.saved_tensors + hist_f = torch.empty((iters + 1, s, b, n), device=H_res.device, dtype=H_res.dtype) + hist_g = torch.empty((iters + 1, s, b, n), device=H_res.device, dtype=H_res.dtype) + else: + H_res, H_res_out, hist_f, hist_g = ctx.saved_tensors + + n = ctx.n + + grad_res_out = grad_out.clone().contiguous().view(M, n * n) + + grad_res = torch.empty_like(H_res) + + # pylint: disable=unnecessary-lambda-assignment + grid = lambda META: (triton.cdiv(M * n * n, META["BLOCK_SIZE"]),) + + if recompute_hist: + _mhc_sinkhorn_bwd_fused_recompute[grid]( + grad_out_ptr=grad_res_out, + output_ptr=H_res_out, + grad_x_ptr=grad_res, + x_ptr=H_res, + hist_f_ptr=hist_f, + hist_g_ptr=hist_g, + stride_grad_out_m=n * n, + stride_grad_out_n=1, + stride_out_m=n * n, + stride_out_n=1, + stride_grad_xm=n * n, + stride_grad_xn=1, + stride_xm=n * n, + stride_xn=1, + M=M, + n=n, + iters=iters, + ) + else: + _mhc_sinkhorn_bwd_fused[grid]( + grad_out_ptr=grad_res_out, + output_ptr=H_res_out, + grad_x_ptr=grad_res, + x_ptr=H_res, + hist_f_ptr=hist_f, + hist_g_ptr=hist_g, + stride_grad_out_m=n * n, + stride_grad_out_n=1, + stride_out_m=n * n, + stride_out_n=1, + stride_grad_xm=n * n, + stride_grad_xn=1, + stride_xm=n * n, + stride_xn=1, + M=M, + n=n, + iters=iters, + ) + + grad_res = grad_res.view(s, b, n, n) + + return grad_res.to(ctx.dtype), None, None, None + + +class mHCAggregateOp(torch.autograd.Function): + """ + PyTorch operator for the aggregate operation in mHC, whose wrapper API is mhc_fused_aggregate. + """ + + @staticmethod + def forward(ctx, x, H_pre, n, use_tf32=True): + """ + The forward pass of the aggregate operation. Merges n activation streams into one by + computing a weighted sum using H_pre: + + out = x @ H_pre: (s, b, C, n) @ (s, b, n, 1) -> (s, b, C) + + Parameters: + ctx : The context object. + x (tensor): The input activation tensor of shape (s, b, C, n). + H_pre (tensor): The pre-connection matrix of shape (s, b, n), used as weights for aggregation. + n (int): The number of hyper connections (only n=4 is supported). + use_tf32 (bool): Whether to use TF32 precision for matmul operations. + + Returns: + tensor: The aggregated output of shape (s, b, C). + """ + + x = x.contiguous() + H_pre = H_pre.contiguous() + + s, b, C, n = x.shape + nC = n * C + M = s * b + + out = torch.empty((s, b, C), device=x.device, dtype=x.dtype) + + # pylint: disable=unnecessary-lambda-assignment + grid = lambda META: ( + triton.cdiv(C, META["BLOCK_SIZE_C"]), + triton.cdiv(M, META["BLOCK_SIZE_M"]), + ) + + _mhc_aggregate_fwd[grid]( + x_ptr=x, + H_pre_ptr=H_pre, + output_ptr=out, + M=M, + C=C, + n=n, + stride_xm=nC, + stride_xCn=1, + stride_output_m=C, + stride_output_c=1, + ) + + ctx.save_for_backward(x, H_pre) + ctx.n = n + ctx.use_tf32 = use_tf32 + + return out + + @staticmethod + def backward(ctx, grad_output): + """ + The backward pass of the aggregate operation. Computes gradients for x and H_pre: + + grad_x[:, :, :, i] = grad_output * H_pre[:, :, i] for each stream i + grad_H_pre[:, :, i] = sum_C(grad_output * x[:, :, :, i]) for each stream i + + Parameters: + ctx : The context object with saved tensors. + grad_output (tensor): The gradient of the loss with respect to the output, of shape (s, b, C). + + Returns: + tuple: A tuple with the gradients (grad_x, grad_H_pre, None, None). + """ + grad_output = grad_output.contiguous() + + x, H_pre = ctx.saved_tensors + n = ctx.n + + s, b, C, n = x.shape + nC = n * C + assert n == 4, "Only n=4 is supported in this implementation" + M = s * b + + grad_x = torch.empty_like(x) + grad_H_pre = torch.zeros( + (s, b, n), dtype=torch.float32, device=H_pre.device + ) # We need to use atomic_add for this so we need higher precision + + # pylint: disable=unnecessary-lambda-assignment + grid = lambda META: ( + triton.cdiv(C, META["BLOCK_SIZE_C"]), + triton.cdiv(M, META["BLOCK_SIZE_M"]), + ) + + _mhc_aggregate_bwd[grid]( + grad_output_ptr=grad_output, + H_pre_ptr=H_pre, + grad_H_pre_ptr=grad_H_pre, + x_ptr=x, + grad_x_ptr=grad_x, + M=M, + C=C, + n=n, + stride_grad_output_m=C, + stride_grad_output_c=1, + stride_xm=nC, + stride_xCn=1, + stride_grad_xm=nC, + stride_grad_xCn=1, + precision="tf32" if ctx.use_tf32 else "ieee", + ) + + grad_H_pre = grad_H_pre.to(H_pre.dtype) # Cast back to the original dtype of H_pre + + return grad_x, grad_H_pre, None, None + + +class mHCExpandCombineOp(torch.autograd.Function): + """ + PyTorch operator for the expand and combine operation in mHC, whose wrapper API is mhc_fused_expand_combine. + """ + + @staticmethod + def forward(ctx, f, bias, H_post, x, H_res, n, use_tf32=True): + """ + The forward pass of the expand and combine operation. Expands the sub-layer output f back + to n streams using H_post, and combines with the residual connections using H_res: + + out = (f [+ bias]) @ H_post + x @ H_res: (s, b, C, 1) @ (s, b, 1, n) + (s, b, C, n) @ (s, b, n, n) -> (s, b, C, n) + + Parameters: + ctx : The context object. + f (tensor): The sub-layer output tensor of shape (s, b, C). + bias (tensor or None): Optional bias tensor of shape (C,) from the last linear layer, fused in this kernel. + H_post (tensor): The post-connection matrix of shape (s, b, n). + x (tensor): The hyper connection input tensor of shape (s, b, C, n) before aggregation. + H_res (tensor): The residual connection matrix of shape (s, b, n, n). + n (int): The number of hyper connections (only n=4 is supported). + use_tf32 (bool): Whether to use TF32 precision for matmul operations. + + Returns: + tensor: The expanded and combined output of shape (s, b, C, n). + """ + + x = x.contiguous() + f = f.contiguous() + if bias is not None: + bias = bias.contiguous() + H_post = H_post.contiguous() + H_res = H_res.contiguous() + + s, b, C, n = x.shape + Cn = C * n + M = s * b + + out = torch.empty((s, b, C, n), device=x.device, dtype=x.dtype) + + # pylint: disable=unnecessary-lambda-assignment + grid = lambda META: ( + triton.cdiv(C, META["BLOCK_SIZE_C"]), + triton.cdiv(M, META["BLOCK_SIZE_M"]), + ) + + if bias is None: + _mhc_expand_combine_fwd[grid]( + f_ptr=f, + H_post_ptr=H_post, + x_ptr=x, + H_res_ptr=H_res, + output_ptr=out, + M=M, + C=C, + n=n, + stride_fm=C, + stride_fc=1, + stride_xm=Cn, + stride_xCn=1, + stride_output_m=Cn, + stride_output_Cn=1, + ) + else: + _mhc_expand_combine_with_bias_fwd[grid]( + f_ptr=f, + bias_ptr=bias, + H_post_ptr=H_post, + x_ptr=x, + H_res_ptr=H_res, + output_ptr=out, + M=M, + C=C, + n=n, + stride_fm=C, + stride_fc=1, + stride_bias=1, + stride_xm=Cn, + stride_xCn=1, + stride_output_m=Cn, + stride_output_Cn=1, + ) + + ctx.n = n + ctx.have_bias = bias is not None + if bias is not None: + ctx.save_for_backward(f, bias, H_post, x, H_res) + else: + ctx.save_for_backward(f, H_post, x, H_res) + ctx.use_tf32 = use_tf32 + + return out + + @staticmethod + def backward(ctx, grad_output): + """ + The backward pass of the expand and combine operation. Computes gradients for f, bias, + H_post, x, and H_res by backpropagating through the outer product and matrix multiply: + + grad_f = sum_n(grad_output * H_post) [+ reduce grad_bias over (s, b)] + grad_H_post[:, :, i] = sum_C(grad_output[:, :, :, i] * (f [+ bias])) + grad_x = grad_output @ H_res^T + grad_H_res[:, :, i, j] = sum_C(grad_output[:, :, :, j] * x[:, :, :, i]) + + Parameters: + ctx : The context object with saved tensors. + grad_output (tensor): The gradient of the loss with respect to the output, of shape (s, b, C, n). + + Returns: + tuple: A tuple with the gradients (grad_f, grad_bias, grad_H_post, grad_x, grad_H_res, None, None). + """ + grad_output = grad_output.contiguous() + s, b, C, n = grad_output.shape + + if ctx.have_bias: + f, bias, H_post, x, H_res = ctx.saved_tensors + else: + bias = None + f, H_post, x, H_res = ctx.saved_tensors + M = s * b + + grad_f = torch.empty_like(f) + grad_bias = torch.zeros_like(bias, dtype=torch.float32) if bias is not None else None + grad_H_post = torch.zeros_like( + H_post, dtype=torch.float32 + ) # We need to use atomic_add for this so we need higher precision + grad_x = torch.empty_like(x) + grad_H_res = torch.zeros_like( + H_res, dtype=torch.float32 + ) # We need to use atomic_add for this so we need higher precision + + # pylint: disable=unnecessary-lambda-assignment + grid = lambda META: ( + triton.cdiv(C, META["BLOCK_SIZE_C"]), + triton.cdiv(M, META["BLOCK_SIZE_M"]), + ) + + if bias is None: + _mhc_expand_combine_bwd[grid]( + grad_output_ptr=grad_output, + f_ptr=f, + H_post_ptr=H_post, + x_ptr=x, + H_res_ptr=H_res, + grad_H_post_ptr=grad_H_post, + grad_f_ptr=grad_f, + grad_H_res_ptr=grad_H_res, + grad_x_ptr=grad_x, + M=M, + C=C, + n=n, + stride_grad_output_m=n * C, + stride_grad_output_Cn=1, + stride_fm=C, + stride_fc=1, + stride_xm=n * C, + stride_xCn=1, + stride_grad_fm=C, + stride_grad_fc=1, + stride_grad_xm=n * C, + stride_grad_xCn=1, + precision="tf32" if ctx.use_tf32 else "ieee", + ) + else: + _mhc_expand_combine_with_bias_bwd[grid]( + grad_output_ptr=grad_output, + f_ptr=f, + bias_ptr=bias, + H_post_ptr=H_post, + x_ptr=x, + H_res_ptr=H_res, + grad_H_post_ptr=grad_H_post, + grad_f_ptr=grad_f, + grad_bias_ptr=grad_bias, + grad_H_res_ptr=grad_H_res, + grad_x_ptr=grad_x, + M=M, + C=C, + n=n, + stride_grad_output_m=n * C, + stride_grad_output_Cn=1, + stride_fm=C, + stride_fc=1, + stride_bias=1, + stride_xm=n * C, + stride_xCn=1, + stride_grad_fm=C, + stride_grad_fc=1, + stride_grad_bias=1, + stride_grad_xm=n * C, + stride_grad_xCn=1, + precision="tf32" if ctx.use_tf32 else "ieee", + ) + + grad_H_post = grad_H_post.to(H_post.dtype) # Cast back to the original dtype of H_post + grad_H_res = grad_H_res.to(H_res.dtype) # Cast back to the original dtype of H_res + if bias is not None: + grad_bias = grad_bias.to(bias.dtype) + + return grad_f, grad_bias, grad_H_post, grad_x, grad_H_res, None, None From b4aeed187f6b8592e6220b7d88962e8544bc3437 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Tue, 28 Apr 2026 19:38:07 -0700 Subject: [PATCH 021/180] [PyTorch] Main_Grad buffer isnt overwritten when overwrite_main_grad=True (#2936) * fix Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add test Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cleanup Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * zero_out should also be tested Signed-off-by: Varun Thumbe --------- Signed-off-by: Varun Thumbe Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: root --- tests/pytorch/test_fusible_ops.py | 299 +++++++++++++++--- .../pytorch/ops/fused/backward_grouped_mlp.py | 21 +- 2 files changed, 259 insertions(+), 61 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index c73f56056..3d6fe704e 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -42,6 +42,7 @@ ) from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor from transformer_engine.pytorch.cpp_extensions.gemm import general_grouped_gemm_for_grouped_tensor +from transformer_engine.pytorch.module.base import get_dummy_wgrad import transformer_engine_torch as tex # Import utility functions @@ -199,6 +200,76 @@ def make_reference_and_test_tensors( return ref, test +class MegatronTrainingHelper: + """Test-side stand-in for the Megatron-Core DDP / MegatronFSDP wrapper. + Megatron's DDP wrapper (and MegatronFSDP) owns the per-parameter + ``main_grad`` buffer and the ``overwrite_main_grad`` / + ``grad_added_to_main_grad`` attributes that coordinate + ``fuse_wgrad_accumulation`` with TE modules. These helpers reproduce the + relevant slice of that protocol so TE tests can exercise the + accumulate-into-``main_grad`` code path without pulling in the full + Megatron-Core dependency. + """ + + @staticmethod + def init_main_grad_buffers( + weight_params: Iterable[torch.nn.Parameter], + *, + fill_value: float, + overwrite_main_grad: bool, + zero_out_wgrad: bool = False, + dtype: torch.dtype = torch.float32, + ) -> None: + """Allocate ``main_grad`` and stamp the wrapper attributes on each + param, mirroring what the Megatron DDP/FSDP wrapper does before + backward.""" + for wp in weight_params: + wp.main_grad = torch.full(wp.size(), fill_value, device=wp.device, dtype=dtype) + wp.overwrite_main_grad = overwrite_main_grad + wp.zero_out_wgrad = zero_out_wgrad + wp.grad_added_to_main_grad = False + + @staticmethod + def verify_main_grad_accumulation( + weight_params: Iterable[torch.nn.Parameter], + *, + expected_main_grads: Iterable[torch.Tensor], + rtol: float = 0.0, + atol: float = 0.0, + ) -> None: + """Check that backward produced what the Megatron wrapper expects: + each ``main_grad`` matches ``expected_main_grads``, + ``grad_added_to_main_grad`` was flipped to ``True`` so the wrapper's + post-backward hooks won't double-accumulate, and ``param.grad`` was + replaced by the cached dummy tensor (so a wrapper hook that did + ``main_grad += grad`` would be a no-op rather than double-counting). + """ + for wp, expected in zip(weight_params, expected_main_grads): + torch.testing.assert_close(wp.main_grad.to(expected), expected, rtol=rtol, atol=atol) + + assert wp.grad_added_to_main_grad is True, ( + "weight.grad_added_to_main_grad was not flipped to True; " + "the Megatron DDP/FSDP wrapper hook will double-accumulate." + ) + + # ``.grad`` should be the cached dummy tensor returned by + # ``get_dummy_wgrad`` -- shared storage, not the real wgrad. + expected_dummy = get_dummy_wgrad(list(wp.size()), wp.dtype) + assert ( + wp.grad is not None + ), "weight.grad is None; the Megatron protocol expects a dummy tensor stand-in here." + assert wp.grad.data_ptr() == expected_dummy.data_ptr(), ( + "weight.grad does not share storage with the cached dummy " + "wgrad; downstream wrapper hooks risk double-accumulating." + ) + if getattr(wp, "zero_out_wgrad", False): + assert torch.all(wp.grad == 0), ( + "weight.zero_out_wgrad=True but the dummy weight.grad " + "was not zeroed; downstream hooks reading .grad would " + "see stale bytes from the previous step." + ) + + class TestSequentialContainer: """Tests for sequential container""" @@ -3537,33 +3608,20 @@ def test_grouped_mlp( getattr(fc1, f"bias{group_idx}").copy_(fc1_bs_test[group_idx]) getattr(fc2, f"bias{group_idx}").copy_(fc2_bs_test[group_idx]) if accumulate_into_main_grad: + # 0.5 sentinel lets us reconstruct ``expected = ref_grad + 0.5`` + # below and detect a missed accumulation. + main_grad_sentinel = 0.5 if single_grouped_weight: - fc1.weight.main_grad = torch.full( - fc1.weight.size(), - 0.5, - device=device, - dtype=torch.float32, - ) - fc2.weight.main_grad = torch.full( - fc2.weight.size(), - 0.5, - device=device, - dtype=torch.float32, - ) + weight_params_for_main_grad = [fc1.weight, fc2.weight] else: - for group_idx in range(group_size): - getattr(fc1, f"weight{group_idx}").main_grad = torch.full( - getattr(fc1, f"weight{group_idx}").size(), - 0.5, - device=device, - dtype=torch.float32, - ) - getattr(fc2, f"weight{group_idx}").main_grad = torch.full( - getattr(fc2, f"weight{group_idx}").size(), - 0.5, - device=device, - dtype=torch.float32, - ) + weight_params_for_main_grad = [ + getattr(fc, f"weight{i}") for fc in (fc1, fc2) for i in range(group_size) + ] + MegatronTrainingHelper.init_main_grad_buffers( + weight_params_for_main_grad, + fill_value=main_grad_sentinel, + overwrite_main_grad=False, + ) del fc1_ws_test, fc1_bs_test, fc2_ws_test, fc2_bs_test # Fuse ops and perform forward and backward pass @@ -3639,32 +3697,24 @@ def test_grouped_mlp( fc1_w_ref_grad = torch.stack([w.grad for w in fc1_ws_ref], dim=0) fc2_w_ref_grad = torch.stack([w.grad for w in fc2_ws_ref], dim=0) if accumulate_into_main_grad: - if single_grouped_weight: - fc1_w_test_grad = fc1.weight.main_grad.to(dtype=torch.float64, device="cpu") - 0.5 - fc2_w_test_grad = fc2.weight.main_grad.to(dtype=torch.float64, device="cpu") - 0.5 - else: - fc1_w_test_grad = torch.stack( - [ - getattr(fc1, f"weight{group_idx}").main_grad.to( - dtype=torch.float64, device="cpu" - ) - - 0.5 - for group_idx in range(group_size) - ], - dim=0, - ) - fc2_w_test_grad = torch.stack( - [ - getattr(fc2, f"weight{group_idx}").main_grad.to( - dtype=torch.float64, device="cpu" - ) - - 0.5 - for group_idx in range(group_size) - ], - dim=0, - ) - assert_close(fc1_w_test_grad, fc1_w_ref_grad, **tols) - assert_close(fc2_w_test_grad, fc2_w_ref_grad, **tols) + # main_grad should accumulate the ref wgrad onto the 0.5 sentinel. + # Per-param expected views must line up with + # ``weight_params_for_main_grad`` registered above. + fc1_expected = ( + [fc1_w_ref_grad + main_grad_sentinel] + if single_grouped_weight + else [g + main_grad_sentinel for g in fc1_w_ref_grad] + ) + fc2_expected = ( + [fc2_w_ref_grad + main_grad_sentinel] + if single_grouped_weight + else [g + main_grad_sentinel for g in fc2_w_ref_grad] + ) + MegatronTrainingHelper.verify_main_grad_accumulation( + weight_params_for_main_grad, + expected_main_grads=fc1_expected + fc2_expected, + **tols, + ) elif single_grouped_weight: assert_close(fc1.weight.grad, fc1_w_ref_grad, **tols) assert_close(fc2.weight.grad, fc2_w_ref_grad, **tols) @@ -3884,6 +3934,153 @@ def _run_case(single_grouped_weight: bool) -> tuple[torch.Tensor, ...]: torch.testing.assert_close(fc1_db_false, fc1_db_true, **bias_tols) torch.testing.assert_close(fc2_db_false, fc2_db_true, **bias_tols) + @pytest.mark.parametrize("single_grouped_weight", (False, True)) + @pytest.mark.parametrize("delay_wgrad_compute", (False, True)) + @pytest.mark.parametrize("zero_out_wgrad", (False, True)) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) + def test_grouped_mlp_overwrite_main_grad( + self, + *, + single_grouped_weight: bool, + delay_wgrad_compute: bool, + zero_out_wgrad: bool, + dtype: torch.dtype = torch.bfloat16, + device: torch.device = "cuda", + group_size: int = 4, + hidden_size: int = 256, + split_alignment: int = 256, + glu_interleave_size: int = 32, + ) -> None: + """End-to-end check that the fused grouped-MLP backward writes the + wgrad into ``weight.main_grad`` correctly under the MegatronFSDP + ``overwrite_main_grad=True`` convention. + ``test_grouped_mlp`` already covers the standard Megatron-LM + ``fuse_wgrad_accumulation`` (DDP) path where the wgrad GEMM + *accumulates* into ``main_grad``. This test focuses exclusively on + the MegatronFSDP variant where the wgrad GEMM must *overwrite* + ``main_grad`` (because FSDP has already ReduceScattered the previous + accumulation), so ``main_grad`` after backward equals ``wgrad`` + regardless of the prior contents. + + Also exercises the MegatronFSDP ``zero_out_wgrad`` flag, which is + independent of ``main_grad`` and only controls whether the dummy + ``param.grad`` returned to autograd is zeroed (so downstream hooks + that read ``.grad`` don't see stale bytes from the cached dummy). + """ + + if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported(): + pytest.skip("MXFP8 fused grouped MLP forward is not supported on this system") + if not te_ops.fused.BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8.is_supported(): + pytest.skip("MXFP8 fused grouped MLP backward is not supported on this system") + + recipe = make_recipe("mxfp8") + split_sizes = [split_alignment * (i + 1) for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int64, device=device) + in_shape = (split_sizes.sum().item(), hidden_size) + x_base = torch.empty(in_shape, device=device, dtype=dtype).uniform_(-0.25, 0.25) + probs_base = torch.empty((in_shape[0],), device=device, dtype=dtype).uniform_(-0.25, 0.25) + dy_base = torch.empty(in_shape, device=device, dtype=dtype).uniform_(-0.25, 0.25) + fc1_ws_base = [ + torch.empty((2 * hidden_size, hidden_size), device=device, dtype=dtype).uniform_( + -0.25, 0.25 + ) + for _ in range(group_size) + ] + fc2_ws_base = [ + torch.empty((hidden_size, hidden_size), device=device, dtype=dtype).uniform_( + -0.25, 0.25 + ) + for _ in range(group_size) + ] + + def _build_module(*, accumulate_into_main_grad: bool): + with te.quantized_model_init(enabled=True, recipe=recipe): + fc1 = te_ops.GroupedLinear( + group_size, + hidden_size, + 2 * hidden_size, + bias=False, + device=device, + dtype=dtype, + single_grouped_weight=single_grouped_weight, + accumulate_into_main_grad=accumulate_into_main_grad, + delay_wgrad_compute=delay_wgrad_compute, + ) + fc2 = te_ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=False, + device=device, + dtype=dtype, + single_grouped_weight=single_grouped_weight, + accumulate_into_main_grad=accumulate_into_main_grad, + delay_wgrad_compute=delay_wgrad_compute, + ) + scaled_act = te_ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) + module = te_ops.Sequential(fc1, scaled_act, fc2) + + with torch.no_grad(): + if single_grouped_weight: + fc1_weights = ( + fc1.weight.quantized_tensors or fc1.weight.split_into_quantized_tensors() + ) + fc2_weights = ( + fc2.weight.quantized_tensors or fc2.weight.split_into_quantized_tensors() + ) + for group_idx in range(group_size): + fc1_weights[group_idx].copy_(fc1_ws_base[group_idx]) + fc2_weights[group_idx].copy_(fc2_ws_base[group_idx]) + else: + for group_idx in range(group_size): + getattr(fc1, f"weight{group_idx}").copy_(fc1_ws_base[group_idx]) + getattr(fc2, f"weight{group_idx}").copy_(fc2_ws_base[group_idx]) + return module, fc1, fc2 + + def _weight_params(fc): + if single_grouped_weight: + return [fc.weight] + return [getattr(fc, f"weight{i}") for i in range(group_size)] + + def _run_backward(module, fc1, fc2): + x = x_base.detach().clone().requires_grad_(True) + probs = probs_base.detach().clone().requires_grad_(True) + with te.autocast(enabled=True, recipe=recipe): + y = module(x, split_sizes, probs, split_sizes) + y.backward(dy_base) + if delay_wgrad_compute: + fc1.backward_dw() + fc2.backward_dw() + + # Reference run: vanilla autograd, no Megatron protocol. + ref_module, ref_fc1, ref_fc2 = _build_module(accumulate_into_main_grad=False) + _run_backward(ref_module, ref_fc1, ref_fc2) + ref_fc1_grads = [wp.grad.detach().clone() for wp in _weight_params(ref_fc1)] + ref_fc2_grads = [wp.grad.detach().clone() for wp in _weight_params(ref_fc2)] + + # Test run: main_grad fusion with overwrite_main_grad=True (MegatronFSDP). + # NaN sentinel makes a missed write loud (would surface as NaN diff). + test_module, test_fc1, test_fc2 = _build_module(accumulate_into_main_grad=True) + for fc in (test_fc1, test_fc2): + MegatronTrainingHelper.init_main_grad_buffers( + _weight_params(fc), + fill_value=float("nan"), + overwrite_main_grad=True, + zero_out_wgrad=zero_out_wgrad, + ) + _run_backward(test_module, test_fc1, test_fc2) + + # main_grad must be overwritten to exactly the ref wgrad (bitwise: + # the wgrad GEMM is deterministic across the two runs because the + # quantized weights and inputs are identical). + MegatronTrainingHelper.verify_main_grad_accumulation( + _weight_params(test_fc1), expected_main_grads=ref_fc1_grads + ) + MegatronTrainingHelper.verify_main_grad_accumulation( + _weight_params(test_fc2), expected_main_grads=ref_fc2_grads + ) + @pytest.mark.parametrize("dtype", _dtypes) @pytest.mark.parametrize("single_grouped_weight", (False, True)) @pytest.mark.parametrize("accumulate_into_main_grad", (False, True)) diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 29273a5b4..510fea0ed 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -173,13 +173,12 @@ def _compute_grad_params( f" {tuple(main_grad.stride())}" ) from e accumulate_into_main_grad = not getattr(weight_param, "overwrite_main_grad", False) - if accumulate_into_main_grad: - grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( - num_tensors=num_groups, - tensor_shape=weight_shape, - rowwise_data=main_grad, - dtype=main_grad.dtype, - ) + grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( + num_tensors=num_groups, + tensor_shape=weight_shape, + rowwise_data=main_grad, + dtype=main_grad.dtype, + ) if grouped_wgrad is None: grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes( @@ -237,7 +236,9 @@ def _compute_grad_params( packed_wgrad = None if not delay_wgrad: packed_wgrad = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) - if accumulate_into_main_grad and hasattr(weight_param, "grad_added_to_main_grad"): + if fc_op._accumulate_into_main_grad and hasattr( + weight_param, "grad_added_to_main_grad" + ): weight_param.grad_added_to_main_grad = True packed_wgrad = get_dummy_wgrad( list(weight_param.size()), @@ -246,9 +247,9 @@ def _compute_grad_params( ) w_list = [packed_wgrad] else: - if delay_wgrad or accumulate_into_main_grad: + if delay_wgrad or fc_op._accumulate_into_main_grad: w_list = [None] * num_groups - if accumulate_into_main_grad: + if fc_op._accumulate_into_main_grad: for idx in range(num_groups): wp = getattr(fc_op, f"weight{idx}") if hasattr(wp, "grad_added_to_main_grad"): From 01aef4fc721bd12fd09cd56d53a314aee1b953d6 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Wed, 29 Apr 2026 12:27:30 -0400 Subject: [PATCH 022/180] Correctly pad scaling factor inverses to satisfy cuteDSL requirements (#2924) * Fix contiguous path for k=2880 Signed-off-by: Kirthi Shankar Sivamani * format Signed-off-by: Kirthi Shankar Sivamani * Review suggestion from @Oleg-Goncharov Signed-off-by: Kirthi Shankar Sivamani * Add test for swizzle + padding fusion Signed-off-by: Kirthi Shankar Sivamani * Address review comments Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- tests/cpp/operator/test_swizzle.cu | 255 +++++++++++++- transformer_engine/common/common.h | 20 ++ transformer_engine/common/swizzle/swizzle.cu | 318 ++++++++++++------ .../pytorch/csrc/extensions/swizzle.cpp | 38 ++- 4 files changed, 516 insertions(+), 115 deletions(-) diff --git a/tests/cpp/operator/test_swizzle.cu b/tests/cpp/operator/test_swizzle.cu index 1ea82f19c..3fec5062f 100644 --- a/tests/cpp/operator/test_swizzle.cu +++ b/tests/cpp/operator/test_swizzle.cu @@ -248,11 +248,11 @@ void performTestGroupedSwizzleMXFP8(const int num_tensors, const size_t M, const const NVTEShape rs = input->rowwise_scale_inv_shape(); zero_scale_inv_padding(input->rowwise_cpu_scale_inv_ptr(), rs.data[0], rs.data[1], - M, (K + BLOCK_SIZE - 1) / BLOCK_SIZE); + M, divide_round_up(K, BLOCK_SIZE)); const NVTEShape cs = input->columnwise_scale_inv_shape(); zero_scale_inv_padding(input->columnwise_cpu_scale_inv_ptr(), cs.data[0], cs.data[1], - (M + BLOCK_SIZE - 1) / BLOCK_SIZE, K); + divide_round_up(M, BLOCK_SIZE), K); input->from_cpu(); input_ptrs.push_back(input.get()); @@ -444,11 +444,11 @@ void performTestGroupedSwizzleUnswizzleRoundtrip(const int num_tensors, const si const NVTEShape rs = orig->rowwise_scale_inv_shape(); zero_scale_inv_padding(orig->rowwise_cpu_scale_inv_ptr(), rs.data[0], rs.data[1], - M, (K + BLOCK_SIZE - 1) / BLOCK_SIZE); + M, divide_round_up(K, BLOCK_SIZE)); const NVTEShape cs = orig->columnwise_scale_inv_shape(); zero_scale_inv_padding(orig->columnwise_cpu_scale_inv_ptr(), cs.data[0], cs.data[1], - (M + BLOCK_SIZE - 1) / BLOCK_SIZE, K); + divide_round_up(M, BLOCK_SIZE), K); orig->from_cpu(); orig_ptrs.push_back(orig.get()); @@ -541,6 +541,253 @@ INSTANTIATE_TEST_SUITE_P( } ); +// Build a "compact" grouped MXFP8 scale_inv buffer for swizzle input. This is +// the layout produced by the grouped MXFP8 quantize kernel: the per-tensor +// stride is `M_per_tensor * padded_K` (rowwise) or `DIVUP(M,32) * padded_K_for_cols` +// (columnwise) -- i.e. NO per-tensor padding rows are inserted. The total buffer +// is rounded up at its very end to a multiple of 128 (rowwise) or 4 (columnwise) +// in the grouped first dim, matching what the C++ allocator hands out. +// +// Each tensor's compact scales are gathered from the unpadded-prefix rows of +// that tensor's per-tensor padded CPU scale buffer. +namespace { + +struct CompactScaleBuffer { + test::CudaPtr<> ptr; + size_t numel{0}; +}; + +CompactScaleBuffer gather_compact_grouped_scale( + const std::vector>& tensors, + size_t M_per_tensor, size_t K_per_tensor, bool rowwise) { + using namespace test; + constexpr size_t BLOCK = 32; + const size_t num_tensors = tensors.size(); + + size_t per_tensor_first_unpadded; + size_t per_tensor_last_padded; + size_t group_first_align; + if (rowwise) { + per_tensor_first_unpadded = M_per_tensor; + per_tensor_last_padded = + round_up_to_nearest_multiple(divide_round_up(K_per_tensor, BLOCK), 4); + group_first_align = 128; + } else { + per_tensor_first_unpadded = divide_round_up(M_per_tensor, BLOCK); + per_tensor_last_padded = round_up_to_nearest_multiple(K_per_tensor, 128); + group_first_align = 4; + } + + const size_t per_tensor_compact_numel = + per_tensor_first_unpadded * per_tensor_last_padded; + const size_t total_first = round_up_to_nearest_multiple( + num_tensors * per_tensor_first_unpadded, group_first_align); + const size_t total_numel = total_first * per_tensor_last_padded; + + std::vector host_buf(total_numel, 0); + for (size_t i = 0; i < num_tensors; ++i) { + tensors[i]->to_cpu(); + const NVTEShape padded_shape = rowwise ? tensors[i]->rowwise_scale_inv_shape() + : tensors[i]->columnwise_scale_inv_shape(); + NVTE_CHECK(padded_shape.data[1] == per_tensor_last_padded, + "Unexpected per-tensor padded last dim in compact gather."); + const uint8_t* src = rowwise + ? tensors[i]->rowwise_cpu_scale_inv_ptr() + : tensors[i]->columnwise_cpu_scale_inv_ptr(); + uint8_t* dst = host_buf.data() + i * per_tensor_compact_numel; + // Per-tensor padded buffer is row-major (padded_first, padded_last); copy + // only the first `per_tensor_first_unpadded` rows. + std::memcpy(dst, src, per_tensor_compact_numel); + } + + CompactScaleBuffer out; + out.ptr = cuda_alloc(total_numel); + NVTE_CHECK_CUDA(cudaMemcpy(out.ptr.get(), host_buf.data(), + total_numel, cudaMemcpyHostToDevice)); + out.numel = total_numel; + return out; +} + +} // namespace + +// Tests that grouped_swizzle_for_gemm correctly handles a COMPACT input +// scale_inv buffer (no per-tensor padding rows), producing an output in the +// per-tensor padded layout with padded regions zeroed out. This is the layout +// produced by the grouped MXFP8 quantize kernel; previously the swizzle kernel +// asserted the input matched the per-tensor padded packed size, which broke +// grouped MLP weights with M not a multiple of 128. +void performTestGroupedSwizzleMXFP8CompactInput(const int num_tensors, const size_t M, + const size_t K) { + using namespace transformer_engine; + using namespace test; + + std::vector> input_tensors; + std::vector> output_tensors; + std::vector input_ptrs, output_ptrs; + input_tensors.reserve(num_tensors); + output_tensors.reserve(num_tensors); + input_ptrs.reserve(num_tensors); + output_ptrs.reserve(num_tensors); + + constexpr size_t BLOCK_SIZE = 32; + const std::vector shape{M, K}; + for (int i = 0; i < num_tensors; ++i) { + auto input = std::make_unique("input_" + std::to_string(i), shape, + DType::kFloat8E4M3, true, true, + NVTE_MXFP8_1D_SCALING); + auto output = std::make_unique("output_" + std::to_string(i), shape, + DType::kFloat8E4M3, true, true, + NVTE_MXFP8_1D_SCALING); + fillUniform(input.get()); + fillUniform(output.get()); + + // Zero the per-tensor padded regions so the reference (which sees the + // padded layout) and the kernel (which sees the compact layout but writes + // zeros into output padding) agree byte-for-byte. + input->to_cpu(); + const NVTEShape rs = input->rowwise_scale_inv_shape(); + zero_scale_inv_padding(input->rowwise_cpu_scale_inv_ptr(), + rs.data[0], rs.data[1], + M, divide_round_up(K, BLOCK_SIZE)); + const NVTEShape cs = input->columnwise_scale_inv_shape(); + zero_scale_inv_padding(input->columnwise_cpu_scale_inv_ptr(), + cs.data[0], cs.data[1], + divide_round_up(M, BLOCK_SIZE), K); + input->from_cpu(); + + input_ptrs.push_back(input.get()); + output_ptrs.push_back(output.get()); + input_tensors.emplace_back(std::move(input)); + output_tensors.emplace_back(std::move(output)); + } + + // Build a per-tensor padded grouped output via the standard helper, and a + // compact-scale grouped input by overriding the scale_inv buffers of a + // padded grouped input with newly allocated compact buffers. + GroupedBuffers grouped_input = build_grouped_tensor(input_ptrs, NVTE_MXFP8_1D_SCALING); + GroupedBuffers grouped_output = build_grouped_tensor(output_ptrs, NVTE_MXFP8_1D_SCALING); + + CompactScaleBuffer compact_row = + gather_compact_grouped_scale(input_tensors, M, K, /*rowwise=*/true); + CompactScaleBuffer compact_col = + gather_compact_grouped_scale(input_tensors, M, K, /*rowwise=*/false); + + grouped_input.scale_inv = std::move(compact_row.ptr); + grouped_input.columnwise_scale_inv = std::move(compact_col.ptr); + { + NVTEShape s = nvte_make_shape(&compact_row.numel, 1); + NVTEBasicTensor t{grouped_input.scale_inv.get(), kNVTEFloat8E8M0, s}; + nvte_set_grouped_tensor_param(grouped_input.get_handle(), + kNVTEGroupedRowwiseScaleInv, &t, sizeof(t)); + } + { + NVTEShape s = nvte_make_shape(&compact_col.numel, 1); + NVTEBasicTensor t{grouped_input.columnwise_scale_inv.get(), kNVTEFloat8E8M0, s}; + nvte_set_grouped_tensor_param(grouped_input.get_handle(), + kNVTEGroupedColumnwiseScaleInv, &t, sizeof(t)); + } + + const uint8_t input_swizzled = 0; + nvte_set_grouped_tensor_param(grouped_input.get_handle(), + kNVTEGroupedWithGEMMSwizzledScales, + &input_swizzled, sizeof(input_swizzled)); + const uint8_t output_swizzled = 1; + nvte_set_grouped_tensor_param(grouped_output.get_handle(), + kNVTEGroupedWithGEMMSwizzledScales, + &output_swizzled, sizeof(output_swizzled)); + + const NVTEShape row_shape = input_tensors[0]->rowwise_scale_inv_shape(); + const NVTEShape col_shape = input_tensors[0]->columnwise_scale_inv_shape(); + const size_t row_numel = row_shape.data[0] * row_shape.data[1]; + const size_t col_numel = col_shape.data[0] * col_shape.data[1]; + + // Memset to a non-zero sentinel so we can detect kernel failures to write + // padded regions (those must be overwritten with zero by the kernel). + NVTE_CHECK_CUDA(cudaMemset(grouped_output.scale_inv.get(), 0xCD, + num_tensors * row_numel)); + NVTE_CHECK_CUDA(cudaMemset(grouped_output.columnwise_scale_inv.get(), 0xCD, + num_tensors * col_numel)); + + nvte_swizzle_grouped_scaling_factors(grouped_input.get_handle(), + grouped_output.get_handle(), 0); + cudaDeviceSynchronize(); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + std::vector output_row(num_tensors * row_numel); + std::vector output_col(num_tensors * col_numel); + NVTE_CHECK_CUDA(cudaMemcpy(output_row.data(), grouped_output.scale_inv.get(), + output_row.size(), cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA(cudaMemcpy(output_col.data(), + grouped_output.columnwise_scale_inv.get(), + output_col.size(), cudaMemcpyDeviceToHost)); + + std::vector ref_row(num_tensors * row_numel); + std::vector ref_col(num_tensors * col_numel); + for (int i = 0; i < num_tensors; ++i) { + compute_ref_swizzle<128, 4, true>( + input_tensors[i]->rowwise_cpu_scale_inv_ptr(), + ref_row.data() + i * row_numel, + row_shape.data[0], row_shape.data[1]); + compute_ref_swizzle<128, 4, false>( + input_tensors[i]->columnwise_cpu_scale_inv_ptr(), + ref_col.data() + i * col_numel, + col_shape.data[1], col_shape.data[0]); + } + + compareResults("grouped_swizzle_compact_rowwise", output_row.data(), + ref_row.data(), num_tensors * row_numel); + compareResults("grouped_swizzle_compact_colwise", output_col.data(), + ref_col.data(), num_tensors * col_numel); +} + +class SwizzleGroupedCompactInputTestSuite + : public ::testing::TestWithParam> {}; + +TEST_P(SwizzleGroupedCompactInputTestSuite, TestGroupedSwizzleMXFP8CompactInput) { + const auto num_tensors = std::get<0>(GetParam()); + const auto M = std::get<1>(GetParam()); + const auto K = std::get<2>(GetParam()); + performTestGroupedSwizzleMXFP8CompactInput(num_tensors, M, K); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + SwizzleGroupedCompactInputTestSuite, + ::testing::Values( + // Aligned M and K. Per-tensor compact stride == per-tensor padded stride, + // so the kernel may use either layout; serves as a sanity check that the + // compact-input plumbing doesn't regress aligned shapes. + std::make_tuple(3, 256, 256), + std::make_tuple(4, 128, 128), + // M NOT divisible by 128 (the original-bug case): per-tensor compact stride + // shrinks vs padded. We pick (num_tensors, M) so that BOTH + // round_up(N * M, 128) != N * round_up(M, 128) (rowwise) + // round_up(N * DIVUP(M,32), 4) != N * round_up(DIVUP(M,32),4) (colwise) + // i.e. compact_total != padded_total on either axis, so the kernel + // unambiguously detects the compact layout. + std::make_tuple(4, 200, 256), + std::make_tuple(4, 65, 256), + std::make_tuple(2, 2880, 2880), // shape from the originally failing workload + // K not divisible by 128 (DIVUP(K,32) padded up to a multiple of 4). + std::make_tuple(3, 256, 160), + std::make_tuple(2, 256, 96), + // Neither M nor K aligned. + std::make_tuple(4, 200, 160), + std::make_tuple(4, 33, 64), + std::make_tuple(2, 1, 32), + // num_tensors * M not aligned to 128 -> exercises trailing alignment slack + // at the end of the compact rowwise buffer. + std::make_tuple(3, 64, 128), + std::make_tuple(5, 33, 96) + ), + [](const testing::TestParamInfo& info) { + return "n" + std::to_string(std::get<0>(info.param)) + + "_M" + std::to_string(std::get<1>(info.param)) + + "_K" + std::to_string(std::get<2>(info.param)); + } +); + class UnswizzleGroupedTestSuite : public ::testing::TestWithParam> {}; diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 68aa0f4c5..c1b3f8f42 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -946,6 +946,26 @@ struct TypeInfo { } \ } +#define TRANSFORMER_ENGINE_VECTORIZED_LOAD_INTEGER_TYPE_SWITCH(INTEGER_ELTS_NUM, type, ...) \ + switch (INTEGER_ELTS_NUM) { \ + case 1: { \ + using type = int; \ + { __VA_ARGS__ } \ + } break; \ + case 2: { \ + using type = int2; \ + { __VA_ARGS__ } \ + } break; \ + case 4: { \ + using type = int4; \ + { __VA_ARGS__ } \ + } break; \ + default: { \ + NVTE_ERROR("Unsupported number of integer elements ", INTEGER_ELTS_NUM, \ + ". Expected one of: 1, 2, or 4."); \ + } \ + } + //////////////////////////////////////////////////////////////////////////////////////////////////// inline int log2_ceil(int value) { diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index de4fdbb04..ad4a13092 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -91,7 +91,11 @@ __device__ inline void regs_unshuffle_with_bit_shifts(LType* regs_vec) { for (int i = 0; i < kVectorSize; i++) regs[i] = new_regs[i]; } -template +// IS_PADDED_K / IS_PADDED_M select the boundary-block specialization at compile +// time so the inner load loop avoids the per-iteration runtime checks. The +// caller computes the runtime predicates from blockIdx/gridDim once per block +// (uniform across the block) and dispatches to the right specialization. +template __device__ void swizzle_col_scaling_kernel_impl(const void* input, void* output, const int M, const int K, const int original_M, const int original_K, const int bid_x, @@ -117,9 +121,6 @@ __device__ void swizzle_col_scaling_kernel_impl(const void* input, void* output, m_tiles_in_tb = (M_i32 / SF_TILE_DIM_M_I32 - 1) % m_tiles_in_tb + 1; } - bool padding_m = (bid_y == grid_dim_y - 1) && (original_M < M); - bool padding_k = (bid_x == grid_dim_x - 1) && (original_K < K); - const int input_offset = bid_x * TB_DIM * SF_TILE_DIM_K_I32 * M_i32 + bid_y * N_TILE_PER_TD * SF_TILE_DIM_M_I32; const int32_t* input_i32 = reinterpret_cast(input) + input_offset; @@ -132,19 +133,37 @@ __device__ void swizzle_col_scaling_kernel_impl(const void* input, void* output, extern __shared__ int slm[]; // load, global -> regs + // Each register read for a given i is along the M direction at K-coord + // (bid_x * TB_DIM * SF_TILE_DIM_K + threadIdx.y * SF_TILE_DIM_K + i). When that + // K-coord is past original_K, the entire register is out of the per-tensor data + // region (which may be the unpadded compact extent), so we must NOT issue the + // __ldg there -- it could read past the per-tensor buffer (and, for the last + // tensor in a grouped allocation, past the end of the allocation entirely). LType regs_vec[N_SF_PER_TD_PER_TILE]; if (threadIdx.x * N_TILE_PER_TD < m_tiles_in_tb * SF_TILE_DIM_M_I32 && threadIdx.y < k_tiles_in_tb) { + const int k_base = bid_x * TB_DIM * SF_TILE_DIM_K + threadIdx.y * SF_TILE_DIM_K; #pragma unroll for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) { const int thread_offset = (threadIdx.y * SF_TILE_DIM_K_I32 + i) * M_i32 + threadIdx.x * N_TILE_PER_TD; + const int k_coord = k_base + i; + if constexpr (IS_PADDED_K) { + if (k_coord >= original_K) { + // Entire register is past original_K: zero directly without loading. + uint8_t* zero_bytes = reinterpret_cast(regs_vec + i); +#pragma unroll + for (int j = 0; j < static_cast(sizeof(LType)); j++) zero_bytes[j] = 0; + continue; + } + } regs_vec[i] = __ldg(reinterpret_cast(input_i32 + thread_offset)); - // Pad zeros - if (padding_m || padding_k) { + // Per-byte M masking is still needed when only part of the register is past + // original_M (i.e. K-coord is in range but the M position spans the boundary). + if constexpr (IS_PADDED_M) { for (int j = 0; j < N_TILE_PER_TD * sizeof(int); j++) { const int index = (input_offset + thread_offset) * sizeof(int) + j; - if (index / M >= original_K || index % M >= original_M) { + if (index % M >= original_M) { reinterpret_cast(regs_vec + i)[j] = 0; } } @@ -183,12 +202,43 @@ __device__ void swizzle_col_scaling_kernel_impl(const void* input, void* output, } } +// Dispatch helper: pick the right (IS_PADDED_K, IS_PADDED_M) col-scaling impl +// specialization at runtime based on the per-block padding predicates. The +// branching here is uniform across all threads in the block, so the indirect +// path each block takes still inlines cleanly. +template +__device__ __forceinline__ void dispatch_swizzle_col_scaling_kernel_impl( + const void* input, void* output, const int M, const int K, const int original_M, + const int original_K, const int bid_x, const int bid_y, const int grid_dim_x, + const int grid_dim_y, const bool padding_k, const bool padding_m) { + if (padding_k && padding_m) { + swizzle_col_scaling_kernel_impl( + input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y); + } else if (padding_k) { + swizzle_col_scaling_kernel_impl( + input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y); + } else if (padding_m) { + swizzle_col_scaling_kernel_impl( + input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y); + } else { + swizzle_col_scaling_kernel_impl( + input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y); + } +} + template __global__ void __launch_bounds__(TB_DIM* TB_DIM) swizzle_col_scaling_kernel(const void* input, void* output, const int M, const int K, const int original_M, const int original_K) { - swizzle_col_scaling_kernel_impl( - input, output, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y); + const bool padding_m = (blockIdx.y == gridDim.y - 1) && (original_M < M); + const bool padding_k = (blockIdx.x == gridDim.x - 1) && (original_K < K); + dispatch_swizzle_col_scaling_kernel_impl( + input, output, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y, + padding_k, padding_m); } template @@ -224,7 +274,11 @@ __device__ inline void regs_unshuffle(LType* regs_vec) { for (int i = 0; i < kVectorSize; i++) ptr[i] = tmp[i]; } -template +// IS_PADDED_K / IS_PADDED_M select the boundary-block specialization at compile +// time so the inner load loop avoids the per-iteration runtime checks. The +// caller computes the runtime predicates from blockIdx/gridDim once per block +// (uniform across the block) and dispatches to the right specialization. +template __device__ void swizzle_row_scaling_kernel_impl(const void* input, void* output, const int M, const int K, const int original_M, const int original_K, const int bid_x, @@ -243,9 +297,6 @@ __device__ void swizzle_row_scaling_kernel_impl(const void* input, void* output, n_tiles_in_tb = (K_i32 - 1) % N_TILES_IN_TB + 1; } - bool padding_m = (bid_y == grid_dim_y - 1) && (original_M < M); - bool padding_k = (bid_x == grid_dim_x - 1) && (original_K < K); - const int input_offset = bid_y * SF_TILE_DIM_M_I32 * K_i32 + bid_x * N_TILES_IN_TB; const int* input_i32 = reinterpret_cast(input) + input_offset; int* output_i32 = reinterpret_cast(output) + bid_y * SF_TILE_DIM_M_I32 * K_i32 + @@ -254,17 +305,35 @@ __device__ void swizzle_row_scaling_kernel_impl(const void* input, void* output, extern __shared__ int4 slm_v4i[]; // load, global -> regs + // Each register read for a given i is along the K direction at row + // (bid_y * SF_TILE_DIM_M + i * TB_DIM + threadIdx.y). When that row is past + // original_M, the entire register is out of the per-tensor data region (which + // may be the unpadded compact extent), so we must NOT issue the __ldg there -- + // it could read past the per-tensor buffer (and, for the last tensor in a + // grouped allocation, past the end of the allocation entirely). LType regs_vec[N_SF_PER_TD_PER_TILE]; if (threadIdx.x * N_TILE_PER_TD < n_tiles_in_tb) { #pragma unroll for (int i = 0; i < N_SF_PER_TD_PER_TILE; i++) { + const int row = bid_y * SF_TILE_DIM_M + i * TB_DIM + threadIdx.y; const int thread_offset = (i * TB_DIM + threadIdx.y) * K_i32 + threadIdx.x * N_TILE_PER_TD; + if constexpr (IS_PADDED_M) { + if (row >= original_M) { + // Entire register is past original_M: zero directly without loading. + uint8_t* zero_bytes = reinterpret_cast(regs_vec + i); +#pragma unroll + for (int j = 0; j < static_cast(sizeof(LType)); j++) zero_bytes[j] = 0; + continue; + } + } regs_vec[i] = __ldg(reinterpret_cast(input_i32 + thread_offset)); - if (padding_m || padding_k) { - // Pad zeros + // Per-byte K masking is still needed when only part of the register is past + // original_K (i.e. row is in range but the K position spans the boundary). + if constexpr (IS_PADDED_K) { +#pragma unroll for (int j = 0; j < N_TILE_PER_TD * sizeof(int); j++) { const int index = (input_offset + thread_offset) * sizeof(int) + j; - if (index / K >= original_M || index % K >= original_K) { + if (index % K >= original_K) { reinterpret_cast(regs_vec + i)[j] = 0; } } @@ -293,12 +362,43 @@ __device__ void swizzle_row_scaling_kernel_impl(const void* input, void* output, } } +// Dispatch helper: pick the right (IS_PADDED_K, IS_PADDED_M) row-scaling impl +// specialization at runtime based on the per-block padding predicates. The +// branching here is uniform across all threads in the block, so the indirect +// path each block takes still inlines cleanly. +template +__device__ __forceinline__ void dispatch_swizzle_row_scaling_kernel_impl( + const void* input, void* output, const int M, const int K, const int original_M, + const int original_K, const int bid_x, const int bid_y, const int grid_dim_x, + const int grid_dim_y, const bool padding_k, const bool padding_m) { + if (padding_k && padding_m) { + swizzle_row_scaling_kernel_impl( + input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y); + } else if (padding_k) { + swizzle_row_scaling_kernel_impl( + input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y); + } else if (padding_m) { + swizzle_row_scaling_kernel_impl( + input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y); + } else { + swizzle_row_scaling_kernel_impl( + input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y); + } +} + template __global__ void __launch_bounds__(TB_DIM* TB_DIM) swizzle_row_scaling_kernel(const void* input, void* output, const int M, const int K, const int original_M, const int original_K) { - swizzle_row_scaling_kernel_impl( - input, output, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y); + const bool padding_m = (blockIdx.y == gridDim.y - 1) && (original_M < M); + const bool padding_k = (blockIdx.x == gridDim.x - 1) && (original_K < K); + dispatch_swizzle_row_scaling_kernel_impl( + input, output, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, gridDim.y, + padding_k, padding_m); } // Narrow-K specialization for row scaling swizzle. @@ -628,14 +728,21 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) grouped_swizzle_row_scaling_uniform_shape_kernel(const void* input, void* output, const int M, const int K, const int original_M, const int original_K, - const size_t scale_stride_bytes) { + const size_t input_stride_bytes, + const size_t output_stride_bytes) { const int tensor_id = blockIdx.z; + // Input and output strides may differ: input is in the kernel-produced "compact" + // layout (per-tensor stride = original_M * padded_k * elem_size) when callers + // pass the unswizzled grouped scale buffer as-is, while the output is always in + // the per-tensor padded ("swizzle-ready") layout (padded_m * padded_k * elem_size). const uint8_t* input_base = - reinterpret_cast(input) + tensor_id * scale_stride_bytes; - uint8_t* output_base = reinterpret_cast(output) + tensor_id * scale_stride_bytes; - swizzle_row_scaling_kernel_impl( + reinterpret_cast(input) + tensor_id * input_stride_bytes; + uint8_t* output_base = reinterpret_cast(output) + tensor_id * output_stride_bytes; + const bool padding_m = (blockIdx.y == gridDim.y - 1) && (original_M < M); + const bool padding_k = (blockIdx.x == gridDim.x - 1) && (original_K < K); + dispatch_swizzle_row_scaling_kernel_impl( input_base, output_base, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, - gridDim.y); + gridDim.y, padding_k, padding_m); } template @@ -643,14 +750,20 @@ __global__ void __launch_bounds__(TB_DIM* TB_DIM) grouped_swizzle_col_scaling_uniform_shape_kernel(const void* input, void* output, const int M, const int K, const int original_M, const int original_K, - const size_t scale_stride_bytes) { + const size_t input_stride_bytes, + const size_t output_stride_bytes) { const int tensor_id = blockIdx.z; + // See the rowwise kernel for stride semantics. For columnwise the per-tensor + // compact stride is DIVUP(original_K, 1) * padded_m * elem_size (i.e. the + // unpadded scale-row count in the K direction times the padded M extent). const uint8_t* input_base = - reinterpret_cast(input) + tensor_id * scale_stride_bytes; - uint8_t* output_base = reinterpret_cast(output) + tensor_id * scale_stride_bytes; - swizzle_col_scaling_kernel_impl( + reinterpret_cast(input) + tensor_id * input_stride_bytes; + uint8_t* output_base = reinterpret_cast(output) + tensor_id * output_stride_bytes; + const bool padding_m = (blockIdx.y == gridDim.y - 1) && (original_M < M); + const bool padding_k = (blockIdx.x == gridDim.x - 1) && (original_K < K); + dispatch_swizzle_col_scaling_kernel_impl( input_base, output_base, M, K, original_M, original_K, blockIdx.x, blockIdx.y, gridDim.x, - gridDim.y); + gridDim.y, padding_k, padding_m); } template @@ -751,8 +864,11 @@ __global__ void multi_tensor_swizzle_row_scaling_kernel(MultiSwizzleArgs kernel_ const int bid_x = (bid - kernel_args.block_range[tensor_id]) / grid_dim_y; const int bid_y = (bid - kernel_args.block_range[tensor_id]) % grid_dim_y; - swizzle_row_scaling_kernel_impl( - input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y); + const bool padding_m = (bid_y == grid_dim_y - 1) && (original_M < M); + const bool padding_k = (bid_x == grid_dim_x - 1) && (original_K < K); + dispatch_swizzle_row_scaling_kernel_impl( + input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y, padding_k, + padding_m); } template @@ -781,8 +897,11 @@ __global__ void multi_tensor_swizzle_col_scaling_kernel(MultiSwizzleArgs kernel_ const int bid_x = (bid - kernel_args.block_range[tensor_id]) / grid_dim_y; const int bid_y = (bid - kernel_args.block_range[tensor_id]) % grid_dim_y; - swizzle_col_scaling_kernel_impl( - input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y); + const bool padding_m = (bid_y == grid_dim_y - 1) && (original_M < M); + const bool padding_k = (bid_x == grid_dim_x - 1) && (original_K < K); + dispatch_swizzle_col_scaling_kernel_impl( + input, output, M, K, original_M, original_K, bid_x, bid_y, grid_dim_x, grid_dim_y, padding_k, + padding_m); } template @@ -1924,23 +2043,56 @@ void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* const size_t padded_m = round_up_to_multiple(m, 128); const size_t padded_k = round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); - const size_t scale_elems = padded_m * padded_k; + // Per-tensor scale-element counts: + // - "padded" layout: each tensor occupies padded_m * padded_k elements + // (total buffer = num_tensors * padded_m * padded_k). + // - "compact" layout (what the grouped MXFP8 quantize kernel actually writes): + // per-tensor stride is m * padded_k (rowwise) or DIVUP(k,32) * padded_m + // (columnwise) and the total buffer the C++ allocator hands out has its + // grouped first dim padded up to a multiple of 128 (rowwise) or 4 + // (columnwise) -- so the buffer may be slightly larger than + // num_tensors * compact_scale_elems, with trailing alignment slack at + // the very end (never read because of the per-tensor row/k guard in the + // kernel impl). + // The output is always written in the padded layout. The input may be in + // either layout; the kernel handles the compact case safely by using + // different per-tensor strides for input vs output and skipping loads past + // the per-tensor extent. + const size_t padded_scale_elems = padded_m * padded_k; + const size_t compact_scale_elems = + rowwise ? m * padded_k : DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)) * padded_m; + const size_t compact_total_scale_elems = + rowwise ? round_up_to_multiple(input->num_tensors * m, 128) * padded_k + : round_up_to_multiple( + input->num_tensors * DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4) * + padded_m; const size_t scale_elem_size = rowwise ? typeToSize(input->scale_inv.dtype) : typeToSize(input->columnwise_scale_inv.dtype); - const size_t scale_stride_bytes = scale_elems * scale_elem_size; - if (rowwise) { - NVTE_CHECK(input->scale_inv.numel() == input->num_tensors * scale_elems, - "Grouped input scale_inv size does not match expected packed size."); - NVTE_CHECK(output->scale_inv.numel() == output->num_tensors * scale_elems, - "Grouped output scale_inv size does not match expected packed size."); + const size_t input_scale_numel = + rowwise ? input->scale_inv.numel() : input->columnwise_scale_inv.numel(); + const size_t output_scale_numel = + rowwise ? output->scale_inv.numel() : output->columnwise_scale_inv.numel(); + + bool input_is_compact; + if (input_scale_numel == input->num_tensors * padded_scale_elems) { + input_is_compact = false; + } else if (input_scale_numel == compact_total_scale_elems) { + input_is_compact = true; } else { - NVTE_CHECK(input->columnwise_scale_inv.numel() == input->num_tensors * scale_elems, - "Grouped input columnwise_scale_inv size does not match expected packed size."); - NVTE_CHECK(output->columnwise_scale_inv.numel() == output->num_tensors * scale_elems, - "Grouped output columnwise_scale_inv size does not match expected packed size."); + NVTE_ERROR("Grouped input ", (rowwise ? "scale_inv" : "columnwise_scale_inv"), + " size does not match expected packed size (got ", input_scale_numel, + ", expected either ", input->num_tensors * padded_scale_elems, + " (per-tensor padded) or ", compact_total_scale_elems, " (compact))."); } + NVTE_CHECK(output_scale_numel == input->num_tensors * padded_scale_elems, "Grouped output ", + (rowwise ? "scale_inv" : "columnwise_scale_inv"), + " size does not match expected per-tensor padded size."); + + const size_t input_stride_bytes = + (input_is_compact ? compact_scale_elems : padded_scale_elems) * scale_elem_size; + const size_t output_stride_bytes = padded_scale_elems * scale_elem_size; const int num_tiles_m = padded_m / SF_TILE_DIM_M; const int num_tiles_k = padded_k / SF_TILE_DIM_K; @@ -1963,69 +2115,25 @@ void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* void* output_ptr = rowwise ? output->scale_inv.dptr : output->columnwise_scale_inv.dptr; if (rowwise) { - switch (vec_load_size) { - case 4: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - grouped_swizzle_row_scaling_uniform_shape_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - grouped_swizzle_row_scaling_uniform_shape_kernel - <<>>(input_ptr, output_ptr, padded_m, - padded_k, original_M, original_K, - scale_stride_bytes); - break; - case 2: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - grouped_swizzle_row_scaling_uniform_shape_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - grouped_swizzle_row_scaling_uniform_shape_kernel - <<>>(input_ptr, output_ptr, padded_m, - padded_k, original_M, original_K, - scale_stride_bytes); - break; - case 1: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - grouped_swizzle_row_scaling_uniform_shape_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - grouped_swizzle_row_scaling_uniform_shape_kernel - <<>>(input_ptr, output_ptr, padded_m, - padded_k, original_M, original_K, - scale_stride_bytes); - break; - default: - NVTE_ERROR("Not valid vec_load_size."); - } + TRANSFORMER_ENGINE_VECTORIZED_LOAD_INTEGER_TYPE_SWITCH(vec_load_size, LType, { + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + grouped_swizzle_row_scaling_uniform_shape_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + grouped_swizzle_row_scaling_uniform_shape_kernel + <<>>(input_ptr, output_ptr, padded_m, + padded_k, original_M, original_K, + input_stride_bytes, output_stride_bytes); + }); } else { - switch (vec_load_size) { - case 4: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - grouped_swizzle_col_scaling_uniform_shape_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - grouped_swizzle_col_scaling_uniform_shape_kernel - <<>>(input_ptr, output_ptr, padded_m, - padded_k, original_M, original_K, - scale_stride_bytes); - break; - case 2: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - grouped_swizzle_col_scaling_uniform_shape_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - grouped_swizzle_col_scaling_uniform_shape_kernel - <<>>(input_ptr, output_ptr, padded_m, - padded_k, original_M, original_K, - scale_stride_bytes); - break; - case 1: - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - grouped_swizzle_col_scaling_uniform_shape_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - grouped_swizzle_col_scaling_uniform_shape_kernel - <<>>(input_ptr, output_ptr, padded_m, - padded_k, original_M, original_K, - scale_stride_bytes); - break; - default: - NVTE_ERROR("Not valid vec_load_size."); - } + TRANSFORMER_ENGINE_VECTORIZED_LOAD_INTEGER_TYPE_SWITCH(vec_load_size, LType, { + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + grouped_swizzle_col_scaling_uniform_shape_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + grouped_swizzle_col_scaling_uniform_shape_kernel + <<>>(input_ptr, output_ptr, padded_m, + padded_k, original_M, original_K, + input_stride_bytes, output_stride_bytes); + }); } NVTE_CHECK_CUDA(cudaGetLastError()); }; diff --git a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp index cbaabaad1..d8ab830c4 100644 --- a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp +++ b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp @@ -403,16 +403,39 @@ std::optional maybe_swizzle_grouped_tensor(GroupedTensorW tensor_offsets.data_ptr, static_cast(tensor_offsets.dtype), tensor_offsets.shape); } + // Per-tensor logical dimensions (uniform-shape grouped tensor). + const size_t num_tensors = input.num_tensors(); + const auto logical_shape_nvte = input.logical_shape(); + NVTE_CHECK(logical_shape_nvte.ndim >= 2, + "Grouped GEMM swizzle expects logical_shape with ndim >= 2."); + const size_t per_tensor_first_dim = logical_shape_nvte.data[0] / num_tensors; + const size_t per_tensor_last_dim = logical_shape_nvte.data[logical_shape_nvte.ndim - 1]; + constexpr size_t kMxfp8BlockSize = 32; + + // Output is always allocated in the per-tensor padded ("swizzle-ready") layout + // so the cuDNN grouped GEMM consumer sees the correct stride between experts. + // The swizzle kernel itself handles converting from the kernel-emitted compact + // layout (per-tensor first dim is the unpadded value) to this padded layout. + auto compute_padded_grouped_scale_shape = [&](bool rowwise) { + const size_t m = rowwise ? per_tensor_first_dim : per_tensor_last_dim; + const size_t k = rowwise ? per_tensor_last_dim : per_tensor_first_dim; + const size_t padded_m = ceildiv(m, size_t{128}) * 128; + const size_t padded_k = ceildiv(ceildiv(k, kMxfp8BlockSize), size_t{4}) * 4; + return std::vector{num_tensors * padded_m, padded_k}; + }; + if (swizzle_rowwise) { const auto data = input.get_rowwise_data(); const auto data_dtype = static_cast(data.dtype); const auto scales_dtype = static_cast(row_scales.dtype); swizzle_input.set_rowwise_data(nullptr, data_dtype, data.shape); swizzle_input.set_rowwise_scale_inv(row_scales.data_ptr, scales_dtype, row_scales.shape); - rowwise_scales_pyt = allocateSpace(row_scales.shape, scales_dtype, false); + const auto padded_shape = compute_padded_grouped_scale_shape(/*rowwise=*/true); + rowwise_scales_pyt = allocateSpace(padded_shape, scales_dtype, false); + NVTEShape padded_shape_nvte = nvte_make_shape(padded_shape.data(), padded_shape.size()); swizzle_output.set_rowwise_data(nullptr, data_dtype, data.shape); swizzle_output.set_rowwise_scale_inv(getDataPtr(*rowwise_scales_pyt), scales_dtype, - row_scales.shape); + padded_shape_nvte); } if (swizzle_columnwise) { const auto data = input.get_columnwise_data(); @@ -420,10 +443,12 @@ std::optional maybe_swizzle_grouped_tensor(GroupedTensorW const auto scales_dtype = static_cast(col_scales.dtype); swizzle_input.set_columnwise_data(nullptr, data_dtype, data.shape); swizzle_input.set_columnwise_scale_inv(col_scales.data_ptr, scales_dtype, col_scales.shape); - columnwise_scales_pyt = allocateSpace(col_scales.shape, scales_dtype, false); + const auto padded_shape = compute_padded_grouped_scale_shape(/*rowwise=*/false); + columnwise_scales_pyt = allocateSpace(padded_shape, scales_dtype, false); + NVTEShape padded_shape_nvte = nvte_make_shape(padded_shape.data(), padded_shape.size()); swizzle_output.set_columnwise_data(nullptr, data_dtype, data.shape); swizzle_output.set_columnwise_scale_inv(getDataPtr(*columnwise_scales_pyt), scales_dtype, - col_scales.shape); + padded_shape_nvte); } swizzle_output.set_with_gemm_swizzled_scales(true); @@ -434,12 +459,13 @@ std::optional maybe_swizzle_grouped_tensor(GroupedTensorW if (swizzle_rowwise) { const auto scales_dtype = static_cast(row_scales.dtype); - input.set_rowwise_scale_inv(getDataPtr(*rowwise_scales_pyt), scales_dtype, row_scales.shape); + input.set_rowwise_scale_inv(getDataPtr(*rowwise_scales_pyt), scales_dtype, + getTensorShape(*rowwise_scales_pyt)); } if (swizzle_columnwise) { const auto scales_dtype = static_cast(col_scales.dtype); input.set_columnwise_scale_inv(getDataPtr(*columnwise_scales_pyt), scales_dtype, - col_scales.shape); + getTensorShape(*columnwise_scales_pyt)); } input.set_with_gemm_swizzled_scales(true); return SwizzledGroupedScales{std::move(rowwise_scales_pyt), std::move(columnwise_scales_pyt)}; From cc05742f47e71d7ad7087ee0d73b3d263ba758d8 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 30 Apr 2026 15:20:06 -0700 Subject: [PATCH 023/180] [JAX] Fix bf16 precision loss in TestGroupedDense reference dbias (#2942) * accumulate bias in fp32 instead of bf16 in ref impl dbias to avoid accumulated numerical error Signed-off-by: tdophung --- tests/jax/test_custom_call_compute.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index d08f5cc11..14d28d95b 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -1914,11 +1914,24 @@ def test_grouped_gemm_fp8(self, fwd_bwd_dtype, scaling_mode, input_shape, layout self._assert_grouped_gemm_output(prim_out, group_sizes, ref_out, allclose_dtype) def _ref_sum_grouped_dense(self, x, kernel, bias, group_sizes, contracting_dims): - out_list = self._ref_grouped_dense(x, kernel, bias, group_sizes, contracting_dims) # Note: we use jnp.sum instead of jnp.mean to make the gradient larger # and prevent them from being clamp to zero in FP8. / sqrt(x.size) is used to # normalize the output and prevent the gradient from being too large for FP8. - out_sum_list = [jnp.sum(out) for out in out_list] + # + # We pass bias=None here and add bias externally in fp32 so the autodiff + # bias-grad (sum over the m axis of the cotangent) accumulates in fp32. + # If bias is added inside _ref_grouped_dense in bf16, JAX lowers the bias + # backward as a bf16 sum-over-m and loses precision on the largest group, + # producing a >bf16-rtol mismatch against the primitive's grouped_dbias + # (which casts the cotangent to fp32 before segment_sum). Bias is required + # for this helper since it is only used by the grad tests below, which all + # set with_bias=True. + assert bias is not None, "_ref_sum_grouped_dense requires a non-None bias" + out_list = self._ref_grouped_dense(x, kernel, None, group_sizes, contracting_dims) + out_sum_list = [] + for out_i, bias_i in zip(out_list, bias): + out_with_bias_fp32 = out_i.astype(jnp.float32) + bias_i.astype(jnp.float32) + out_sum_list.append(jnp.sum(out_with_bias_fp32)) return jnp.sum(jnp.asarray(out_sum_list)) / jnp.sqrt(x.size) def _primitive_sum_grouped_dense( @@ -1927,7 +1940,9 @@ def _primitive_sum_grouped_dense( out = grouped_dense( x, kernel, group_sizes, contracting_dims, bias=bias, quantizer_set=quantizer_set ) - return jnp.sum(jnp.asarray(out)) / jnp.sqrt(x.size) + # Match the fp32 accumulation in _ref_sum_grouped_dense so loss values are + # comparable and the cotangent dtype on `out` is unambiguous. + return jnp.sum(out.astype(jnp.float32)) / jnp.sqrt(x.size) @pytest_parametrize_wrapper("dtype", [jnp.bfloat16, jnp.float16]) def test_grouped_dense_grad_fp16(self, dtype, input_shape): From d156fa6c6a7fa9d48372a32f1b77b7a37c07db5d Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 30 Apr 2026 15:21:59 -0700 Subject: [PATCH 024/180] [JAX] Fix MNIST L2 jax test instability (#2933) * loosen up thresholds. Also only check min loss of last 10% of steps to avoid failing by noise near convergence Signed-off-by: tdophung * add deterministic flag for mnist run Signed-off-by: tdophung Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- examples/jax/mnist/test_single_gpu_mnist.py | 55 +++++++++++++++++---- qa/L2_jax_unittest/test.sh | 4 +- 2 files changed, 48 insertions(+), 11 deletions(-) diff --git a/examples/jax/mnist/test_single_gpu_mnist.py b/examples/jax/mnist/test_single_gpu_mnist.py index ef85f4a7a..5a058cbcc 100644 --- a/examples/jax/mnist/test_single_gpu_mnist.py +++ b/examples/jax/mnist/test_single_gpu_mnist.py @@ -3,6 +3,7 @@ # See LICENSE for license information. """MNIST training on single GPU""" import argparse +import math import unittest from functools import partial import sys @@ -223,6 +224,10 @@ def train_and_evaluate(args): print("PASSED") return None + train_losses = [] + train_accuracies = [] + test_losses = [] + test_accuracies = [] for epoch in range(1, args.epochs + 1): rng, input_rng = jax.random.split(rng) rng, dropout_rng = jax.random.split(rng) @@ -233,6 +238,11 @@ def train_and_evaluate(args): ) test_loss, test_accuracy = eval_model(state, test_ds, args.test_batch_size, var_collect) + train_losses.append(train_loss) + train_accuracies.append(train_accuracy) + test_losses.append(test_loss) + test_accuracies.append(test_accuracy) + print( f"Epoch: {epoch:>2} " f"Train Loss: {train_loss:.6f} " @@ -241,7 +251,7 @@ def train_and_evaluate(args): f"Test Accuracy: {test_accuracy:.6f} " ) - return [train_loss, train_accuracy, test_loss, test_accuracy] + return [train_losses, train_accuracies, test_losses, test_accuracies] def mnist_parser(args): @@ -324,15 +334,42 @@ def setUpClass(cls): @staticmethod def verify(actual): - """Check If loss and accuracy match target""" - desired_traing_loss = 0.055 + """Check that loss and accuracy match target. + + ``actual`` is ``[train_losses, train_accuracies, test_losses, test_accuracies]``, + i.e. per-epoch lists of metrics. To avoid flakiness from stochastic noise in + the final epoch near convergence (especially under FP8), the check considers + a tail window of the last ~10% of epochs (at least 2) and asserts on the + best metric within that window. + """ + train_losses, train_accuracies, test_losses, test_accuracies = actual + epochs = len(train_losses) + tail = max(2, math.ceil(epochs * 0.1)) + tail = min(tail, epochs) + + best_train_loss = min(train_losses[-tail:]) + best_train_accuracy = max(train_accuracies[-tail:]) + best_test_loss = min(test_losses[-tail:]) + best_test_accuracy = max(test_accuracies[-tail:]) + + desired_traing_loss = 0.06 desired_traing_accuracy = 0.98 - desired_test_loss = 0.045 - desired_test_accuracy = 0.098 - assert actual[0] < desired_traing_loss - assert actual[1] > desired_traing_accuracy - assert actual[2] < desired_test_loss - assert actual[3] > desired_test_accuracy + desired_test_loss = 0.05 + desired_test_accuracy = 0.98 + assert ( + best_train_loss < desired_traing_loss + ), f"best train loss over last {tail} epochs {best_train_loss} >= {desired_traing_loss}" + assert best_train_accuracy > desired_traing_accuracy, ( + f"best train accuracy over last {tail} epochs {best_train_accuracy} " + f"<= {desired_traing_accuracy}" + ) + assert ( + best_test_loss < desired_test_loss + ), f"best test loss over last {tail} epochs {best_test_loss} >= {desired_test_loss}" + assert best_test_accuracy > desired_test_accuracy, ( + f"best test accuracy over last {tail} epochs {best_test_accuracy} " + f"<= {desired_test_accuracy}" + ) @unittest.skipIf(not is_bf16_supported(), "Device compute capability 8.0+ is required for BF16") def test_te_bf16(self): diff --git a/qa/L2_jax_unittest/test.sh b/qa/L2_jax_unittest/test.sh index 582267566..8441486e2 100644 --- a/qa/L2_jax_unittest/test.sh +++ b/qa/L2_jax_unittest/test.sh @@ -31,11 +31,11 @@ mkdir -p "$XML_LOG_DIR" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_not_distributed.xml $TE_PATH/tests/jax -k 'not distributed' || test_fail "tests/jax/*not_distributed_*" pip3 install -r $TE_PATH/examples/jax/mnist/requirements.txt || error_exit "Failed to install mnist requirements" +# Make mnist and encoder tests run-to-run deterministic for stable CI results +export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_mnist.xml $TE_PATH/examples/jax/mnist || test_fail "mnist" pip3 install -r $TE_PATH/examples/jax/encoder/requirements.txt || error_exit "Failed to install encoder requirements" -# Make encoder tests to have run-to-run deterministic to have the stable CI results -export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_single_gpu_encoder.xml $TE_PATH/examples/jax/encoder/test_single_gpu_encoder.py || test_fail "test_single_gpu_encoder.py" # Test without custom calls export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops" From a7a2b3bbff4b9cc6487e5da4efbf4e463285f1cd Mon Sep 17 00:00:00 2001 From: int-smart Date: Thu, 30 Apr 2026 16:51:07 -0700 Subject: [PATCH 025/180] Variable Grouped Swizzle (#2914) * feat: add support for grouped GEMM swizzling with variable shapes and update C++ operator interface Signed-off-by: Abhishek * Added confirmation with uniformity in one of the dimensions Signed-off-by: Abhishek * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Using single kernel for variable m and k Signed-off-by: Abhishek * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Cached blocks per sm for device and removed redundant checks Signed-off-by: Abhishek * Updated the code with newer changes in main Signed-off-by: Abhishek * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Abhishek Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: vthumbe1503 --- tests/cpp/operator/test_swizzle.cu | 136 ++++++ tests/cpp/test_common.cu | 14 +- transformer_engine/common/swizzle/swizzle.cu | 425 +++++++++++++----- .../pytorch/csrc/extensions/swizzle.cpp | 8 +- 4 files changed, 456 insertions(+), 127 deletions(-) diff --git a/tests/cpp/operator/test_swizzle.cu b/tests/cpp/operator/test_swizzle.cu index 3fec5062f..8990ce8db 100644 --- a/tests/cpp/operator/test_swizzle.cu +++ b/tests/cpp/operator/test_swizzle.cu @@ -506,6 +506,142 @@ void performTestGroupedSwizzleUnswizzleRoundtrip(const int num_tensors, const si num_tensors * col_numel); } +void performTestGroupedSwizzleMXFP8Variable(const std::vector>& shapes) { + using namespace transformer_engine; + using namespace test; + + int num_tensors = shapes.size(); + std::vector> input_tensors; + std::vector> output_tensors; + std::vector input_ptrs; + std::vector output_ptrs; + input_tensors.reserve(num_tensors); + output_tensors.reserve(num_tensors); + input_ptrs.reserve(num_tensors); + output_ptrs.reserve(num_tensors); + + constexpr size_t BLOCK_SIZE = 32; + for (int i = 0; i < num_tensors; ++i) { + const std::vector shape{shapes[i].first, shapes[i].second}; + auto input = std::make_unique("input_" + std::to_string(i), shape, + DType::kFloat8E4M3, true, true, + NVTE_MXFP8_1D_SCALING); + auto output = std::make_unique("output_" + std::to_string(i), shape, + DType::kFloat8E4M3, true, true, + NVTE_MXFP8_1D_SCALING); + fillUniform(input.get()); + fillUniform(output.get()); + + // Zero padding + input->to_cpu(); + const NVTEShape rs = input->rowwise_scale_inv_shape(); + zero_scale_inv_padding(input->rowwise_cpu_scale_inv_ptr(), + rs.data[0], rs.data[1], + shapes[i].first, (shapes[i].second + BLOCK_SIZE - 1) / BLOCK_SIZE); + const NVTEShape cs = input->columnwise_scale_inv_shape(); + zero_scale_inv_padding(input->columnwise_cpu_scale_inv_ptr(), + cs.data[0], cs.data[1], + (shapes[i].first + BLOCK_SIZE - 1) / BLOCK_SIZE, shapes[i].second); + input->from_cpu(); + + input_ptrs.push_back(input.get()); + output_ptrs.push_back(output.get()); + input_tensors.emplace_back(std::move(input)); + output_tensors.emplace_back(std::move(output)); + } + + GroupedBuffers grouped_input = build_grouped_tensor(input_ptrs, NVTE_MXFP8_1D_SCALING); + GroupedBuffers grouped_output = build_grouped_tensor(output_ptrs, NVTE_MXFP8_1D_SCALING); + + const uint8_t input_swizzled = 0; + nvte_set_grouped_tensor_param(grouped_input.get_handle(), + kNVTEGroupedWithGEMMSwizzledScales, + &input_swizzled, sizeof(input_swizzled)); + const uint8_t output_swizzled = 1; + nvte_set_grouped_tensor_param(grouped_output.get_handle(), + kNVTEGroupedWithGEMMSwizzledScales, + &output_swizzled, sizeof(output_swizzled)); + + nvte_swizzle_grouped_scaling_factors(grouped_input.get_handle(), + grouped_output.get_handle(), + 0); + + cudaDeviceSynchronize(); + NVTE_CHECK_CUDA(cudaGetLastError()); + + // Verification + size_t row_offset = 0; + size_t col_offset = 0; + for (int i = 0; i < num_tensors; ++i) { + const NVTEShape row_shape = input_tensors[i]->rowwise_scale_inv_shape(); + const NVTEShape col_shape = input_tensors[i]->columnwise_scale_inv_shape(); + const size_t row_numel = row_shape.data[0] * row_shape.data[1]; + const size_t col_numel = col_shape.data[0] * col_shape.data[1]; + + std::vector output_row_host(row_numel); + std::vector output_col_host(col_numel); + NVTE_CHECK_CUDA(cudaMemcpy(output_row_host.data(), + static_cast(grouped_output.scale_inv.get()) + row_offset, + row_numel, cudaMemcpyDeviceToHost)); + NVTE_CHECK_CUDA(cudaMemcpy(output_col_host.data(), + static_cast(grouped_output.columnwise_scale_inv.get()) + col_offset, + col_numel, cudaMemcpyDeviceToHost)); + + std::vector ref_row(row_numel); + std::vector ref_col(col_numel); + compute_ref_swizzle<128, 4, true>(input_tensors[i]->rowwise_cpu_scale_inv_ptr(), + ref_row.data(), + row_shape.data[0], row_shape.data[1]); + compute_ref_swizzle<128, 4, false>( + input_tensors[i]->columnwise_cpu_scale_inv_ptr(), + ref_col.data(), + col_shape.data[1], col_shape.data[0]); + + compareResults("grouped_swizzle_variable_rowwise_" + std::to_string(i), + output_row_host.data(), ref_row.data(), row_numel); + compareResults("grouped_swizzle_variable_colwise_" + std::to_string(i), + output_col_host.data(), ref_col.data(), col_numel); + + row_offset += row_numel; + col_offset += col_numel; + } +} + +class SwizzleGroupedVariableTestSuite + : public ::testing::TestWithParam>> {}; + +TEST_P(SwizzleGroupedVariableTestSuite, TestGroupedSwizzleMXFP8Variable) { + const auto shapes = GetParam(); + performTestGroupedSwizzleMXFP8Variable(shapes); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + SwizzleGroupedVariableTestSuite, + ::testing::Values( + // Case 1: num_tensors = 1 (n+3 = 4, even). Check simple alignment. + std::vector>{{1024, 1024}}, + + // Case 2: num_tensors = 2 (n+3 = 5, odd). Forces padding logic to trigger. + std::vector>{{128, 128}, {256, 256}}, + + // Case 3: Mixed small/irregular shapes. + std::vector>{{200, 160}, {33, 64}, {1, 32}}, + + // Case 4: Large workload to verify persistent grid + std::vector>(10, {4096, 4096}), + + // Case 5: Variable M, Uniform K (Semi-variable) + std::vector>{{128, 256}, {512, 256}, {64, 256}}, + + // Case 6: Uniform M, Variable K (Semi-variable) + std::vector>{{512, 128}, {512, 1024}, {512, 32}} + ), + [](const testing::TestParamInfo& info) { + return "VariableShapes_" + std::to_string(info.index) + "_N" + std::to_string(info.param.size()); + } +); + class SwizzleGroupedTestSuite : public ::testing::TestWithParam> {}; diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index 519668411..b8bc38935 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -1099,7 +1099,7 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, const bool same_last = std::all_of(last_dims.begin(), last_dims.end(), [&](int64_t v) { return v == last_dims[0]; }); - std::vector offsets(num_tensors, 0); + std::vector offsets(num_tensors + 1, 0); auto random_padding = [&]() -> int64_t { // Random padding ensuring 16-byte alignment regardless of element size // cuBLAS requires aligned pointers for vectorized loads @@ -1118,12 +1118,11 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, const bool need_offsets = !same_first || !same_last; const bool use_random_padding = need_offsets && scaling_mode != NVTE_MXFP8_1D_SCALING; if (need_offsets) { - offsets[0] = 0; - for (size_t i = 1; i < num_tensors; ++i) { + for (size_t i = 1; i < num_tensors + 1; ++i) { offsets[i] = offsets[i - 1] + numel(i - 1) + (use_random_padding ? random_padding() : 0); } } else { - for (size_t i = 0; i < num_tensors; ++i) { + for (size_t i = 0; i < num_tensors + 1; ++i) { offsets[i] = static_cast(i) * numel(0); } } @@ -1211,10 +1210,11 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, } if (!same_first || !same_last) { - grouped.offsets_dev = cuda_alloc(num_tensors * sizeof(int64_t)); + size_t num_off = num_tensors + 1; + grouped.offsets_dev = cuda_alloc(num_off * sizeof(int64_t)); NVTE_CHECK_CUDA(cudaMemcpy(grouped.offsets_dev.get(), offsets.data(), - num_tensors * sizeof(int64_t), cudaMemcpyHostToDevice)); - NVTEShape off_shape = nvte_make_shape(&num_tensors, 1); + num_off * sizeof(int64_t), cudaMemcpyHostToDevice)); + NVTEShape off_shape = nvte_make_shape(&num_off, 1); NVTEBasicTensor off_tensor{grouped.offsets_dev.get(), kNVTEInt64, off_shape}; nvte_set_grouped_tensor_param(h, kNVTEGroupedTensorOffsets, &off_tensor, sizeof(off_tensor)); } diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index ad4a13092..c7ed407a5 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -8,10 +8,13 @@ #include #include +#include #include #include +#include #include "../common.h" +#include "../util/cuda_runtime.h" #include "../util/logging.h" #include "transformer_engine/transformer_engine.h" @@ -2001,6 +2004,154 @@ void nvte_multi_tensor_unswizzle_scaling_factors(const NVTETensor* inputs, NVTET namespace transformer_engine { +template +__global__ void __launch_bounds__(TB_DIM* TB_DIM) + grouped_swizzle_scaling_variable_shape_kernel(const void* input, void* output, + const int64_t* m_array, const int64_t* k_array, + int num_tensors, bool rowwise, + size_t scale_elem_size, size_t common_m, + size_t common_k) { + extern __shared__ int s_metadata[]; + int* s_total_blocks = &s_metadata[0]; + + // Warp reduction to compute total workload + if (threadIdx.x < 32 && threadIdx.y == 0) { + int local_blocks = 0; + for (int i = threadIdx.x; i < num_tensors; i += 32) { + size_t m = rowwise ? (m_array ? m_array[i] : common_m) : (k_array ? k_array[i] : common_k); + size_t k = rowwise ? (k_array ? k_array[i] : common_k) : (m_array ? m_array[i] : common_m); + + size_t padded_m = round_up_to_multiple(m, 128); + size_t padded_k = round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); + + int num_tiles_m = padded_m / SF_TILE_DIM_M; + int num_tiles_k = padded_k / SF_TILE_DIM_K; + + int vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1)); + if (vec_load_size == 3) vec_load_size = 1; + int n_tiles_in_tb = TB_DIM * vec_load_size; + + int grid_dim_x = rowwise ? DIVUP(num_tiles_k, n_tiles_in_tb) : DIVUP(num_tiles_k, TB_DIM); + int grid_dim_y = rowwise ? num_tiles_m : DIVUP(num_tiles_m, vec_load_size); + local_blocks += grid_dim_x * grid_dim_y; + } + + for (int offset = 16; offset > 0; offset /= 2) { + local_blocks += __shfl_down_sync(0xffffffff, local_blocks, offset); + } + if (threadIdx.x == 0) *s_total_blocks = local_blocks; + } + __syncthreads(); + + const int total_blocks = *s_total_blocks; + + // Persistent-grid loop + for (int linear_block_id = blockIdx.x; linear_block_id < total_blocks; + linear_block_id += gridDim.x) { + // Discover tensor_id and local_block_id via linear scan + int tensor_id = 0; + int current_block_base = 0; + size_t current_scale_base = 0; + int grid_dim_x = 0; + int grid_dim_y = 0; + size_t M = 0, K = 0; + int vec_load_size = 0; + + for (int i = 0; i < num_tensors; ++i) { + M = rowwise ? (m_array ? m_array[i] : common_m) : (k_array ? k_array[i] : common_k); + K = rowwise ? (k_array ? k_array[i] : common_k) : (m_array ? m_array[i] : common_m); + + size_t padded_m = round_up_to_multiple(M, 128); + size_t padded_k = round_up_to_multiple(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4); + + int num_tiles_m = padded_m / SF_TILE_DIM_M; + int num_tiles_k = padded_k / SF_TILE_DIM_K; + + vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1)); + if (vec_load_size == 3) vec_load_size = 1; + int n_tiles_in_tb = TB_DIM * vec_load_size; + + grid_dim_x = rowwise ? DIVUP(num_tiles_k, n_tiles_in_tb) : DIVUP(num_tiles_k, TB_DIM); + grid_dim_y = rowwise ? num_tiles_m : DIVUP(num_tiles_m, vec_load_size); + int blocks_i = grid_dim_x * grid_dim_y; + + if (linear_block_id < current_block_base + blocks_i) { + tensor_id = i; + break; + } + current_block_base += blocks_i; + current_scale_base += padded_m * padded_k * scale_elem_size; + } + + int local_block_id = linear_block_id - current_block_base; + int block_x = local_block_id % grid_dim_x; + int block_y = local_block_id / grid_dim_x; + + const uint8_t* input_base = reinterpret_cast(input) + current_scale_base; + uint8_t* output_base = reinterpret_cast(output) + current_scale_base; + + const int padded_m = static_cast(round_up_to_multiple(M, 128)); + const int padded_k = + static_cast(round_up_to_multiple(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE)), 4)); + const int original_M = static_cast(M); + const int original_K = static_cast(DIVUP(K, static_cast(MXFP8_BLOCK_SIZE))); + const bool padding_m = (block_y == grid_dim_y - 1) && (original_M < padded_m); + const bool padding_k = (block_x == grid_dim_x - 1) && (original_K < padded_k); + + if (rowwise) { + if (vec_load_size == 4) { + dispatch_swizzle_row_scaling_kernel_impl( + input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y, + grid_dim_x, grid_dim_y, padding_k, padding_m); + } else if (vec_load_size == 2) { + dispatch_swizzle_row_scaling_kernel_impl( + input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y, + grid_dim_x, grid_dim_y, padding_k, padding_m); + } else { + dispatch_swizzle_row_scaling_kernel_impl( + input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y, + grid_dim_x, grid_dim_y, padding_k, padding_m); + } + } else { + if (vec_load_size == 4) { + dispatch_swizzle_col_scaling_kernel_impl( + input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y, + grid_dim_x, grid_dim_y, padding_k, padding_m); + } else if (vec_load_size == 2) { + dispatch_swizzle_col_scaling_kernel_impl( + input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y, + grid_dim_x, grid_dim_y, padding_k, padding_m); + } else { + dispatch_swizzle_col_scaling_kernel_impl( + input_base, output_base, padded_m, padded_k, original_M, original_K, block_x, block_y, + grid_dim_x, grid_dim_y, padding_k, padding_m); + } + } + } +} + +template +int grouped_swizzle_variable_max_active_blocks_per_sm(int device_id) { + static std::vector cache(cuda::num_devices(), -1); + static std::vector flags(cuda::num_devices()); + NVTE_CHECK(0 <= device_id && device_id < cuda::num_devices(), "invalid CUDA device ID"); + + auto init = [&]() { + constexpr int metadata_shmem = sizeof(int); // s_total_blocks + constexpr int dynamic_smem_size = + TB_DIM * 4 * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t) + metadata_shmem; + int max_active_blocks_per_sm; + NVTE_CHECK_CUDA(cudaOccupancyMaxActiveBlocksPerMultiprocessor( + &max_active_blocks_per_sm, + grouped_swizzle_scaling_variable_shape_kernel, + TB_DIM * TB_DIM, dynamic_smem_size)); + NVTE_CHECK(max_active_blocks_per_sm > 0, "Occupancy query returned 0 blocks per SM."); + cache[device_id] = max_active_blocks_per_sm; + }; + std::call_once(flags[device_id], init); + return cache[device_id]; +} + void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* output, cudaStream_t stream) { // Check scaling mode @@ -2022,127 +2173,175 @@ void swizzle_grouped_scaling_factors(const GroupedTensor* input, GroupedTensor* return; } - // Only support uniform shapes for graph-safe grouped swizzle - NVTE_CHECK(input->all_same_shape(), "Grouped swizzle requires uniform tensor shapes."); - NVTE_CHECK(input->all_same_last_dim() && input->all_same_first_dim(), - "Grouped swizzle requires uniform tensor shapes."); + const int64_t* m_array = reinterpret_cast(input->first_dims.dptr); + const int64_t* k_array = reinterpret_cast(input->last_dims.dptr); + const bool is_variable_shape = !input->all_same_shape(); + + if (!is_variable_shape) { + // Fallback to uniform shape implementation + // Assumption is that all the tensors share the same shapes and are contgiuous. + // And so we dont need to pass array of input/output pointers(due to conttiguity) + // as well as array of shapes(due to uniform shapes). + const size_t first_dim = input->get_common_first_dim(); + const size_t last_dim = input->get_common_last_dim(); + + constexpr int SF_TILE_DIM_M = 128; + constexpr int SF_TILE_DIM_K = 4; + const dim3 block_size(TB_DIM, TB_DIM); + + auto launch_grouped_swizzle = [&](bool rowwise) { + const size_t m = rowwise ? first_dim : last_dim; + const size_t k = rowwise ? last_dim : first_dim; + const size_t padded_m = round_up_to_multiple(m, 128); + const size_t padded_k = + round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); + // Per-tensor scale-element counts: + // - "padded" layout: each tensor occupies padded_m * padded_k elements + // (total buffer = num_tensors * padded_m * padded_k). + // - "compact" layout (what the grouped MXFP8 quantize kernel actually writes): + // per-tensor stride is m * padded_k (rowwise) or DIVUP(k,32) * padded_m + // (columnwise) and the total buffer the C++ allocator hands out has its + // grouped first dim padded up to a multiple of 128 (rowwise) or 4 + // (columnwise) -- so the buffer may be slightly larger than + // num_tensors * compact_scale_elems, with trailing alignment slack at + // the very end (never read because of the per-tensor row/k guard in the + // kernel impl). + // The output is always written in the padded layout. The input may be in + // either layout; the kernel handles the compact case safely by using + // different per-tensor strides for input vs output and skipping loads past + // the per-tensor extent. + const size_t padded_scale_elems = padded_m * padded_k; + const size_t compact_scale_elems = + rowwise ? m * padded_k : DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)) * padded_m; + const size_t compact_total_scale_elems = + rowwise ? round_up_to_multiple(input->num_tensors * m, 128) * padded_k + : round_up_to_multiple( + input->num_tensors * DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4) * + padded_m; + + const size_t scale_elem_size = rowwise ? typeToSize(input->scale_inv.dtype) + : typeToSize(input->columnwise_scale_inv.dtype); + + const size_t input_scale_numel = + rowwise ? input->scale_inv.numel() : input->columnwise_scale_inv.numel(); + const size_t output_scale_numel = + rowwise ? output->scale_inv.numel() : output->columnwise_scale_inv.numel(); + + bool input_is_compact; + if (input_scale_numel == input->num_tensors * padded_scale_elems) { + input_is_compact = false; + } else if (input_scale_numel == compact_total_scale_elems) { + input_is_compact = true; + } else { + NVTE_ERROR("Grouped input ", (rowwise ? "scale_inv" : "columnwise_scale_inv"), + " size does not match expected packed size (got ", input_scale_numel, + ", expected either ", input->num_tensors * padded_scale_elems, + " (per-tensor padded) or ", compact_total_scale_elems, " (compact))."); + } + NVTE_CHECK(output_scale_numel == input->num_tensors * padded_scale_elems, "Grouped output ", + (rowwise ? "scale_inv" : "columnwise_scale_inv"), + " size does not match expected per-tensor padded size."); - // Assumption is that all the tensors share the same shapes and are contgiuous. - // And so we dont need to pass array of input/output pointers(due to conttiguity) - // as well as array of shapes(due to uniform shapes). - const size_t first_dim = input->get_common_first_dim(); - const size_t last_dim = input->get_common_last_dim(); + const size_t input_stride_bytes = + (input_is_compact ? compact_scale_elems : padded_scale_elems) * scale_elem_size; + const size_t output_stride_bytes = padded_scale_elems * scale_elem_size; - constexpr int SF_TILE_DIM_M = 128; - constexpr int SF_TILE_DIM_K = 4; - const dim3 block_size(TB_DIM, TB_DIM); + const int num_tiles_m = padded_m / SF_TILE_DIM_M; + const int num_tiles_k = padded_k / SF_TILE_DIM_K; + int vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1)); + if (vec_load_size == 3) vec_load_size = 1; + const int n_tiles_in_tb = TB_DIM * vec_load_size; - auto launch_grouped_swizzle = [&](bool rowwise) { - const size_t m = rowwise ? first_dim : last_dim; - const size_t k = rowwise ? last_dim : first_dim; - const size_t padded_m = round_up_to_multiple(m, 128); - const size_t padded_k = - round_up_to_multiple(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4); - // Per-tensor scale-element counts: - // - "padded" layout: each tensor occupies padded_m * padded_k elements - // (total buffer = num_tensors * padded_m * padded_k). - // - "compact" layout (what the grouped MXFP8 quantize kernel actually writes): - // per-tensor stride is m * padded_k (rowwise) or DIVUP(k,32) * padded_m - // (columnwise) and the total buffer the C++ allocator hands out has its - // grouped first dim padded up to a multiple of 128 (rowwise) or 4 - // (columnwise) -- so the buffer may be slightly larger than - // num_tensors * compact_scale_elems, with trailing alignment slack at - // the very end (never read because of the per-tensor row/k guard in the - // kernel impl). - // The output is always written in the padded layout. The input may be in - // either layout; the kernel handles the compact case safely by using - // different per-tensor strides for input vs output and skipping loads past - // the per-tensor extent. - const size_t padded_scale_elems = padded_m * padded_k; - const size_t compact_scale_elems = - rowwise ? m * padded_k : DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)) * padded_m; - const size_t compact_total_scale_elems = - rowwise ? round_up_to_multiple(input->num_tensors * m, 128) * padded_k - : round_up_to_multiple( - input->num_tensors * DIVUP(k, static_cast(MXFP8_BLOCK_SIZE)), 4) * - padded_m; - - const size_t scale_elem_size = rowwise ? typeToSize(input->scale_inv.dtype) - : typeToSize(input->columnwise_scale_inv.dtype); - - const size_t input_scale_numel = - rowwise ? input->scale_inv.numel() : input->columnwise_scale_inv.numel(); - const size_t output_scale_numel = - rowwise ? output->scale_inv.numel() : output->columnwise_scale_inv.numel(); - - bool input_is_compact; - if (input_scale_numel == input->num_tensors * padded_scale_elems) { - input_is_compact = false; - } else if (input_scale_numel == compact_total_scale_elems) { - input_is_compact = true; - } else { - NVTE_ERROR("Grouped input ", (rowwise ? "scale_inv" : "columnwise_scale_inv"), - " size does not match expected packed size (got ", input_scale_numel, - ", expected either ", input->num_tensors * padded_scale_elems, - " (per-tensor padded) or ", compact_total_scale_elems, " (compact))."); - } - NVTE_CHECK(output_scale_numel == input->num_tensors * padded_scale_elems, "Grouped output ", - (rowwise ? "scale_inv" : "columnwise_scale_inv"), - " size does not match expected per-tensor padded size."); + dim3 num_blocks; + if (rowwise) { + num_blocks = dim3(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m, input->num_tensors); + } else { + num_blocks = + dim3(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size), input->num_tensors); + } + const int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); - const size_t input_stride_bytes = - (input_is_compact ? compact_scale_elems : padded_scale_elems) * scale_elem_size; - const size_t output_stride_bytes = padded_scale_elems * scale_elem_size; + const int original_M = static_cast(rowwise ? first_dim : last_dim); + const int original_K = static_cast(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE))); + const void* input_ptr = rowwise ? input->scale_inv.dptr : input->columnwise_scale_inv.dptr; + void* output_ptr = rowwise ? output->scale_inv.dptr : output->columnwise_scale_inv.dptr; - const int num_tiles_m = padded_m / SF_TILE_DIM_M; - const int num_tiles_k = padded_k / SF_TILE_DIM_K; - int vec_load_size = (rowwise ? ((num_tiles_k - 1) % 4 + 1) : ((num_tiles_m - 1) % 4 + 1)); - if (vec_load_size == 3) vec_load_size = 1; - const int n_tiles_in_tb = TB_DIM * vec_load_size; + if (rowwise) { + TRANSFORMER_ENGINE_VECTORIZED_LOAD_INTEGER_TYPE_SWITCH(vec_load_size, LType, { + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + grouped_swizzle_row_scaling_uniform_shape_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + grouped_swizzle_row_scaling_uniform_shape_kernel + <<>>( + input_ptr, output_ptr, padded_m, padded_k, original_M, original_K, + input_stride_bytes, output_stride_bytes); + }); + } else { + TRANSFORMER_ENGINE_VECTORIZED_LOAD_INTEGER_TYPE_SWITCH(vec_load_size, LType, { + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + grouped_swizzle_col_scaling_uniform_shape_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); + grouped_swizzle_col_scaling_uniform_shape_kernel + <<>>( + input_ptr, output_ptr, padded_m, padded_k, original_M, original_K, + input_stride_bytes, output_stride_bytes); + }); + } + NVTE_CHECK_CUDA(cudaGetLastError()); + }; - dim3 num_blocks; - if (rowwise) { - num_blocks = dim3(DIVUP(num_tiles_k, n_tiles_in_tb), num_tiles_m, input->num_tensors); - } else { - num_blocks = - dim3(DIVUP(num_tiles_k, TB_DIM), DIVUP(num_tiles_m, vec_load_size), input->num_tensors); + if (has_rowwise_scale_inv) { + launch_grouped_swizzle(true); } - const int slm_size = n_tiles_in_tb * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); - - const int original_M = static_cast(rowwise ? first_dim : last_dim); - const int original_K = static_cast(DIVUP(k, static_cast(MXFP8_BLOCK_SIZE))); - const void* input_ptr = rowwise ? input->scale_inv.dptr : input->columnwise_scale_inv.dptr; - void* output_ptr = rowwise ? output->scale_inv.dptr : output->columnwise_scale_inv.dptr; - - if (rowwise) { - TRANSFORMER_ENGINE_VECTORIZED_LOAD_INTEGER_TYPE_SWITCH(vec_load_size, LType, { - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - grouped_swizzle_row_scaling_uniform_shape_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - grouped_swizzle_row_scaling_uniform_shape_kernel - <<>>(input_ptr, output_ptr, padded_m, - padded_k, original_M, original_K, - input_stride_bytes, output_stride_bytes); - }); - } else { - TRANSFORMER_ENGINE_VECTORIZED_LOAD_INTEGER_TYPE_SWITCH(vec_load_size, LType, { - NVTE_CHECK_CUDA(cudaFuncSetAttribute( - grouped_swizzle_col_scaling_uniform_shape_kernel, - cudaFuncAttributeMaxDynamicSharedMemorySize, slm_size)); - grouped_swizzle_col_scaling_uniform_shape_kernel - <<>>(input_ptr, output_ptr, padded_m, - padded_k, original_M, original_K, - input_stride_bytes, output_stride_bytes); - }); + if (has_columnwise_scale_inv) { + launch_grouped_swizzle(false); + } + } else { + // Variable shape implementation using Device-Side Block Scheduler + size_t num_tensors = input->num_tensors; + + constexpr int SF_TILE_DIM_M = 128; + constexpr int SF_TILE_DIM_K = 4; + const dim3 block_size(TB_DIM, TB_DIM); + const int max_slm_size = TB_DIM * 4 * SF_TILE_DIM_M * SF_TILE_DIM_K * sizeof(int8_t); + const int metadata_shmem = sizeof(int); // s_total_blocks + const int dynamic_smem_size = max_slm_size + metadata_shmem; + + size_t common_m = input->all_same_first_dim() ? input->get_common_first_dim() : 0; + size_t common_k = input->all_same_last_dim() ? input->get_common_last_dim() : 0; + + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + grouped_swizzle_scaling_variable_shape_kernel, + cudaFuncAttributeMaxDynamicSharedMemorySize, dynamic_smem_size)); + + const int device_id = cuda::current_device(); + const int num_SMs = cuda::sm_count(device_id); + const int max_active_blocks_per_sm = + grouped_swizzle_variable_max_active_blocks_per_sm(device_id); + const int persistent_blocks = num_SMs * max_active_blocks_per_sm; + const dim3 num_blocks(persistent_blocks); + + auto launch_grouped_swizzle_variable = [&](bool rowwise) { + const size_t scale_elem_size = rowwise ? typeToSize(input->scale_inv.dtype) + : typeToSize(input->columnwise_scale_inv.dtype); + + const void* input_ptr = rowwise ? input->scale_inv.dptr : input->columnwise_scale_inv.dptr; + void* output_ptr = rowwise ? output->scale_inv.dptr : output->columnwise_scale_inv.dptr; + + grouped_swizzle_scaling_variable_shape_kernel + <<>>( + input_ptr, output_ptr, m_array, k_array, num_tensors, rowwise, scale_elem_size, + common_m, common_k); + + NVTE_CHECK_CUDA(cudaGetLastError()); + }; + + if (has_rowwise_scale_inv) { + launch_grouped_swizzle_variable(true); + } + if (has_columnwise_scale_inv) { + launch_grouped_swizzle_variable(false); } - NVTE_CHECK_CUDA(cudaGetLastError()); - }; - - if (has_rowwise_scale_inv) { - launch_grouped_swizzle(true); - } - if (has_columnwise_scale_inv) { - launch_grouped_swizzle(false); } } diff --git a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp index d8ab830c4..7f7f8a435 100644 --- a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp +++ b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp @@ -379,13 +379,6 @@ std::optional maybe_swizzle_grouped_tensor(GroupedTensorW if (!swizzle_rowwise && !swizzle_columnwise) { return std::nullopt; } - const auto first_dims = input.get_first_dims(); - const auto last_dims = input.get_last_dims(); - if (first_dims.data_ptr != nullptr || last_dims.data_ptr != nullptr) { - NVTE_ERROR( - "Grouped GEMM swizzle requires uniform shapes for now (first_dims/last_dims must be " - "absent)."); - } std::optional rowwise_scales_pyt; std::optional columnwise_scales_pyt; @@ -452,6 +445,7 @@ std::optional maybe_swizzle_grouped_tensor(GroupedTensorW } swizzle_output.set_with_gemm_swizzled_scales(true); + NVTE_SCOPED_GIL_RELEASE({ nvte_swizzle_grouped_scaling_factors(swizzle_input.data(), swizzle_output.data(), at::cuda::getCurrentCUDAStream()); From 88e607186400f8546908c829783c69599c74aac5 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Thu, 30 Apr 2026 19:33:16 -0700 Subject: [PATCH 026/180] [PyTorch] Fusible ops preserve usages in quantized weight tensors (#2929) * Avoid removing usages from quantized weight in linear op Quantized weight tensor may be used across steps, so removing a usage is not safe. Signed-off-by: Tim Moon * Tweak test to catch bug when alternating train and infer steps Signed-off-by: Tim Moon * Avoid removing usages from quantized weights in grouped linear op Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Restore pre-forward quantizer config in ops Turns out we still need this in case the quantizer is used before the forward, e.g. in previous ops or CPU offloading. Signed-off-by: Tim Moon * Blindly preserve quantizer usages in quantized weight params. Signed-off-by: Tim Moon --------- Signed-off-by: Tim Moon Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/test_fusible_ops.py | 264 ++++++++++++++++-- .../pytorch/ops/basic/basic_linear.py | 40 +-- .../pytorch/ops/basic/grouped_linear.py | 66 +++-- 3 files changed, 311 insertions(+), 59 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 3d6fe704e..10baae0d9 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -4,7 +4,7 @@ from __future__ import annotations -from collections.abc import Iterable +from collections.abc import Iterable, Sequence import functools import io import math @@ -200,6 +200,18 @@ def make_reference_and_test_tensors( return ref, test +def to_cpu(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: + """Convert to an FP64 CPU tensor""" + if tensor is None: + return None + out = tensor.detach() + if isinstance(out, QuantizedTensor): + out = out.dequantize() + out = out.to(dtype=torch.float64, device="cpu") + out = out.requires_grad_(requires_grad=tensor.requires_grad) + return out + + class MegatronTrainingHelper: """Test-side stand-in for the Megatron-Core DDP / MegatronFSDP wrapper. Megatron's DDP wrapper (and MegatronFSDP) owns the per-parameter @@ -3368,25 +3380,17 @@ def test_layernorm_mlp( y_test = forward(x_test) y_test.backward(dy_test) - def to_cpu(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: - """Convert to FP64 CPU tensor""" - if tensor is None: - return None - out = tensor.detach().to(dtype=torch.float64, device="cpu") - out = out.requires_grad_(requires_grad=tensor.requires_grad) - return out - # Check values tols = {"rtol": 0.25, "atol": 0.5} # Loose tols for sanity checking - torch.testing.assert_close(to_cpu(y_test), y_ref, **tols) - torch.testing.assert_close(to_cpu(x_test.grad), x_ref.grad, **tols) - torch.testing.assert_close(to_cpu(norm.weight.grad), norm_w_ref.grad, **tols) - torch.testing.assert_close(to_cpu(norm.bias.grad), norm_b_ref.grad, **tols) - torch.testing.assert_close(to_cpu(ffn2.weight.grad), w2_ref.grad, **tols) - torch.testing.assert_close(to_cpu(ffn1.weight.grad), w1_ref.grad, **tols) + assert_close(y_test, y_ref, **tols) + assert_close(x_test.grad, x_ref.grad, **tols) + assert_close_grads(norm.weight, norm_w_ref, **tols) + assert_close_grads(norm.bias, norm_b_ref, **tols) + assert_close_grads(ffn2.weight, w2_ref, **tols) + assert_close_grads(ffn1.weight, w1_ref, **tols) if bias: - torch.testing.assert_close(to_cpu(ffn1.bias.grad), b1_ref.grad, **tols) - torch.testing.assert_close(to_cpu(ffn2.bias.grad), b2_ref.grad, **tols) + assert_close_grads(ffn1.bias, b1_ref, **tols) + assert_close_grads(ffn2.bias, b2_ref, **tols) @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("dtype", _dtypes) @@ -4740,6 +4744,232 @@ def fuse_ops( torch.testing.assert_close(dw_test, w_ref.grad, **tols) +class TestTrainingLoops: + + def _linear_train_stage( + self, + module: te.ops.Linear, + *, + steps: int = 3, + in_shape: Sequence[int], + out_shape: Sequence[int], + dtype: torch.type, + device: torch.device, + quantization: Optional[str], + recipe: Optional[transformer_engine.common.recipe.Recipe], + ) -> None: + """Perform training steps with linear op""" + + # Expected numerical error + tols = dtype_tols(dtype) + if dtype == torch.float32: + tols = dtype_tols(torch.float16) # TF32 GEMM + if quantization is not None: + tols = quantization_tols(quantization) + + for _ in range(steps): + # Update parameters with random values to simulate + # optimizer step or FSDP param all-gather + with torch.no_grad(): + module.weight.copy_(torch.empty_like(module.weight).uniform_()) + module.bias.copy_(torch.empty_like(module.bias).uniform_()) + for param in module.parameters(): + param.grad = None + + # Random data + x_ref, x_test = make_reference_and_test_tensors( + in_shape, + quantization=quantization, + test_dtype=dtype, + test_device=device, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + out_shape, + quantization=quantization, + test_dtype=dtype, + test_device=device, + ) + w_ref = to_cpu(module.weight) + b_ref = to_cpu(module.bias) + + # Plain PyTorch implementation + y_ref = torch.nn.functional.linear(x_ref, w_ref, bias=b_ref) + y_ref.backward(dy_ref) + + # Implementation with linear op + with te.autocast(enabled=quantization is not None, recipe=recipe): + y_test = module(x_test) + y_test.backward(dy_test) + + # Check results + assert_close(y_test, y_ref, **tols) + assert_close_grads(x_test, x_ref, **tols) + assert_close_grads(module.weight, w_ref, **tols) + assert_close_grads(module.bias, b_ref, **tols) + + @torch.inference_mode + def _linear_infer_stage( + self, + module: te.ops.Linear, + *, + steps: int = 3, + in_shape: Sequence[int], + dtype: torch.type, + device: torch.device, + quantization: Optional[str], + recipe: Optional[transformer_engine.common.recipe.Recipe], + ) -> None: + """Perform inference steps with linear op""" + + # Parameter reference values + w_ref = to_cpu(module.weight) + b_ref = to_cpu(module.bias) + + # Expected numerical error + tols = dtype_tols(dtype) + if dtype == torch.float32: + tols = dtype_tols(torch.float16) # TF32 GEMM + if quantization is not None: + tols = quantization_tols(quantization) + + for _ in range(steps): + # Random data + x_ref, x_test = make_reference_and_test_tensors( + in_shape, + quantization=quantization, + test_dtype=dtype, + test_device=device, + ) + + # Plain PyTorch implementation + y_ref = torch.nn.functional.linear(x_ref, w_ref, bias=b_ref) + + # Implementation with linear op + with te.autocast(enabled=quantization is not None, recipe=recipe): + y_test = module(x_test) + + # Check results + assert_close(y_test, y_ref, **tols) + + @pytest.mark.parametrize("stages", (["train", "infer"] * 2, ["infer", "train"] * 2)) + @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize("quantized_weight", (False, True)) + def test_linear_training_loop( + self, + *, + stages: Sequence[str], + weight_shape: tuple[int, int] = (32, 32), + in_shape: Sequence[int] = (32, -1), + dtype: Optional[torch.dtype] = None, + device: torch.device = "cuda", + quantization: Optional[str], + quantized_weight: bool, + ) -> None: + """Training loops with linear op""" + if dtype is None: + dtype = torch.bfloat16 if is_bf16_available() else torch.float32 + + # Make input and weight shapes consistent + out_features, in_features = weight_shape + in_shape = list(in_shape)[:-1] + [in_features] + out_shape = in_shape[:-1] + [out_features] + + # Skip invalid configurations + maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) + maybe_skip_quantization(quantization, dims=out_shape) + if quantization is None and quantized_weight: + pytest.skip("Quantization scheme is not specified") + + # Construct module with random weights + recipe = make_recipe(quantization) + with te.quantized_model_init(enabled=quantized_weight, recipe=recipe): + module = te.ops.Linear( + in_features, + out_features, + device=device, + dtype=dtype, + ) + with torch.no_grad(): + for param in module.parameters(): + param.copy_(torch.empty_like(param).uniform_()) + + # Training loop stages + for stage in stages: + if stage == "train": + self._linear_train_stage( + module, + in_shape=in_shape, + out_shape=out_shape, + dtype=dtype, + device=device, + quantization=quantization, + recipe=recipe, + ) + elif stage == "infer": + self._linear_infer_stage( + module, + in_shape=in_shape, + dtype=dtype, + device=device, + quantization=quantization, + recipe=recipe, + ) + else: + raise ValueError(f"Unrecognized stage ({stage})") + + @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize("quantized_weight", (False, True)) + def test_linear_inference_loop( + self, + *, + weight_shape: tuple[int, int] = (32, 32), + in_shape: Sequence[int] = (32, -1), + dtype: Optional[torch.dtype] = None, + device: torch.device = "cuda", + quantization: Optional[str], + quantized_weight: bool, + ) -> None: + """Inference loop with linear op""" + if dtype is None: + dtype = torch.bfloat16 if is_bf16_available() else torch.float32 + + # Make input and weight shapes consistent + out_features, in_features = weight_shape + in_shape = list(in_shape)[:-1] + [in_features] + out_shape = in_shape[:-1] + [out_features] + + # Skip invalid configurations + maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) + maybe_skip_quantization(quantization, dims=out_shape) + if quantization is None and quantized_weight: + pytest.skip("Quantization scheme is not specified") + + # Construct module with random weights + recipe = make_recipe(quantization) + with ( + torch.inference_mode(), + te.quantized_model_init(enabled=quantized_weight, recipe=recipe), + ): + module = te.ops.Linear( + in_features, + out_features, + device=device, + dtype=dtype, + ) + for param in module.parameters(): + param.copy_(torch.empty_like(param).uniform_()) + + # Inference loop + self._linear_infer_stage( + module, + in_shape=in_shape, + dtype=dtype, + device=device, + quantization=quantization, + recipe=recipe, + ) + + def test_grouped_gemm_quant_cute_matches_mxfp8_quantized() -> None: if not mxfp8_available: pytest.skip(reason_for_no_mxfp8) diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index 17594726c..19fcf62ce 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -329,8 +329,6 @@ def pre_fuser_forward(self, *, requires_grad: bool) -> None: super().pre_fuser_forward(requires_grad=requires_grad) if FP8GlobalStateManager.is_fp8_enabled(): # Configure quantizer usages - # Note: We cache the quantized input for backward pass, - # but discard the quantized weights. weight_requires_grad = requires_grad and self.weight.requires_grad columnwise_usage = weight_requires_grad if FP8GlobalStateManager.get_fp8_recipe().backward_override is not None: @@ -339,13 +337,13 @@ def pre_fuser_forward(self, *, requires_grad: bool) -> None: weight_quantizer = self.get_quantizer("forward", 1) grad_output_quantizer = self.get_quantizer("backward", 0) input_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) - weight_quantizer.set_usage(rowwise=True, columnwise=False) + weight_quantizer.set_usage(rowwise=True, columnwise=requires_grad) grad_output_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: super().reset_recipe_state(recipe=recipe) - # Configure input/grad output tensor + # Configure input/grad output quantizers # Note: These tensors are only used internally. If there is no # tensor-parallel communication, they are only used for GEMM. input_quantizer = self.get_quantizer("forward", 0) @@ -370,21 +368,15 @@ def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: # Configure weight quantizer # Note: This function may be called in base class constructor, - # before any basic linear attrs have been set. + # before basic linear attrs have been set. weight_quantizer = self.get_quantizer("forward", 1) - if weight_quantizer is None: - pass - elif is_quantized_tensor(getattr(self, "weight", None)): - # Make sure weight param has correct quantizer - weight_quantizer.set_usage(rowwise=True, columnwise=torch.is_grad_enabled()) - weight_quantizer.internal = False - self.weight.update_quantizer(weight_quantizer.copy()) - else: - # Use internal tensors if quantized weights will not be - # exposed externally - weight_quantizer.internal = ( - not FP8GlobalStateManager.with_fp8_parameters() - and not getattr(self, "_with_quantized_weight", False) + weight = getattr(self, "weight", None) + if weight_quantizer is not None: + # Determine if quantized weight is exposed as parameter + weight_quantizer.internal = not ( + FP8GlobalStateManager.with_fp8_parameters() + or getattr(self, "_with_quantized_weight", False) + or is_quantized_tensor(weight) ) # Recipe-specific configuration @@ -416,6 +408,18 @@ def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: grad_output_quantizer.with_amax_reduction = True grad_output_quantizer.amax_reduction_group = self.tensor_parallel_group + # Update quantizer in quantized weight tensor + if weight_quantizer is not None and is_quantized_tensor(weight): + if weight._quantizer is not None: + # Preserve existing usages in weight tensor. Even if a + # usage is currently unnecessary, the weight tensor + # may be used elsewhere. + weight_quantizer.set_usage( + rowwise=weight._quantizer.rowwise_usage, + columnwise=weight._quantizer.columnwise_usage, + ) + weight.update_quantizer(weight_quantizer.copy()) + @staticmethod def _functional_forward( input: torch.Tensor, # pylint: disable=redefined-builtin diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index fe5997a71..b503cb186 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -619,14 +619,12 @@ def pre_fuser_forward(self, *, requires_grad: bool) -> None: weight_requires_grad = requires_grad and weight_requires_grad # Configure quantizer usages - # Note: We cache the quantized input for backward pass, - # but discard the quantized weights. for group_idx in range(self.num_groups): input_quantizer = self.get_quantizer("forward", 2 * group_idx) weight_quantizer = self.get_quantizer("forward", 2 * group_idx + 1) grad_output_quantizer = self.get_quantizer("backward", group_idx) input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) - weight_quantizer.set_usage(rowwise=True, columnwise=False) + weight_quantizer.set_usage(rowwise=True, columnwise=requires_grad) grad_output_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: @@ -641,32 +639,29 @@ def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: if grad_output_quantizer is not None: grad_output_quantizer.internal = True - # Handle weight quantizer + # Get weight tensor # Note: This function may be called in base class constructor, - # before any basic linear attrs have been set. - weight_quantizer = self.get_quantizer("forward", 2 * group_idx + 1) - if weight_quantizer is None: - pass - elif is_quantized_tensor(getattr(self, f"weight{group_idx}", None)): - # Make sure weight param has correct quantizer - weight_quantizer.set_usage(rowwise=True, columnwise=torch.is_grad_enabled()) - weight_quantizer.internal = False - if self.single_grouped_weight: - self.weight.quantizer = weight_quantizer.copy() - else: - getattr(self, f"weight{group_idx}").update_quantizer(weight_quantizer.copy()) + # before any grouped linear attrs have been set. + weight = None + weight_is_quantized = False + if getattr(self, "single_grouped_weight", False): + weight = getattr(self, "weight", None) + weight_is_quantized = weight is not None and weight.quantizer is not None else: - # Use internal tensors if quantized weights will not be - # exposed externally - weight_quantizer.internal = ( - not FP8GlobalStateManager.with_fp8_parameters() - and not getattr(self, "_with_quantized_weight", False) - and not self.single_grouped_weight + weight = getattr(self, f"weight{group_idx}", None) + weight_is_quantized = is_quantized_tensor(weight) + + # Configure weight quantizer + weight_quantizer = self.get_quantizer("forward", 2 * group_idx + 1) + if weight_quantizer is not None: + # Determine if quantized weight is exposed as parameter + weight_quantizer.internal = not ( + FP8GlobalStateManager.with_fp8_parameters() + or getattr(self, "_with_quantized_weight", False) + or weight_is_quantized ) # Recipe-specific configuration - # Note: This function may be called in base class constructor, - # before any basic linear attrs have been set. if recipe is not None: if recipe.float8_current_scaling(): input_quantizer.force_pow_2_scales = recipe.fp8_quant_fwd_inp.power_2_scale @@ -680,6 +675,29 @@ def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: recipe.fp8_quant_bwd_grad.amax_epsilon ) + # Update quantizer in quantized weight tensor + if weight_quantizer is not None and weight_is_quantized: + # Get quantizer from weight tensor + weight_tensor_quantizer = ( + weight.quantizer if self.single_grouped_weight else weight._quantizer + ) + + # Preserve existing usages in weight tensor. Even if a + # usage is currently unnecessary, the weight tensor + # may be used elsewhere. + if weight_tensor_quantizer is not None: + weight_quantizer.set_usage( + rowwise=weight_tensor_quantizer.rowwise_usage, + columnwise=weight_tensor_quantizer.columnwise_usage, + ) + + # Update weight tensor + if self.single_grouped_weight: + if group_idx == 0: + weight.quantizer = weight_quantizer.copy() + else: + weight.update_quantizer(weight_quantizer.copy()) + def op_forward(self, *args, **kwargs): raise RuntimeError( f"{self.__class__.__name__} operation has " From 4fafdf2a330c067badc3a9c25124ff1ce4e9ac9f Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Fri, 1 May 2026 05:22:56 +0200 Subject: [PATCH 027/180] [Common] Fix incorrect amax initialization in non-RHT NVFP4 C++ tests (#2943) * Patch for NVFP4 test suite Signed-off-by: Oleg Goncharov * C++ tests fix Signed-off-by: Oleg Goncharov * Cleanup Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Removed dead code Signed-off-by: Oleg Goncharov * Set the golden value for amax in tests Signed-off-by: Oleg Goncharov * Fixed memory leakage Signed-off-by: Oleg Goncharov --------- Signed-off-by: Oleg Goncharov Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../cpp/operator/test_cast_nvfp4_transpose.cu | 17 +++---- tests/cpp/test_common.cu | 45 +++++++++++++++---- tests/cpp/test_common.h | 40 +++++++++++++++-- 3 files changed, 79 insertions(+), 23 deletions(-) diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index d8d495d61..15d7c695c 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -565,17 +565,12 @@ void performTest(float (*OP)(const float), fillCase(&input, InputsFillCase::uniform); - // Find global amax - float amax = 0.0f; - const InputType* input_dptr = input.rowwise_cpu_dptr(); - for (size_t i = 0; i < rows; ++i) { - for (size_t j = 0; j < cols; ++j) { - const size_t idx = i * cols + j; - amax = fmaxf(amax, static_cast(input_dptr[idx])); - } - } + // Golden value of amax chosen to make the 2nd-stage scaling mantissa zero and avoid rounding issues + const float amax = 448.0f * 6.0f * 8.0f; + // Set 2nd stage NVFP4 scaling factor - output.set_scale(amax); + output.set_tensor_amax(amax); + output.set_tensor_amax_columnwise(amax); bool use_2d_quantization = false; @@ -585,7 +580,7 @@ void performTest(float (*OP)(const float), ref_output_t.get(), ref_scales.get(), ref_scales_t.get(), - output.scale(), + amax, rows, cols, scales_stride, diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index b8bc38935..c756b8381 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -364,8 +364,11 @@ Tensor::Tensor(const std::string& name, } // Configure scales, amaxes, and other tensor buffers - float *amax = nullptr, *scale = nullptr; - float *rowwise_scale_inv = nullptr, *columnwise_scale_inv = nullptr; + float *amax = nullptr; + float *amax_columnwise = nullptr; + float *scale = nullptr; + float *rowwise_scale_inv = nullptr; + float *columnwise_scale_inv = nullptr; if (isFp8Type(type) || isFp4Type(type)) { if (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) { cudaMalloc((void**)&amax, sizeof(float)); // NOLINT(*) @@ -392,10 +395,14 @@ Tensor::Tensor(const std::string& name, } else { if (scaling_mode == NVTE_NVFP4_1D_SCALING) { // Used for NVFP4 second stage scaling - cudaMalloc((void**)&scale, sizeof(float)); // NOLINT(*) - cudaMemset(scale, 0, sizeof(float)); - scale_cpu_data_ = std::make_shared(0); - tensor_.set_scale(scale, DType::kFloat32, std::vector{1}); + amax_cpu_data_ = std::make_shared(0); + amax_cpu_data_columnwise_ = std::make_shared(0); + cudaMalloc((void**)&amax, sizeof(float)); // NOLINT(*) + cudaMalloc((void**)&amax_columnwise, sizeof(float)); // NOLINT(*) + cudaMemset(amax, 0, sizeof(float)); + cudaMemset(amax_columnwise, 0, sizeof(float)); + tensor_.set_amax(amax, DType::kFloat32, std::vector{1}); + tensor_.set_columnwise_amax(amax_columnwise, DType::kFloat32, std::vector{1}); } auto [rowwise_scale_meta, colwise_scale_meta] = get_scales(flattened_shape, tensor_.scaling_mode()); auto rowwise_scale_size = rowwise_scale_meta.bytes(); @@ -441,7 +448,7 @@ void Tensor::to_cpu() const { cudaMemcpyDeviceToHost); } if (isFp8Type(dtype()) || isFp4Type(dtype())) { - if ((tensor_.scaling_mode() == NVTE_DELAYED_TENSOR_SCALING)) { + if (tensor_.scaling_mode() == NVTE_DELAYED_TENSOR_SCALING) { if (tensor_.amax() != nullptr){ cudaMemcpy(amax_cpu_data_.get(), tensor_.amax(), @@ -452,6 +459,19 @@ void Tensor::to_cpu() const { tensor_.scale(), sizeof(float), cudaMemcpyDeviceToHost); + } else if (tensor_.scaling_mode() == NVTE_NVFP4_1D_SCALING) { + if (rowwise_ && (tensor_.amax() != nullptr)){ + cudaMemcpy(amax_cpu_data_.get(), + tensor_.amax(), + sizeof(float), + cudaMemcpyDeviceToHost); + } + if (columnwise_ && (tensor_.get_columnwise_amax().data_ptr != nullptr)){ + cudaMemcpy(amax_cpu_data_columnwise_.get(), + tensor_.get_columnwise_amax().data_ptr, + sizeof(float), + cudaMemcpyDeviceToHost); + } } auto [rowwise_scale_meta, colwise_scale_meta] = get_scales(s, tensor_.scaling_mode()); if (rowwise_) { @@ -483,12 +503,19 @@ void Tensor::from_cpu() const { cudaMemcpyHostToDevice); } if (isFp8Type(dtype()) || isFp4Type(dtype())) { - if ((tensor_.scaling_mode() == NVTE_DELAYED_TENSOR_SCALING) - || (tensor_.scaling_mode() == NVTE_NVFP4_1D_SCALING)) { + if (tensor_.scaling_mode() == NVTE_DELAYED_TENSOR_SCALING) { if (tensor_.amax() != nullptr){ cudaMemcpy(tensor_.amax(), amax_cpu_data_.get(), sizeof(float), cudaMemcpyHostToDevice); } cudaMemcpy(tensor_.scale(), scale_cpu_data_.get(), sizeof(float), cudaMemcpyHostToDevice); + } else if (tensor_.scaling_mode() == NVTE_NVFP4_1D_SCALING) { + if (rowwise_ && (tensor_.amax() != nullptr)) { + cudaMemcpy(tensor_.amax(), amax_cpu_data_.get(), sizeof(float), cudaMemcpyHostToDevice); + } + if (columnwise_ && (tensor_.get_columnwise_amax().data_ptr != nullptr)) { + cudaMemcpy(tensor_.get_columnwise_amax().data_ptr, amax_cpu_data_columnwise_.get(), + sizeof(float), cudaMemcpyHostToDevice); + } } auto [rowwise_scale_meta, colwise_scale_meta] = get_scales(s, tensor_.scaling_mode()); if (rowwise_) { diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index b5a7f26d1..b8389d583 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -146,6 +146,9 @@ class Tensor { void *scale_inv = tensor_.scale_inv(); void *columnwise_data_ptr = tensor_.get_columnwise_data().data_ptr; void *columnwise_scale_inv = tensor_.get_columnwise_scale_inv().data_ptr; + void *amax = tensor_.amax(); + void *columnwise_amax_ptr = tensor_.get_columnwise_amax().data_ptr; + void *scale = tensor_.scale(); if (columnwise_data_ptr == data_ptr) { columnwise_data_ptr = nullptr; } @@ -164,6 +167,15 @@ class Tensor { if (columnwise_scale_inv != nullptr) { cudaFree(columnwise_scale_inv); } + if (amax != nullptr) { + cudaFree(amax); + } + if (columnwise_amax_ptr != nullptr) { + cudaFree(columnwise_amax_ptr); + } + if (scale != nullptr) { + cudaFree(scale); + } } NVTETensor data() const noexcept { return tensor_.data(); } @@ -223,11 +235,18 @@ class Tensor { } } + float amax_columnwise() const { + if(amax_cpu_data_columnwise_) { + to_cpu(); + return *amax_cpu_data_columnwise_; + } else { + return 0; + } + } + float scale() const { if(scale_cpu_data_) { - NVTE_CHECK((tensor_.scaling_mode() == NVTE_DELAYED_TENSOR_SCALING) - || (tensor_.scaling_mode() == NVTE_NVFP4_1D_SCALING), - "Invalid scaling_mode!"); + NVTE_CHECK(tensor_.scaling_mode() == NVTE_DELAYED_TENSOR_SCALING, "Invalid scaling_mode!"); to_cpu(); return *scale_cpu_data_; } else { @@ -282,6 +301,20 @@ class Tensor { return columnwise_; } + void set_tensor_amax(const float amax) { + if (amax_cpu_data_) { + *amax_cpu_data_ = amax; + from_cpu(); + } + } + + void set_tensor_amax_columnwise(const float amax) { + if (amax_cpu_data_columnwise_) { + *amax_cpu_data_columnwise_ = amax; + from_cpu(); + } + } + void set_tensor_amax_nullptr(){ tensor_.set_amax(nullptr, DType::kFloat32, tensor_.defaultShape); } @@ -303,6 +336,7 @@ class Tensor { std::unique_ptr cpu_data_rowwise_; std::unique_ptr cpu_data_columnwise_; std::shared_ptr amax_cpu_data_; + std::shared_ptr amax_cpu_data_columnwise_; std::shared_ptr scale_cpu_data_; std::unique_ptr rowwise_scale_inv_cpu_data_; std::unique_ptr columnwise_scale_inv_cpu_data_; From 0e9020d4f1d9ecd88bddc67d6fcbc7394dd47013 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Fri, 1 May 2026 02:23:09 -0400 Subject: [PATCH 028/180] [PyTorch] Cleanup `cudnn-frontend` requirements for fused grouped MLP (#2948) * Switch to cuDNN-FE min version 1.23.0 to enable fused grouped MLP Signed-off-by: Kirthi Shankar Sivamani * Fix tests Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani --- tests/pytorch/test_fusible_ops.py | 19 +------- transformer_engine/pytorch/ops/_common.py | 28 ++---------- .../pytorch/ops/fused/backward_grouped_mlp.py | 36 ++++------------ .../pytorch/ops/fused/forward_grouped_mlp.py | 43 ++----------------- 4 files changed, 18 insertions(+), 108 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 10baae0d9..47507dc38 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -20,7 +20,7 @@ import transformer_engine.pytorch as te import transformer_engine.pytorch.ops as te_ops from transformer_engine.pytorch.ops._common import ( - _nvidia_cudnn_frontend_supports_scaled_clamped_qgeglu, + _cudnn_frontend_version_supported, ) from transformer_engine.pytorch.ops.fused import ( @@ -3642,10 +3642,7 @@ def test_grouped_mlp( quantization == "mxfp8" and dtype in (torch.bfloat16, torch.float16) and glu_interleave_size == 32 - and ( - activation != "scaled_clamped_qgeglu" - or _nvidia_cudnn_frontend_supports_scaled_clamped_qgeglu() - ) + and _cudnn_frontend_version_supported() ): if te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported(): forward_ops = module._module_groups[0]._forward_ops @@ -3748,12 +3745,6 @@ def test_grouped_mlp_single_weight_numerics( pytest.skip("MXFP8 fused grouped MLP forward is not supported on this system") if not te_ops.fused.BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8.is_supported(): pytest.skip("MXFP8 fused grouped MLP backward is not supported on this system") - if activation == "scaled_clamped_qgeglu" and not ( - _nvidia_cudnn_frontend_supports_scaled_clamped_qgeglu() - ): - pytest.skip( - "ScaledClampedQGeGLU fused grouped MLP requires nvidia-cudnn-frontend >= 1.23.0" - ) split_sizes = [split_alignment * (i + 1) for i in range(group_size)] random.shuffle(split_sizes) @@ -4110,12 +4101,6 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8( pytest.skip("MXFP8 fused grouped MLP is not supported on this system") if dtype not in (torch.bfloat16, torch.float16): pytest.skip("MXFP8 fused grouped MLP is only supported with BF16/FP16") - if activation == "scaled_clamped_qgeglu" and not ( - _nvidia_cudnn_frontend_supports_scaled_clamped_qgeglu() - ): - pytest.skip( - "ScaledClampedQGeGLU fused grouped MLP requires nvidia-cudnn-frontend >= 1.23.0" - ) split_sizes = [split_alignment * (i + 1) for i in range(group_size)] random.shuffle(split_sizes) diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index e21915a5a..beef6fe52 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -21,17 +21,11 @@ @functools.lru_cache(maxsize=1) -def _nvidia_cudnn_frontend_supports_scaled_clamped_qgeglu() -> bool: - """Check cuDNN FE min version with fixed numerics for qgeglu.""" - try: - return PkgVersion(get_pkg_version("nvidia-cudnn-frontend")) >= PkgVersion("1.23.0") - except PackageNotFoundError: - return False - +def _cudnn_frontend_version_supported() -> bool: + """Check cuDNN frontend is at least 1.23.0. -@functools.lru_cache(maxsize=1) -def _nvidia_cudnn_frontend_supports_wgrad() -> bool: - """Check cuDNN FE min version for grouped GEMM wgrad kernel.""" + All grouped MLP fused-kernel features require cuDNN frontend 1.23.0. + """ try: return PkgVersion(get_pkg_version("nvidia-cudnn-frontend")) >= PkgVersion("1.23.0") except PackageNotFoundError: @@ -140,8 +134,6 @@ def fuse_grouped_mlp_ops( constructor accepting ``fc1``, ``glu_op``, ``fc2`` keyword args. The ``glu_op`` must be :class:`~transformer_engine.pytorch.ops.basic.swiglu.ScaledSwiGLU` or :class:`~transformer_engine.pytorch.ops.basic.swiglu.ScaledClampedQGeGLU`. - May also expose ``is_fc1_bias_supported()`` and/or - ``is_fc2_bias_supported()`` classmethods for bias eligibility. Returns ------- @@ -159,13 +151,6 @@ def fuse_grouped_mlp_ops( if recipe is None or not recipe.mxfp8(): return ops - fc1_bias_ok = ( - not hasattr(fused_op_cls, "is_fc1_bias_supported") or fused_op_cls.is_fc1_bias_supported() - ) - fc2_bias_ok = ( - not hasattr(fused_op_cls, "is_fc2_bias_supported") or fused_op_cls.is_fc2_bias_supported() - ) - out = [] window, ops = ops[:3], ops[3:] while len(window) == 3: @@ -179,7 +164,6 @@ def fuse_grouped_mlp_ops( matches_pattern = False elif isinstance(window[1], ScaledClampedQGeGLU) and ( abs(window[1]._clamped.alpha - 1.702) > 0.001 - or not _nvidia_cudnn_frontend_supports_scaled_clamped_qgeglu() ): matches_pattern = False elif window[0].num_groups != window[2].num_groups: @@ -193,10 +177,6 @@ def fuse_grouped_mlp_ops( matches_pattern = False elif window[1].glu_interleave_size != 32: matches_pattern = False - elif window[0].has_bias and not fc1_bias_ok: - matches_pattern = False - elif window[2].has_bias and not fc2_bias_ok: - matches_pattern = False if matches_pattern: op = fused_op_cls( diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 510fea0ed..70d0d7469 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -7,7 +7,6 @@ from __future__ import annotations from collections.abc import Callable import functools -import inspect import math import os from typing import Optional @@ -15,7 +14,6 @@ import torch import transformer_engine_torch as tex -from ...module.base import get_dummy_wgrad from ...quantization import Recipe from ...tensor.grouped_tensor import GroupedTensor from ...tensor.mxfp8_tensor import MXFP8Quantizer @@ -25,13 +23,13 @@ from ..fuser import register_backward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext from .._common import ( - _nvidia_cudnn_frontend_supports_wgrad, + _cudnn_frontend_version_supported, fuse_grouped_mlp_ops, maybe_dequantize, validate_grouped_mlp_dims, ) from ...cpp_extensions import general_grouped_gemm_for_grouped_tensor -from ...module.base import _2X_ACC_WGRAD +from ...module.base import _2X_ACC_WGRAD, get_dummy_wgrad from ...triton.grouped_dbias_dscales import compute_grouped_dbias_dscales @@ -109,20 +107,6 @@ def _cudnn_compute_wgrad( ) -@functools.lru_cache(maxsize=1) -def _dglu_wrapper_has_generate_dbias_arg() -> bool: - """True if cudnn-frontend SM100 dGLU wrapper accepts ``generate_dbias``.""" - try: - from cudnn import grouped_gemm_dglu_wrapper_sm100 # pylint: disable=import-outside-toplevel - except ImportError: - return False - try: - params = inspect.signature(grouped_gemm_dglu_wrapper_sm100).parameters - except (TypeError, ValueError): - return False - return "generate_dbias" in params - - def _compute_grad_params( fc_op, ctx, @@ -300,10 +284,11 @@ def grouped_gemm_quant_kernel(cls) -> Callable: @functools.lru_cache(maxsize=None) def grouped_gemm_wgrad_kernel(cls) -> Optional[Callable]: """CuTe DSL kernel for grouped GEMM wgrad on SM100+. - Returns ``None`` when the cuDNN front-end package is older than - 1.23.0. + + Returns ``None`` when the environment variable + ``NVTE_DISABLE_CUTEDSL_WGRAD_FUSED_GROUPED_MLP`` is set to ``1``. """ - if not _nvidia_cudnn_frontend_supports_wgrad(): + if int(os.environ.get("NVTE_DISABLE_CUTEDSL_WGRAD_FUSED_GROUPED_MLP", "0")) >= 1: return None from cudnn import grouped_gemm_wgrad_wrapper_sm100 # pylint: disable=no-name-in-module @@ -317,6 +302,8 @@ def is_supported(cls) -> bool: return False if get_device_compute_capability()[0] != 10: return False + if not _cudnn_frontend_version_supported(): + return False try: cls.grouped_gemm_dglu_kernel() cls.grouped_gemm_quant_kernel() @@ -324,13 +311,6 @@ def is_supported(cls) -> bool: return False return True - @classmethod - def is_fc1_bias_supported(cls) -> bool: - """Whether cudnn-frontend exposes ``generate_dbias`` on the dGLU SM100 wrapper (FC1 bias grad only).""" - if not cls.is_supported(): - return False - return _dglu_wrapper_has_generate_dbias_arg() - def __init__( self, *, diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index cad31e2c5..599e5f96a 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -7,7 +7,6 @@ from __future__ import annotations from collections.abc import Callable, Iterable import functools -import inspect import os from typing import Any, Optional @@ -24,6 +23,7 @@ from ..fuser import register_forward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext from .._common import ( + _cudnn_frontend_version_supported, fuse_grouped_mlp_ops, is_quantized_tensor, maybe_dequantize, @@ -76,6 +76,8 @@ def is_supported(cls) -> bool: return False if get_device_compute_capability()[0] != 10: return False + if not _cudnn_frontend_version_supported(): + return False try: cls.grouped_gemm_glu_kernel() cls.grouped_gemm_quant_kernel() @@ -83,42 +85,6 @@ def is_supported(cls) -> bool: return False return True - @classmethod - @functools.lru_cache(maxsize=1) - def is_fc1_bias_supported(cls) -> bool: - """Whether cudnn-frontend exposes ``bias_tensor`` on the grouped GEMM GLU SM100 wrapper (FC1).""" - if not cls.is_supported(): - return False - try: - from cudnn import ( - grouped_gemm_glu_wrapper_sm100, - ) # pylint: disable=import-outside-toplevel - except ImportError: - return False - try: - params = inspect.signature(grouped_gemm_glu_wrapper_sm100).parameters - except (TypeError, ValueError): - return False - return "bias_tensor" in params - - @classmethod - @functools.lru_cache(maxsize=1) - def is_fc2_bias_supported(cls) -> bool: - """Whether cudnn-frontend exposes ``bias_tensor`` on the grouped GEMM Quant SM100 wrapper (FC2).""" - if not cls.is_supported(): - return False - try: - from cudnn import ( - grouped_gemm_quant_wrapper_sm100, - ) # pylint: disable=import-outside-toplevel - except ImportError: - return False - try: - params = inspect.signature(grouped_gemm_quant_wrapper_sm100).parameters - except (TypeError, ValueError): - return False - return "bias_tensor" in params - def __init__( self, *, @@ -433,6 +399,7 @@ def fuser_forward( "sfa_tensor": fc1_kernel_out["sfd_row_tensor"], "padded_offsets": split_points, "alpha_tensor": alpha_tensor.float(), + "bias_tensor": fc2_bias_packed, "norm_const_tensor": None, "prob_tensor": fc2_scales_tensor, "acc_dtype": torch.float32, @@ -442,8 +409,6 @@ def fuser_forward( "current_stream": current_stream, "use_dynamic_sched": True, } - if self.is_fc2_bias_supported(): - fc2_quant_kwargs["bias_tensor"] = fc2_bias_packed if fc2_op.single_grouped_weight: # Clone and swizzle scales for GEMM (original stays unmodified for save_for_backward) From 36fc33603ce4baa6a2054f31b4151013ff48374a Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Fri, 1 May 2026 02:23:28 -0400 Subject: [PATCH 029/180] [PyTorch] Add workaround for cuteDSL stride requirement for zero-token expert (#2947) Add workaround for cuteDSL stride requirement for zero token expert Signed-off-by: Kirthi Shankar Sivamani --- .../pytorch/ops/fused/backward_grouped_mlp.py | 44 ++++++++++++++----- 1 file changed, 33 insertions(+), 11 deletions(-) diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 70d0d7469..b07ebb73e 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -56,19 +56,41 @@ def _cudnn_compute_wgrad( fp8_dtype = torch.float8_e4m3fn - # a_tensor = DY^T = (out_features, total_tokens) row-major - a_tensor = grouped_dy.columnwise_data.view(dtype=fp8_dtype).view(total_tokens, out_features).T - # b_tensor = X = (total_tokens, in_features) column-major - b_tensor = grouped_x.columnwise_data.view(dtype=fp8_dtype).view(total_tokens, in_features) - sfa_leading_dim = ((out_features + 127) // 128) * 128 sfb_leading_dim = ((in_features + 127) // 128) * 128 - sfa_tensor = grouped_dy.columnwise_scale_inv.view(sfa_leading_dim, -1).view( - dtype=torch.float8_e8m0fnu - ) - sfb_tensor = grouped_x.columnwise_scale_inv.view(sfb_leading_dim, -1).view( - dtype=torch.float8_e8m0fnu - ) + + if total_tokens == 0: + # A workaround for the case with zero-token experts. + # Even for this case, cuteDSL still requires the same + # stride requirements for the input and scale tensors. + device = grouped_dy.columnwise_data.device + a_tensor = torch.empty_strided((out_features, 0), (16, 1), dtype=fp8_dtype, device=device) + b_tensor = torch.empty_strided( + (0, in_features), (in_features, 1), dtype=fp8_dtype, device=device + ) + sfa_tensor = torch.empty_strided( + (sfa_leading_dim, 0), + (16, 1), + dtype=torch.float8_e8m0fnu, + device=device, + ) + sfb_tensor = torch.empty_strided( + (sfb_leading_dim, 0), + (16, 1), + dtype=torch.float8_e8m0fnu, + device=device, + ) + else: + a_tensor = ( + grouped_dy.columnwise_data.view(dtype=fp8_dtype).view(total_tokens, out_features).T + ) + b_tensor = grouped_x.columnwise_data.view(dtype=fp8_dtype).view(total_tokens, in_features) + sfa_tensor = grouped_dy.columnwise_scale_inv.view(sfa_leading_dim, -1).view( + dtype=torch.float8_e8m0fnu + ) + sfb_tensor = grouped_x.columnwise_scale_inv.view(sfb_leading_dim, -1).view( + dtype=torch.float8_e8m0fnu + ) # Prepare wgrad output if single_grouped_weight: From 7e8bc98b26bfd28f733bc74f6054e5fdeb655145 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Fri, 1 May 2026 01:30:45 -0700 Subject: [PATCH 030/180] [Core] Remove unused NVFP4 quantize kernel (#2946) Remove unused NVFP4 quantize kernel Signed-off-by: Tim Moon Co-authored-by: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> --- .../common/cast/dispatch/quantize.cuh | 1 - .../common/cast/nvfp4/quantize_nvfp4.cuh | 681 ------------------ 2 files changed, 682 deletions(-) delete mode 100644 transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 8d985f64f..5d0d3c28e 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -21,7 +21,6 @@ #include "../mxfp8/group_quantize_mxfp8.cuh" #include "../mxfp8/quantize_mxfp8.cuh" #include "../nvfp4/group_quantize_transpose_nvfp4.cuh" -#include "../nvfp4/quantize_nvfp4.cuh" #include "../nvfp4/quantize_transpose_nvfp4.cuh" namespace transformer_engine { diff --git a/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh deleted file mode 100644 index ec80924df..000000000 --- a/transformer_engine/common/cast/nvfp4/quantize_nvfp4.cuh +++ /dev/null @@ -1,681 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -/*! \file quantize_nvfp4.cuh - * \brief CUDA kernels to cast to NVFP4. - */ - -#ifndef TRANSFORMER_ENGINE_QUANTIZE_NVFP4_CUH_ -#define TRANSFORMER_ENGINE_QUANTIZE_NVFP4_CUH_ - -#include -#include -#include -#include - -#include "../../common.h" -#include "../../util/math.h" -#include "../../util/ptx.cuh" -#include "../../utils.cuh" -#include "core_nvfp4.cuh" - -namespace transformer_engine { -namespace dispatch { -namespace nvfp4 { -namespace quantize_kernel { - -using namespace ptx; -using namespace quantization_SF; -using namespace core; - -constexpr size_t SCALE_DIM_Y = 32; -constexpr size_t SCALE_DIM_X = 16; - -constexpr size_t BUFFS_NUM = 2; -constexpr size_t BUFF_DIM_Y = 32; - -constexpr size_t PACK_SIZE = 8; -constexpr size_t WAVES = SCALE_DIM_X / PACK_SIZE; - -// Number of 4-bit elements that span 32 banks (4-byte each) of shared memory -constexpr size_t TOTAL_BANKS_WIDTH = (32 * 4 * 8) / 4; // 256 - -// Number of threads (rowwise scaling) that span 32 banks (4-byte banks) of shared memory -constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM_X; // 8 = 128 / 16 - -#define DIRECT_SCALING_FACTORS_STORE 1 - -template -__global__ void __launch_bounds__(THREADS_PER_CHUNK) - quantize_nvfp4_kernel(const __grid_constant__ CUtensorMap tensor_map_input, - const __grid_constant__ CUtensorMap tensor_map_output_rowwise, - const __grid_constant__ CUtensorMap tensor_map_output_colwise, - fp8e4m3 *const scales_rowwise_e4m3, e8m0_t *const scales_colwise_e8m0, - const float *noop, float *const amax_ptr, - const float *const nvfp4_second_stage_scale_ptr, const size_t rows, - const size_t cols, const size_t scale_stride_rowwise, - const size_t scale_stride_colwise) { -#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) - constexpr bool ROWWISE_SCALING = true; - constexpr bool NO_ACTIVATIONS_NOT_FP32_INPUT = - (!COMPUTE_ACTIVATIONS) && (!std::is_same_v); - - using IType2 = typename ptx::FPx2; - - if constexpr (!COMPUTE_ACTIVATIONS) { - if (noop != nullptr && noop[0] == 1.0f) { - return; - } - } - constexpr size_t NVFP4_SCALING_FACTORS_PER_CHUNK_ROW = CHUNK_DIM_X / SCALE_DIM_X; - constexpr size_t THREADS_X_ROWWISE = NVFP4_SCALING_FACTORS_PER_CHUNK_ROW; - constexpr size_t THREADS_Y_ROWWISE = THREADS_PER_CHUNK / THREADS_X_ROWWISE; - - static_assert(BUFF_DIM_Y >= SCALE_DIM_Y && - "Number of buffer rows must be greater or equal to the size of the columwise " - "scaling block\0"); - static_assert(CHUNK_DIM_Y >= BUFF_DIM_Y); - static_assert(BUFF_DIM_Y >= THREADS_Y_ROWWISE && - "Number of buffer rows must be greater or equal to the number of rowwise " - "processing threads in Y dimension\0"); - - constexpr size_t BUFF_IN_DIM_X = CHUNK_DIM_X; - constexpr size_t BUFF_OUT_DIM_X = (CHUNK_DIM_X * 4) / 8; // Holds 2 elements of 4-bit size - constexpr size_t BUFF_IN_DIM = BUFF_DIM_Y * BUFF_IN_DIM_X; - constexpr size_t BUFF_OUT_DIM = BUFF_DIM_Y * BUFF_OUT_DIM_X; - - constexpr size_t STAGES = CHUNK_DIM_Y / BUFF_DIM_Y; - - constexpr size_t ITERATIONS_ROWWISE = BUFF_DIM_Y / THREADS_Y_ROWWISE; - // static_assert(THREADS_PER_CHUNK >= CHUNK_DIM_X); // there should be a sufficient number of - // // threads to process one row in a single iteration - - constexpr bool IS_CACHED_ACT_OP = COMPUTE_ACTIVATIONS && ROWWISE_SCALING && COLWISE_SCALING; - - const int block_offset_Y = blockIdx.y * CHUNK_DIM_Y; - const int block_offset_X = blockIdx.x * CHUNK_DIM_X; - const int scales_block_offset_Y_rowwise = blockIdx.y * CHUNK_DIM_Y; - const int scales_block_offset_X_rowwise = blockIdx.x * CHUNK_DIM_X / SCALE_DIM_X; - const int scales_block_offset_Y_colwise = blockIdx.y * CHUNK_DIM_Y / SCALE_DIM_Y; - const int scales_block_offset_X_colwise = blockIdx.x * CHUNK_DIM_X; - - const int tid_Y_rowwise = threadIdx.x / THREADS_X_ROWWISE; - const int tid_X_rowwise = threadIdx.x % THREADS_X_ROWWISE; - const int tid_Y_colwise = 0; - const int tid_X_colwise = threadIdx.x; - - const int thread_offset_Y_rowwise = tid_Y_rowwise; - const int thread_offset_X_rowwise = tid_X_rowwise * SCALE_DIM_X; - const int thread_offset_Y_colwise = tid_Y_colwise; - const int thread_offset_X_colwise = tid_X_colwise; // Each thread processes two adjacent elements - - const int row_base_rowwise = block_offset_Y + thread_offset_Y_rowwise; - const int row_base_colwise = block_offset_Y + thread_offset_Y_colwise; - const int col_base_colwise = block_offset_X + thread_offset_X_colwise; - - const bool col_out_of_bounds_colwise = (col_base_colwise >= cols); - - const int scales_offset_Y_rowwise = scales_block_offset_Y_rowwise + tid_Y_rowwise; - const int scales_offset_X_rowwise = scales_block_offset_X_rowwise + tid_X_rowwise; - const int scales_offset_Y_colwise = scales_block_offset_Y_colwise + tid_Y_colwise; - const int scales_offset_X_colwise = scales_block_offset_X_colwise + tid_X_colwise; - - const bool rowwise_scale_is_within_bounds = scales_offset_X_rowwise < cols; - const bool colwise_scale_is_within_bounds = scales_offset_X_colwise < cols; - - // helps resolving bank conflicts in shmem - const int thread_lane = threadIdx.x % THREADS_PER_WARP; - const int bank_group = thread_lane / THREADS_PER_BANK; - - constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_IN_DIM_X; - constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; - - constexpr size_t buff_size_aligned_in = - DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(IType), TMA_SHMEM_ALIGNMENT); - constexpr size_t buff_size_aligned_out_nvfp4 = - DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); - constexpr size_t buff_size_aligned_out_mxfp8 = - DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(OType), TMA_SHMEM_ALIGNMENT); - - constexpr size_t in_mem = buff_size_aligned_in; - - constexpr size_t out_mem_rowwise_data = (ROWWISE_SCALING ? buff_size_aligned_out_nvfp4 : 0); - constexpr size_t out_mem_colwise_data = (COLWISE_SCALING ? buff_size_aligned_out_mxfp8 : 0); - - extern __shared__ char dynamic_shmem[]; - uintptr_t base_shmem_ptr = reinterpret_cast(dynamic_shmem); - // Manually align dynamic SHMEM per TMA requirements using padding - // __align__(128) Does not guarantee the pointer to be aligned! - uintptr_t dshmem = (base_shmem_ptr + TMA_SHMEM_ALIGNMENT - 1) & - ~(static_cast(TMA_SHMEM_ALIGNMENT - 1)); - - // The destination shared memory buffer of a bulk tensor operation should be 16-byte aligned - IType *in_sh = reinterpret_cast(dshmem); - fp4e2m1x2 *out_rowwise_data_sh = reinterpret_cast(dshmem + in_mem); - OType *out_colwise_data_sh = reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data); - fp8e4m3 *out_rowwise_scales_sh = - reinterpret_cast(dshmem + in_mem + out_mem_rowwise_data + out_mem_colwise_data); - (void)out_rowwise_scales_sh; // Suppress unused variable warning - IType *cached_act_sh = in_sh; // in_sh is used as a cache buffer - - constexpr int shmem_buff_size = buff_size_aligned_in / BUFFS_NUM; - - const bool is_master_thread = (threadIdx.x == 0); - - // Compute a global encoding/decoding scaling factor for all S_dec_b - const float S_enc = - (nvfp4_second_stage_scale_ptr == nullptr) ? 1.0f : 1.0f / (*nvfp4_second_stage_scale_ptr); - - float thread_amax = 0.0f; - -// Initialize shared memory barrier with the number of threads participating in the barrier. -#pragma nv_diag_suppress static_var_with_dynamic_init - __shared__ alignas(8) uint64_t mbar[STAGES]; - - initialize_barriers(mbar, is_master_thread); - - copy_2d_to_shared(&in_sh[0], &tensor_map_input, block_offset_X, block_offset_Y, shmem_buff_size, - &mbar[0], is_master_thread); - -#pragma unroll - for (int stage = 0; stage < STAGES; ++stage) { - const int buff = stage % BUFFS_NUM; - const int next_stage = stage + 1; - const int stage_offset_Y = stage * BUFF_DIM_Y; - - const int buff_offset_in = buff * BUFF_IN_DIM; - const int buff_offset_out = buff * BUFF_OUT_DIM; - - if (next_stage < STAGES) { - // Wait for TMA transfer to have finished reading shared memory. - // I.e. the buffer is ready to be written to - ptx::cp_async_bulk_wait_group_read<1>(); - - const int next_buff = next_stage % BUFFS_NUM; - const int next_stage_offset_Y = next_stage * BUFF_DIM_Y; - const int global_offset_Y = block_offset_Y + next_stage_offset_Y; - const int global_offset_X = block_offset_X; - const int next_buff_offset = next_buff * BUFF_IN_DIM; - - copy_2d_to_shared(&in_sh[next_buff_offset], &tensor_map_input, global_offset_X, - global_offset_Y, shmem_buff_size, &mbar[next_stage], is_master_thread); - } - - ptx::fence_proxy_async_shared_cta(); - - // Wait for the data to have arrived - ptx::mbarrier_wait_parity(&mbar[stage], 0); - - float block_amax = 0.0f; - if constexpr (COLWISE_SCALING) { - const int shmem_offset_base_colwise = buff_offset_in + tid_X_colwise; - - block_amax = 0.0f; - float in_compute_colwise[SCALE_DIM_Y]; - IType in_colwise_IType[SCALE_DIM_Y]; - - // 1. Read/Compute elements. Find MXFP8-block AMAX - if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { - IType block_amax_f16 = static_cast(0.0f); -#pragma unroll - for (int i = 0; i < SCALE_DIM_Y; ++i) { - const int shmem_offset_colwise = shmem_offset_base_colwise + i * BUFF_IN_DIM_X; - in_colwise_IType[i] = in_sh[shmem_offset_colwise]; - block_amax_f16 = __hmax(block_amax_f16, __habs(in_colwise_IType[i])); - } - block_amax = static_cast(block_amax_f16); - } else { -#pragma unroll - for (int i = 0; i < SCALE_DIM_Y; ++i) { - const int shmem_offset_colwise = shmem_offset_base_colwise + i * BUFF_IN_DIM_X; - - float elt = static_cast(in_sh[shmem_offset_colwise]); - if constexpr (COMPUTE_ACTIVATIONS) { - elt = OP(elt, {}); - } - // Numerical truncation: Downcast to IType (BF16/FP16), then upcast it back to FP32 - if constexpr (!std::is_same_v) { - elt = static_cast(static_cast(elt)); - } - // Cache computed activations to avoid computing them again in the 2nd pass along another dimension - if constexpr (IS_CACHED_ACT_OP) { - cached_act_sh[shmem_offset_colwise] = static_cast(elt); - } - - if constexpr (COMPUTE_ACTIVATIONS) { - const bool row_out_of_bounds_colwise = (row_base_colwise + stage_offset_Y + i >= rows); - const bool out_of_bounds = (col_out_of_bounds_colwise || row_out_of_bounds_colwise); - if (!out_of_bounds) { - block_amax = fmaxf(block_amax, fabsf(elt)); - } - } else { - // If no activation, elt is 0 so we can safely do this - block_amax = fmaxf(block_amax, fabsf(elt)); - } - in_compute_colwise[i] = elt; - } - } - // 2. Compute E8M0 scaling factor - const e8m0_t biased_exponent = - ptx::float_to_e8m0(block_amax * Quantized_Limits::max_norm_rcp); - - const int global_scales_offset_Y = scales_offset_Y_colwise + stage; - const int global_scales_offset_X = scales_offset_X_colwise; - const int scale_idx = global_scales_offset_Y * scale_stride_colwise + global_scales_offset_X; - if (colwise_scale_is_within_bounds) { - scales_colwise_e8m0[scale_idx] = biased_exponent; - } - const float block_scale_inverse = ptx::exp2f_rcp(biased_exponent); - -// 3. Scale elements -#pragma unroll - for (int i = 0; i < SCALE_DIM_Y; ++i) { - float in; - if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { - in = static_cast(in_colwise_IType[i]); - } else { - in = in_compute_colwise[i]; - } - const float scaled_out = in * block_scale_inverse; - - const int shmem_offset_elt = shmem_offset_base_colwise + i * BUFF_IN_DIM_X; - out_colwise_data_sh[shmem_offset_elt] = static_cast(scaled_out); - } - } - - if constexpr (ROWWISE_SCALING) { - const int stage_rowwise_scales_offset_Y = stage * BUFF_DIM_Y; -#pragma unroll - for (int it = 0; it < ITERATIONS_ROWWISE; ++it) { - const int it_thread_offset_Y_rowwise = thread_offset_Y_rowwise + it * THREADS_Y_ROWWISE; - - const int shmem_offset_base_rowwise_in = - buff_offset_in + it_thread_offset_Y_rowwise * BUFF_IN_DIM_X; - const int shmem_offset_base_rowwise_out = - buff_offset_out + it_thread_offset_Y_rowwise * BUFF_OUT_DIM_X; - - const int it_offset_Y = stage_offset_Y + it * THREADS_Y_ROWWISE; - - block_amax = 0.0f; - float in_compute_rowwise[SCALE_DIM_X]; - Vec in_cached[WAVES]; - - // used as an IType container for BF16/FP16 --> NVFP4 CAST ONLY - Vec in_IType[WAVES]; - - // 1. Read/Compute elements. Find NVFP4-block AMAX - if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { - IType2 thread_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; -#pragma unroll - for (int w = 0; w < WAVES; ++w) { - const int swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; - const int swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const int shmem_offset_rowwise = shmem_offset_base_rowwise_in + swizzled_thread_idx; - // Load elements - in_IType[w].load_from(&in_sh[shmem_offset_rowwise]); -#pragma unroll - for (int e = 0; e < PACK_SIZE / 2; ++e) { - ptx::abs_max_2x(thread_amax_2x, thread_amax_2x, in_IType[w].data.elt[e]); - } - } - block_amax = - static_cast(__hmax(__habs(thread_amax_2x.x), __habs(thread_amax_2x.y))); - } else if constexpr (IS_CACHED_ACT_OP) { - // ensures that all writes to cache made in the section above are visible to all threads - __syncthreads(); - IType2 thread_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; -#pragma unroll - for (int w = 0; w < WAVES; ++w) { - const int swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; - const int swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const int shmem_offset_rowwise = shmem_offset_base_rowwise_in + swizzled_thread_idx; - - const bool row_out_of_bounds_rowwise = (row_base_rowwise + it_offset_Y >= rows); - const bool swizzled_col_out_of_bounds = (block_offset_X + swizzled_thread_idx >= cols); - const bool out_of_bounds = (row_out_of_bounds_rowwise || swizzled_col_out_of_bounds); - - // Load cached elements - in_cached[w].load_from(&cached_act_sh[shmem_offset_rowwise]); - // Since TMA requirement for the data alignment is 16B (i.e. cols % 8 == 0, in case of BF16 elements) - // only single check (w.r.t. column direction) is sufficient to be sure the entire wave is inside the boundaries - if (!out_of_bounds) { - if constexpr (std::is_same_v) { -#pragma unroll - for (int e = 0; e < PACK_SIZE; ++e) { - block_amax = fmaxf(block_amax, fabsf(in_cached[w].data.elt[e])); - } - } else { -#pragma unroll - for (int e = 0; e < PACK_SIZE; e += 2) { - const IType2 in_cached_2x = {in_cached[w].data.elt[e], - in_cached[w].data.elt[e + 1]}; - ptx::abs_max_2x(thread_amax_2x, thread_amax_2x, in_cached_2x); - } - } - } - } - if constexpr (!std::is_same_v) { - block_amax = - static_cast(__hmax(__habs(thread_amax_2x.x), __habs(thread_amax_2x.y))); - } - } else { -#pragma unroll - for (int w = 0; w < WAVES; ++w) { - const int swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; - const int swizzled_thread_idx = thread_offset_X_rowwise + swizzled_group_idx; - const int shmem_offset_rowwise = shmem_offset_base_rowwise_in + swizzled_thread_idx; - - Vec in; - Vec act_in; - - in.load_from(&in_sh[shmem_offset_rowwise]); -#pragma unroll - for (int e = 0; e < PACK_SIZE; ++e) { - const int j = w * PACK_SIZE + e; - // Compute element - float elt = static_cast(in.data.elt[e]); - if constexpr (COMPUTE_ACTIVATIONS) { - elt = OP(elt, {}); - } - // Numerical truncation: Downcast to IType (BF16/FP16), then upcast it back to FP32 - if constexpr (!std::is_same_v) { - elt = static_cast(static_cast(elt)); - } - if constexpr (COMPUTE_ACTIVATIONS) { - const bool row_out_of_bounds_rowwise = (row_base_rowwise + it_offset_Y >= rows); - const bool swizzled_col_out_of_bounds = - (block_offset_X + swizzled_thread_idx >= cols); - const bool out_of_bounds = - (row_out_of_bounds_rowwise || swizzled_col_out_of_bounds); - if (!out_of_bounds) { - block_amax = fmaxf(block_amax, fabsf(elt)); - } - } else { - // If no activation, elt is 0 so we can safely do this - block_amax = fmaxf(block_amax, fabsf(elt)); - } - in_compute_rowwise[j] = elt; - } - } - } - - // 2. Compute E4M3 scaling factor - const fp8e4m3 S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax, S_enc); - -#if DIRECT_SCALING_FACTORS_STORE - // Check boundaries - if (rowwise_scale_is_within_bounds) { - const int scales_offset_Y = - scales_offset_Y_rowwise + stage_rowwise_scales_offset_Y + it * THREADS_Y_ROWWISE; - const int scales_offset_X = scales_offset_X_rowwise; - const int scale_idx_global = scales_offset_Y * scale_stride_rowwise + scales_offset_X; - scales_rowwise_e4m3[scale_idx_global] = S_dec_b_fp8; - } -#else - const int shmem_scales_offset_Y = - stage_rowwise_scales_offset_Y + it * THREADS_Y_ROWWISE + tid_Y_rowwise; - const int shmem_scales_offset_X = tid_X_rowwise; - const int scale_idx = - shmem_scales_offset_Y * NVFP4_SCALING_FACTORS_PER_CHUNK_ROW + shmem_scales_offset_X; - out_rowwise_scales_sh[scale_idx] = S_dec_b_fp8; -#endif - // Compute "correct" per-block encoding scaling factor - const float block_scale_inverse = - __fdiv_rn(S_enc, static_cast(S_dec_b_fp8)); // S_enc_b_fp8 - -// 3. Scale elements -#pragma unroll - for (int w = 0; w < WAVES; ++w) { - Vec out; // Vec out; -#pragma unroll - for (int e = 0; e < PACK_SIZE / 4; ++e) { - IType2 in01; - IType2 in23; - if constexpr (NO_ACTIVATIONS_NOT_FP32_INPUT) { - in01 = in_IType[w].data.elt[2 * e]; - in23 = in_IType[w].data.elt[2 * e + 1]; - } else if constexpr (IS_CACHED_ACT_OP) { - in01.x = in_cached[w].data.elt[4 * e]; - in01.y = in_cached[w].data.elt[4 * e + 1]; - in23.x = in_cached[w].data.elt[4 * e + 2]; - in23.y = in_cached[w].data.elt[4 * e + 3]; - } else { - const int j = w * PACK_SIZE + 4 * e; - in01.x = in_compute_rowwise[j]; - in01.y = in_compute_rowwise[j + 1]; - in23.x = in_compute_rowwise[j + 2]; - in23.y = in_compute_rowwise[j + 3]; - } - fp4e2m1x4 &out_quad = reinterpret_cast(out.data.elt[e]); - ptx::mul_cvt_4x(out_quad, in01, in23, block_scale_inverse); - } - const int swizzled_group_idx = ((w + bank_group) * PACK_SIZE) % SCALE_DIM_X; - const int swizzled_idx = swizzled_group_idx + thread_offset_X_rowwise; - const int shmem_offset_rowwise = shmem_offset_base_rowwise_out + swizzled_idx / 2; - out.store_to(&out_rowwise_data_sh[shmem_offset_rowwise]); - } - } - } - - __builtin_assume(thread_amax >= 0); - __builtin_assume(block_amax >= 0); - thread_amax = fmaxf(thread_amax, block_amax); - - // Wait for shared memory writes to be visible to TMA engine. - ptx::fence_proxy_async_shared_cta(); - __syncthreads(); - // After syncthreads, writes by all threads are visible to TMA engine. - - // Initiate TMA transfer to copy shared memory to global memory - if (is_master_thread) { - const int global_offset_Y = block_offset_Y + stage_offset_Y; - const int global_offset_X = block_offset_X; - const int buff_offset_nvfp4 = buff * BUFF_OUT_DIM; - const int buff_offset_mxfp8 = buff * BUFF_IN_DIM; - - if constexpr (ROWWISE_SCALING) { - ptx::cp_async_bulk_tensor_2d_shared_to_global( - reinterpret_cast(&tensor_map_output_rowwise), global_offset_X, - global_offset_Y, reinterpret_cast(&out_rowwise_data_sh[buff_offset_nvfp4])); - } - if constexpr (COLWISE_SCALING) { - ptx::cp_async_bulk_tensor_2d_shared_to_global( - reinterpret_cast(&tensor_map_output_colwise), global_offset_X, - global_offset_Y, reinterpret_cast(&out_colwise_data_sh[buff_offset_mxfp8])); - } - - // Create a "bulk async-group" out of the previous bulk copy operation. - ptx::cp_async_bulk_commit_group(); - } - } - -#if !DIRECT_SCALING_FACTORS_STORE - // Vectorized store of scaling factors. - // Each thread stores multiple scaling factors in one store instruction. - if constexpr (ROWWISE_SCALING) { - // Number of scaling factors = CHUNK_DIM_X / SCALE_DIM_X - const int scales_offset_Y_rowwise = scales_block_offset_Y_rowwise + threadIdx.x; - const int scales_offset_X_rowwise = scales_block_offset_X_rowwise; - const int scale_idx_global = - scales_offset_Y_rowwise * scale_stride_rowwise + scales_offset_X_rowwise; - const int scale_idx_shmem = threadIdx.x * NVFP4_SCALING_FACTORS_PER_CHUNK_ROW; - - if ((threadIdx.x < CHUNK_DIM_Y) && (scales_offset_Y_rowwise < rows) && - (scales_offset_X_rowwise < (cols / SCALE_DIM_X))) { - using ScalesVec_t = Vec; - const ScalesVec_t &scales = - *reinterpret_cast(&out_rowwise_scales_sh[scale_idx_shmem]); - scales.store_to(&scales_rowwise_e4m3[scale_idx_global]); - } - } -#endif - - float chunk_amax = 0.0f; - if (amax_ptr != nullptr) { - const int warp_id = threadIdx.x / THREADS_PER_WARP; - // Reduce the amax over the block - chunk_amax = reduce_max(thread_amax, warp_id); - } - - if (is_master_thread && amax_ptr != nullptr) { - atomicMaxFloat(amax_ptr, chunk_amax); - } - - destroy_barriers(mbar, is_master_thread); -#endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) -} -} // namespace quantize_kernel - -// This kernel supports only two scaling cases: -// 1. r16c0 - Rowwise NVFP4 -// 2. r16c32 - Rowwise NVFP4 AND Colwise MXFP8 -inline void quantize(const Tensor &input, const Tensor *noop, Tensor *output, cudaStream_t stream) { -#if FP4_TYPE_SUPPORTED - using namespace quantize_kernel; - using namespace ptx; - checkCuDriverContext(stream); - - constexpr bool COMPUTE_ACTIVATIONS = false; - using ParamOP = Empty; - constexpr float (*OP)(float, const ParamOP &) = nullptr; - - NVTE_CHECK(output->has_data(), "NVFP4 Output tensor must be allocated."); - NVTE_CHECK(input.has_data(), "Cannot quantize tensor without rowwise data."); - - NVTE_CHECK(is_fp4_dtype(output->data.dtype), "Output must have FP4 type."); - NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); - NVTE_CHECK(!output->with_gemm_swizzled_scales, "Output must have scales in compact format."); - - bool use_colwise_scaling = output->has_columnwise_data(); - if (use_colwise_scaling) { - NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, - "Columnwise scaling tensor must be allocated"); - } - CheckNoopTensor(*noop, "cast_noop"); - - const size_t rows = input.flat_first_dim(); - const size_t cols = input.flat_last_dim(); - - constexpr size_t CHUNK_DIM_Y = 128; - constexpr size_t CHUNK_DIM_X = 128; - constexpr size_t THREADS_PER_CHUNK = 128; - - constexpr size_t BUFF_DIM_X = CHUNK_DIM_X; - - const size_t blocks_Y = DIVUP(rows, CHUNK_DIM_Y); - const size_t blocks_X = DIVUP(cols, CHUNK_DIM_X); - const dim3 grid(blocks_X, blocks_Y); - const size_t block_size = THREADS_PER_CHUNK; - - const size_t scale_stride_rowwise = output->scale_inv.shape[1]; - const size_t scale_stride_colwise = - use_colwise_scaling ? output->columnwise_scale_inv.shape[1] : 1; - - fp8e4m3 *const scales_rowwise_e4m3_ptr = reinterpret_cast(output->scale_inv.dptr); - e8m0_t *const scales_colwise_e8m0_ptr = - use_colwise_scaling ? reinterpret_cast(output->columnwise_scale_inv.dptr) : nullptr; - - const ScalingType scaling_type = - use_colwise_scaling ? ScalingType::BIDIMENSIONAL : ScalingType::ROWWISE; - - float *const amax_ptr = reinterpret_cast(output->amax.dptr); - const float *noop_ptr = reinterpret_cast(noop->data.dptr); - const float *const nvfp4_second_stage_scale_ptr = - reinterpret_cast(output->scale.dptr); - - // Output data type is only required for the column-wise MXFP8 scaling. - // It has no effect for the row-wise NVFP4 scaling, but is set to the default E4M3 for the macros to work - const DType output_data_type = - use_colwise_scaling ? output->columnwise_data.dtype : DType::kFloat8E4M3; - - TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( - input.dtype(), IType, - TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY( - output_data_type, OType, alignas(64) CUtensorMap tensor_map_input{}; - alignas(64) CUtensorMap tensor_map_output_rowwise{}; - alignas(64) CUtensorMap tensor_map_output_colwise{}; - - create_2D_tensor_map(tensor_map_input, input.data, rows, cols, BUFF_DIM_Y, BUFF_DIM_X, - cols, 0, sizeof(IType) * 8); - - create_2D_tensor_map(tensor_map_output_rowwise, output->data, rows, cols, BUFF_DIM_Y, - BUFF_DIM_X, cols, 0, 4); - - if (use_colwise_scaling) { - create_2D_tensor_map(tensor_map_output_colwise, output->columnwise_data, rows, cols, - BUFF_DIM_Y, BUFF_DIM_X, cols, 0, sizeof(OType) * 8); - } - - constexpr size_t buff_elems = BUFF_DIM_Y * BUFF_DIM_X; - constexpr size_t buff_elems_total = BUFFS_NUM * buff_elems; - constexpr size_t buff_size_aligned_in = - DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(IType), TMA_SHMEM_ALIGNMENT); - constexpr size_t buff_size_aligned_out_nvfp4 = - DIVUP_TO_MULTIPLE((buff_elems_total * 4) / 8, TMA_SHMEM_ALIGNMENT); - constexpr size_t buff_size_aligned_out_mxfp8 = - DIVUP_TO_MULTIPLE(buff_elems_total * sizeof(OType), TMA_SHMEM_ALIGNMENT); - constexpr size_t buff_size_nvfp4_scales = - (CHUNK_DIM_Y * CHUNK_DIM_X) / 16 * sizeof(fp8e4m3); - constexpr size_t buff_size_mxfp8_scales = - (CHUNK_DIM_Y * CHUNK_DIM_X) / 32 * sizeof(e8m0_t); - - constexpr size_t in_mem = buff_size_aligned_in; - - const size_t out_rowwise_data_mem = buff_size_aligned_out_nvfp4; - const size_t out_colwise_data_mem = use_colwise_scaling ? buff_size_aligned_out_mxfp8 : 0; - - const size_t out_rowwise_scales_mem = buff_size_nvfp4_scales; - const size_t out_colwise_scales_mem = use_colwise_scaling ? buff_size_mxfp8_scales : 0; - - const size_t out_mem = out_rowwise_data_mem + out_colwise_data_mem + - out_rowwise_scales_mem + out_colwise_scales_mem + - TMA_SHMEM_ALIGNMENT; - - const size_t dshmem_size = in_mem + out_mem; - - switch (scaling_type) { - case ScalingType::ROWWISE: { - auto kernel = - quantize_nvfp4_kernel; - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - dshmem_size); - - kernel<<>>( - tensor_map_input, tensor_map_output_rowwise, tensor_map_output_colwise, - scales_rowwise_e4m3_ptr, scales_colwise_e8m0_ptr, noop_ptr, amax_ptr, - nvfp4_second_stage_scale_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); - break; - } - case ScalingType::BIDIMENSIONAL: { - auto kernel = - quantize_nvfp4_kernel; - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - dshmem_size); - - kernel<<>>( - tensor_map_input, tensor_map_output_rowwise, tensor_map_output_colwise, - scales_rowwise_e4m3_ptr, scales_colwise_e8m0_ptr, noop_ptr, amax_ptr, - nvfp4_second_stage_scale_ptr, rows, cols, scale_stride_rowwise, - scale_stride_colwise); - break; - } - } NVTE_CHECK_CUDA(cudaGetLastError());); // NOLINT(*) - ); // NOLINT(*) -#else - NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); -#endif // FP4_TYPE_SUPPORTED -} - -} // namespace nvfp4 -} // namespace dispatch -} // namespace transformer_engine - -#endif // TRANSFORMER_ENGINE_QUANTIZE_NVFP4_CUH_ From 360779bab99cb55612ac2361860b5ebe06e493ee Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Fri, 1 May 2026 14:28:31 -0700 Subject: [PATCH 031/180] [JAX] Calculate seqlens and offsets in O(T) space instead of O(T*T) space for THD sequences (#2522) * Get seqlens and offsets in O(N) space instead of O(N*N) space Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Re enable fast causal path Signed-off-by: Kshitij Lakhani * Fix: seqoffsets calculation for THD Signed-off-by: Kshitij Janardan Lakhani * Clean up code. Add new comments. Fix unecessary pasing of seg pos to the seqoffsets calculation API Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Optimize and fix the slow O(T*T) path for seqlens and seqoffsets calculation for THD non-cp and Cp p2p ring - Newer path is O(T*max_segments) per seq - Newer path works well with CP p2p ring Fix BRCM cross attn by routing to new slow path rather than fast causal path Signed-off-by: Kshitij Janardan Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix lint failure Signed-off-by: Kshitij Janardan Lakhani --------- Signed-off-by: Kshitij Lakhani Signed-off-by: Kshitij Janardan Lakhani Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Kshitij Janardan Lakhani Co-authored-by: JAX Toolbox --- transformer_engine/jax/attention.py | 273 +++++++++++++++++++++------- 1 file changed, 203 insertions(+), 70 deletions(-) diff --git a/transformer_engine/jax/attention.py b/transformer_engine/jax/attention.py index 29d084838..f54a043fd 100644 --- a/transformer_engine/jax/attention.py +++ b/transformer_engine/jax/attention.py @@ -11,7 +11,6 @@ from jax.ad_checkpoint import checkpoint_name import jax import jax.numpy as jnp -from flax.linen import make_attention_mask from transformer_engine_jax import NVTE_Bias_Type from transformer_engine_jax import NVTE_Mask_Type @@ -541,6 +540,149 @@ def run_length_fill(segment_ids) -> jnp.ndarray: return run_length_segment_id_shape.reshape(orig_shape) +def _get_seqlens_offsets_thd( + segment_ids_q, + segment_ids_kv, + segment_pos_q, + segment_pos_kv, + attn_mask_type, + max_segments_per_seq, +): + """O(T * max_segments_per_seq) replacement for the older O(T^2) mask-based slow path. + Returns (q_seqlen, kv_seqlen, q_offset, kv_offset) values to match the reference older mask-based path: + segment_mask = make_attention_mask(q_ids, kv_ids, equal) + segment_mask_with_id = make_attention_mask(q_ids, kv_ids, equal * q_id) + attn_mask = segment_mask AND (causal_or_brcm_or_none) + attn_mask_with_id = where(attn_mask, segment_mask_with_id, 0) + row_ids = reduce_max(attn_mask_with_id, axis=kv) # [B, T_q] + col_ids = reduce_max(attn_mask_with_id, axis=q) # [B, T_kv] + seqlens/offsets = bincount(...) / find_offsets(...) + The two reductions are expressed equivalently as per-segment aggregates: + - causal: row_ids[q] = q_seg_id iff seg_pos_q[q] >= min(seg_pos_kv over same-seg KV) + - brcm: row_ids[q] = q_seg_id iff (run_len_q - seg_pos_q) >= + min(run_len_kv - seg_pos_kv over same-seg KV) + - padding: row_ids[q] = q_seg_id iff q_seg_id appears in KV + (and symmetrically for col_ids with max/<=). + """ + + # Example: For striping P2P causal attention (but this logic also applies for non-CP fused attn) + # pre-striping and sharding: segment_ids = [[1 1 1 1 2 2 2 2]], segment_pos = [[0 1 2 3 0 1 2 3]] + # post-striping and sharding (striped CP=2, Q from rank 0 × KV from rank 1, max_segments_per_seq=2): + # segment_ids_q = [1 1 2 2] segment_pos_q = [0 2 0 2] → q_key = [0 2 0 2] + # segment_ids_kv = [1 1 2 2] segment_pos_kv = [1 3 1 3] → kv_key = [1 3 1 3] + # Q-side — kv_agg[s] = min(kv_key over same-seg KV), fill = max_fill_val = 5 (assumed to be large enough): + # scatter (rows = kv tokens, cols = segs): + # [5 1 5 / 5 3 5 / 5 5 1 / 5 5 3] → reduce min → kv_agg = [5 1 1] + # q_ok = q_key >= kv_agg[seg_ids_q] = [0 2 0 2] >= [1 1 1 1] = [F T F T] + # KV-side — q_agg[s] = max(q_key over same-seg Q), fill = neg_fill_val = -1 (assumed to be small enough): + # scatter: [-1 0 -1 / -1 2 -1 / -1 -1 0 / -1 -1 2] → reduce max → q_agg = [-1 2 2] + # kv_ok = kv_key <= q_agg[seg_ids_kv] = [1 3 1 3] <= [2 2 2 2] = [T F T F] + # Outer combiner: + # row_ids = [0 1 0 2] col_ids = [1 0 2 0] + # q_seqlen = [1 1] kv_seqlen = [1 1] + # q_offset = [1 3 -1] kv_offset = [0 2 -1] + def _row_and_col_ids(): + if attn_mask_type.is_bottom_right(): + # BRCM: mask[q][kv] = (same seg) AND (q_key <= kv_key). + rl_q = run_length_fill(segment_ids_q) + rl_kv = run_length_fill(segment_ids_kv) + q_key = (rl_q - segment_pos_q).astype(jnp.int32) + kv_key = (rl_kv - segment_pos_kv).astype(jnp.int32) + + # Use large positive and negative values as fill values for the KV keys and Q keys respectively + max_fill_val = jnp.asarray(jnp.iinfo(jnp.int32).max, dtype=jnp.int32) + neg_fill_val = jnp.asarray(-1, dtype=jnp.int32) + # Creates a one-hot encoding mask of the KV segment ids (size [B, T_kv, max_segments_per_seq+1]) + # i.e. each row has only one True value, which is the segment id of the row. + kv_oh = jax.nn.one_hot(segment_ids_kv, max_segments_per_seq + 1, dtype=jnp.bool_) + # Mask the KV keys with the valid segment ids (size [B, T_kv, 1]) + kv_key_masked = jnp.where(segment_ids_kv != 0, kv_key, neg_fill_val)[..., None] + # Scatter each KV key (i.e. seg pos) into it's own segment column + kv_agg = jnp.where(kv_oh, kv_key_masked, neg_fill_val) + kv_agg = jnp.max(kv_agg, axis=-2) + # Define causal relationship: Q is attended iff q_key <= max(kv_key over same-seg KV) + q_has_match = q_key <= jnp.take_along_axis( + kv_agg, segment_ids_q.astype(jnp.int32), axis=-1 + ) + + # Symmetric to the Q case, but with KV and Q swapped + q_oh = jax.nn.one_hot(segment_ids_q, max_segments_per_seq + 1, dtype=jnp.bool_) + q_key_masked = jnp.where(segment_ids_q != 0, q_key, max_fill_val)[..., None] + q_agg = jnp.where(q_oh, q_key_masked, max_fill_val) + q_agg = jnp.min(q_agg, axis=-2) + # Define causal relationship: KV is attended iff kv_key >= min(q_key over same-seg Q) + kv_has_match = kv_key >= jnp.take_along_axis( + q_agg, segment_ids_kv.astype(jnp.int32), axis=-1 + ) + elif attn_mask_type.is_causal(): + # CM: mask[q][kv] = (same_seg) AND (q_pos >= kv_pos). + q_key = segment_pos_q.astype(jnp.int32) + kv_key = segment_pos_kv.astype(jnp.int32) + + # Use large positive and negative values as a fill value for the KV keys and Q keys respectively + max_fill_val = jnp.asarray(jnp.iinfo(jnp.int32).max, dtype=jnp.int32) + neg_fill_val = jnp.asarray(-1, dtype=jnp.int32) + + # Creates a one-hot encoding mask of the KV segment ids (size [B, T_kv, max_segments_per_seq+1]) + # i.e. each row has only one True value, which is the segment id of the row. + kv_oh = jax.nn.one_hot(segment_ids_kv, max_segments_per_seq + 1, dtype=jnp.bool_) + # Mask the KV keys with the valid segment ids (size [B, T_kv, 1]) + kv_key_masked = jnp.where(segment_ids_kv != 0, kv_key, max_fill_val)[..., None] + # Scatter each KV key (i.e. seg pos) into it's own segment column + kv_agg = jnp.where(kv_oh, kv_key_masked, max_fill_val) + kv_agg = jnp.min(kv_agg, axis=-2) + # Define causal relationship: Q is attended iff q_key >= min(kv_key over same-seg KV) + q_has_match = q_key >= jnp.take_along_axis( + kv_agg, segment_ids_q.astype(jnp.int32), axis=-1 + ) + + # Symmetric to the Q case, but with KV and Q swapped + q_oh = jax.nn.one_hot(segment_ids_q, max_segments_per_seq + 1, dtype=jnp.bool_) + q_key_masked = jnp.where(segment_ids_q != 0, q_key, neg_fill_val)[..., None] + q_agg = jnp.where(q_oh, q_key_masked, neg_fill_val) + q_agg = jnp.max(q_agg, axis=-2) + # Define causal relationship: KV is attended iff kv_key <= max(q_key over same-seg Q) + kv_has_match = kv_key <= jnp.take_along_axis( + q_agg, segment_ids_kv.astype(jnp.int32), axis=-1 + ) + else: + # Padding-only: row_ids[q] = q_seg_id iff q_seg_id is present in KV (and q not pad). + kv_seg_ids_present = jax.nn.one_hot( + segment_ids_kv, max_segments_per_seq + 1, dtype=jnp.bool_ + ).any(axis=-2) + q_seg_ids_present = jax.nn.one_hot( + segment_ids_q, max_segments_per_seq + 1, dtype=jnp.bool_ + ).any(axis=-2) + q_has_match = jnp.take_along_axis( + kv_seg_ids_present, segment_ids_q.astype(jnp.int32), axis=-1 + ) & (segment_ids_q != 0) + kv_has_match = jnp.take_along_axis( + q_seg_ids_present, segment_ids_kv.astype(jnp.int32), axis=-1 + ) & (segment_ids_kv != 0) + + row_ids = jnp.where(q_has_match, segment_ids_q, 0).astype(jnp.int32) + col_ids = jnp.where(kv_has_match, segment_ids_kv, 0).astype(jnp.int32) + return row_ids, col_ids + + row_ids, col_ids = _row_and_col_ids() + + bincount_vmap = jax.vmap(partial(jnp.bincount, length=max_segments_per_seq + 1)) + q_seqlen = bincount_vmap(row_ids)[..., 1:] + kv_seqlen = bincount_vmap(col_ids)[..., 1:] + + def _find_offsets(x): + same_as_previous = jnp.logical_and(x[..., 1:] != x[..., :-1], x[..., 1:] != 0) + first_column = x[..., :1] != 0 + boundaries = jnp.concatenate([first_column, same_as_previous], axis=-1) + return jax.vmap(partial(jnp.argwhere, size=(max_segments_per_seq + 1), fill_value=-1))( + boundaries + ).squeeze(-1) + + q_offset = _find_offsets(row_ids) + kv_offset = _find_offsets(col_ids) + return q_seqlen, kv_seqlen, q_offset, kv_offset + + def _segment_ids_pos_to_seqlens_offsets( segment_ids_q, segment_ids_kv, @@ -550,9 +692,52 @@ def _segment_ids_pos_to_seqlens_offsets( window_size, max_segments_per_seq, ): + """Compute per-segment seqlens and start offsets(currently only used for THD) + Given segment-id and segment-position tensors for Q and KV, + returns the four metadata tensors cuDNN needed for variable-length attention: + q_seqlen : [..., max_segments_per_seq] # valid Q tokens per segment + kv_seqlen : [..., max_segments_per_seq] # valid KV tokens per segment + q_offset : [..., max_segments_per_seq + 1] # start index of each Q segment + kv_offset : [..., max_segments_per_seq + 1] # start index of each KV segment + + Args: + segment_ids_q: int32 [..., T_q] per-token segment id; 0 == padding + segment_ids_kv: int32 [..., T_kv] same convention as segment_ids_q + segment_pos_q: int32 [..., T_q] per-token position inside its segment + segment_pos_kv: int32 [..., T_kv] same convention as segment_pos_q + attn_mask_type: AttnMaskType. Selects the mask predicate used to decide + which positions are valid (top-left causal vs + bottom-right causal vs. padding-only) + window_size: Optional sliding-window tuple ``(left, right)`` or None + Used here only as a fast-path eligibility hint + max_segments_per_seq: maximum number of segments expected per row + Used to size the bincount / argwhere outputs + + Routing (only invoked for THD qkv_layout): + 1. Fast path -- ``_segment_ids_pos_to_seqlens_offsets_fast_causal_path``. + O(T) per row. Counts all segment tokens via bincount on + segment_ids and trims at most one token per segment at the + boundary. Used for: + - top-left CAUSAL / PADDING_CAUSAL with ``window_size is None`` + - SWA with ``window_size == (-1, -1)`` and not bottom-right + Bottom-right causal cross-attention is excluded: the boundary + trim leaves kv_seqlen short by one per active segment, which + shifts the BRCM bottom-right alignment by one KV per Q row. + + 2. Slow path -- ``_get_seqlens_offsets_thd``. + O(T * max_segments_per_seq) per row. Per-segment min/max + aggregation that is equivalent to the older O(T^2) + mask-based reference for top-left causal, bottom-right causal, + and padding-only masks. Required under ring attention where + ``segment_ids_q != segment_ids_kv`` in rotated steps. + + Returns: + Tuple ``(q_seqlen, kv_seqlen, q_offset, kv_offset)`` with shapes as + above. Inactive segment slots are filled with 0 in seqlens and -1 + in offsets. + """ # TODO(mgoldfarb-nvidia): Consider an opt-in for arbitrary masking if needed here. # Computing the full mask is expensive due to quadratic expansion of Q * KV masking. - # Assumptions for cudnn causal mask correctness. # 1. Segments are monotonic [4 4 4 0 0 5 5 5 6 6 0 0] # 2. No intra-segment padding, only inter-segment paddding allowed @@ -561,82 +746,30 @@ def _segment_ids_pos_to_seqlens_offsets( # 0 x x # 4 x x x x x # 8 x x x x x x x x - # # This fast path avoids expanding the mask to Q * KV matrix and instead allows us to # examine only O(Q+KV) elements. - - # For seqlens and seqoffsets calculations, the intermediate(temp) attn_mask creation - # using the segment ids and pos along with mask type (causal or brcm) is sufficient. - # It does not need to involve SW for this mask's creation - - # Currently, this function is only exercised for THD qkv_layout. - - # TODO(KshitijLakhani): Try exercising the fast path for BRCM as well - if (attn_mask_type.is_causal() and window_size is None) or ( - window_size == (-1, -1) and not attn_mask_type.is_bottom_right() - ): + # The fast causal path encodes TOP-LEFT causal semantics via + # valid[q][kv] = (segment_pos_q >= segment_pos_kv) + # which is only equivalent to BRCM when s_q == s_kv (self-attention). For + # cross-attention (s_q != s_kv), BRCM diverges from top-left causal, so we + # must route bottom-right masks to the slow path. + + # Fast path: O(T) per row. + if ( + attn_mask_type.is_causal() and not attn_mask_type.is_bottom_right() and window_size is None + ) or (window_size == (-1, -1) and not attn_mask_type.is_bottom_right()): return _segment_ids_pos_to_seqlens_offsets_fast_causal_path( segment_ids_q, segment_ids_kv, segment_pos_q, segment_pos_kv, max_segments_per_seq ) - - # (1 = attend, 0 = masked) - segment_mask = make_attention_mask( - segment_ids_q, - segment_ids_kv, - jnp.equal, - ) - segment_mask_with_id = make_attention_mask( + # Slow path: O(T * max_segments_per_seq) per row. + return _get_seqlens_offsets_thd( segment_ids_q, segment_ids_kv, - lambda x, y: jnp.equal(x, y) * x, - ) - # TE JAX Attn expects the THD segments to have q_token <= kv_tokens so that a correct cross-attn type BRCM can be applied - attn_mask = segment_mask - if attn_mask_type.is_bottom_right(): - run_length_out_q = run_length_fill(segment_ids_q) - run_length_out_kv = run_length_fill(segment_ids_kv) - # Example for brcm: - # run_length_out_q: [3 3 3 0 4 4 4 4] - # segment_pos_q: [0 1 2 3 0 1 2 3] - # segment_ids_q: [1 1 1 0 2 2 2 2] - # run_length_out_kv: [4 4 4 4 0 0 10 10 10 10 10 10 10 10 10 10] - # segment_pos_kv: [0 1 2 3 4 5 0 1 2 3 4 5 6 7 8 9] - # segment_ids_kv: [1 1 1 1 0 0 2 2 2 2 2 2 2 2 2 2] - # brcm: [[[1 1 0 0 0 0 1 1 1 1 1 1 1 1 0 0] - # [1 1 1 0 0 0 1 1 1 1 1 1 1 1 1 0] - # [1 1 1 1 0 0 1 1 1 1 1 1 1 1 1 1] - # [1 1 1 1 0 0 1 1 1 1 1 1 1 1 1 1] - # [1 0 0 0 0 0 1 1 1 1 1 1 1 0 0 0] - # [1 1 0 0 0 0 1 1 1 1 1 1 1 1 0 0] - # [1 1 1 0 0 0 1 1 1 1 1 1 1 1 1 0] - # [1 1 1 1 0 0 1 1 1 1 1 1 1 1 1 1]]] - # attn_mask(noswa):[[[1 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0] - # [1 1 1 0 0 0 0 0 0 0 0 0 0 0 0 0] - # [1 1 1 1 0 0 0 0 0 0 0 0 0 0 0 0] - # [0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0] - # [0 0 0 0 0 0 1 1 1 1 1 1 1 0 0 0] - # [0 0 0 0 0 0 1 1 1 1 1 1 1 1 0 0] - # [0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 0] - # [0 0 0 0 0 0 1 1 1 1 1 1 1 1 1 1]]] - bottom_right_causal_mask = make_attention_mask( - run_length_out_q - segment_pos_q, - run_length_out_kv - segment_pos_kv, - jnp.less_equal, - ) - attn_mask = jnp.logical_and(segment_mask, bottom_right_causal_mask) - elif attn_mask_type.is_causal(): - causal_mask = make_attention_mask( - segment_pos_q, - segment_pos_kv, - jnp.greater_equal, - ) - attn_mask = jnp.logical_and(segment_mask, causal_mask) - - attn_mask_with_id = jnp.where(attn_mask, segment_mask_with_id, 0) - q_seqlen, q_offset, kv_seqlen, kv_offset = _mask_to_seqlens_offset( - attn_mask_with_id, max_segments_per_seq + segment_pos_q, + segment_pos_kv, + attn_mask_type, + max_segments_per_seq, ) - return q_seqlen, kv_seqlen, q_offset, kv_offset def _segment_ids_to_seqlens(segment_ids_q, segment_ids_kv, attn_mask_type): From 0803102751cfb099f71348efddc50cc0664cda89 Mon Sep 17 00:00:00 2001 From: Yigong Qin <62076142+YigongQin@users.noreply.github.com> Date: Sat, 2 May 2026 02:41:25 -0700 Subject: [PATCH 032/180] Optimizations for MXFP8/NVFP4 dequantize kernels (#2865) * Handle empty tensors in dequantize for CUDA graph compatibility Signed-off-by: YigongQin * dequant with swizzled scales Signed-off-by: YigongQin * pass nvfp4 dequant tests Signed-off-by: YigongQin * cleanup unit tests Signed-off-by: YigongQin * remove allocation in set amax Signed-off-by: YigongQin * Drop disabling `optimize_for_gemm` introduced in PR 2644 Signed-off-by: Ziang Li Signed-off-by: YigongQin * Drop `optimize_for_gemm` in basic linear Signed-off-by: Ziang Li Signed-off-by: YigongQin * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Apply suggestions from code review Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * remove redundant set scale Signed-off-by: YigongQin * rebase on nvfp4 test fix Signed-off-by: YigongQin * remove redundant line Signed-off-by: YigongQin * add missing from_cpu() for scale Signed-off-by: YigongQin * Remove unnecessary scale from NVFP4 C++ tests Signed-off-by: Tim Moon --------- Signed-off-by: YigongQin Signed-off-by: Ziang Li Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon Co-authored-by: Ziang Li Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Tim Moon --- tests/cpp/operator/CMakeLists.txt | 1 + tests/cpp/operator/test_dequantize_mxfp8.cu | 155 +++++++++ tests/cpp/operator/test_dequantize_nvfp4.cu | 297 ++++++++++++++++++ .../common/cast/dispatch/dequantize.cuh | 4 + .../common/cast/mxfp8/dequantize_mxfp8.cuh | 60 ++-- .../common/cast/nvfp4/dequantize_nvfp4.cuh | 31 +- .../pytorch/module/grouped_linear.py | 31 +- .../pytorch/module/layernorm_linear.py | 7 - transformer_engine/pytorch/module/linear.py | 7 - .../pytorch/ops/basic/basic_linear.py | 9 - 10 files changed, 523 insertions(+), 79 deletions(-) create mode 100644 tests/cpp/operator/test_dequantize_nvfp4.cu diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt index a5ea74171..9b67c09f3 100644 --- a/tests/cpp/operator/CMakeLists.txt +++ b/tests/cpp/operator/CMakeLists.txt @@ -16,6 +16,7 @@ add_executable(test_operator test_cast_float8blockwise.cu test_dequantize_mxfp8.cu test_dequantize_mxfp8_grouped.cu + test_dequantize_nvfp4.cu test_transpose.cu test_cast_transpose.cu test_cast_transpose_current_scaling.cu diff --git a/tests/cpp/operator/test_dequantize_mxfp8.cu b/tests/cpp/operator/test_dequantize_mxfp8.cu index a529f93d7..5950fba98 100644 --- a/tests/cpp/operator/test_dequantize_mxfp8.cu +++ b/tests/cpp/operator/test_dequantize_mxfp8.cu @@ -18,6 +18,7 @@ #include #include +#include #include "../test_common.h" #include "transformer_engine/transformer_engine.h" @@ -369,7 +370,95 @@ void performTest_x2(const size_t rows, compareResults("output_colwise", output, ref_output_colwise.get(), false, atol, rtol); } +// Dequantize with GEMM-swizzled scales (single dimension) +template +void performTest_x1_swizzled(const size_t rows, + const size_t cols, + const bool rowwise, + const bool colwise) +{ + using namespace test; + DType itype = TypeInfo::dtype; + DType otype = TypeInfo::dtype; + + const size_t block_size_rows = rowwise ? 1 : 32; + const size_t block_size_cols = colwise ? 1 : 32; + + const size_t unpadded_blocks_Y_rowwise = rows; + const size_t unpadded_blocks_X_rowwise = divide_round_up(cols, block_size_cols); + const size_t unpadded_blocks_Y_colwise = divide_round_up(rows, block_size_rows); + const size_t unpadded_blocks_X_colwise = cols; + + const size_t blocks_Y_rowwise = round_up_to_nearest_multiple(unpadded_blocks_Y_rowwise, + scale_tensor_alignment_Y_rowwise); + const size_t blocks_X_rowwise = round_up_to_nearest_multiple(unpadded_blocks_X_rowwise, + scale_tensor_alignment_X_rowwise); + const size_t blocks_Y_colwise = round_up_to_nearest_multiple(unpadded_blocks_Y_colwise, + scale_tensor_alignment_Y_colwise); + const size_t blocks_X_colwise = round_up_to_nearest_multiple(unpadded_blocks_X_colwise, + scale_tensor_alignment_X_colwise); + + const size_t blocks_num_rowwise = blocks_Y_rowwise * blocks_X_rowwise; + const size_t blocks_num_colwise = blocks_Y_colwise * blocks_X_colwise; + + const size_t blocks_num = rowwise ? blocks_num_rowwise : blocks_num_colwise; + const size_t scales_stride = rowwise ? blocks_X_rowwise : blocks_X_colwise; + + Tensor input_compact_scales("input_compact_scales", std::vector{ rows, cols }, itype, + rowwise, colwise, NVTE_MXFP8_1D_SCALING); + + Tensor input_swizzled_scales("input_swizzled_scales", std::vector{ rows, cols }, itype, + rowwise, colwise, NVTE_MXFP8_1D_SCALING); + input_swizzled_scales.set_with_gemm_swizzled_scales(true); + + Tensor output("output", std::vector{ rows, cols }, otype, true, false); + + std::unique_ptr ref_output = std::make_unique(rows * cols); + std::unique_ptr scales = std::make_unique(blocks_num); + + fill_tensor_data(input_compact_scales, scales.get(), scales.get(), rowwise, colwise, + rows, cols, blocks_num_rowwise, blocks_num_colwise); + + const size_t data_bytes = rows * cols * sizeof(InputType); + if (rowwise && data_bytes > 0) { + cudaMemcpy(input_swizzled_scales.rowwise_dptr(), input_compact_scales.rowwise_dptr(), + data_bytes, cudaMemcpyDeviceToDevice); + } + if (colwise && data_bytes > 0) { + cudaMemcpy(input_swizzled_scales.columnwise_dptr(), input_compact_scales.columnwise_dptr(), + data_bytes, cudaMemcpyDeviceToDevice); + } + + if (data_bytes > 0) { + nvte_swizzle_scaling_factors(input_compact_scales.data(), input_swizzled_scales.data(), 0); + } + + nvte_dequantize(input_swizzled_scales.data(), output.data(), 0); + + cudaDeviceSynchronize(); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + InputType *data_ptr = rowwise + ? input_compact_scales.rowwise_cpu_dptr() + : input_compact_scales.columnwise_cpu_dptr(); + + compute_ref_x1(data_ptr, + ref_output.get(), + scales.get(), + rows, + cols, + block_size_rows, + block_size_cols, + scales_stride); + + auto [atol, rtol] = getTolerances(otype); + compareResults("output_swizzled", output, ref_output.get(), true, atol, rtol); +} + std::vector> tensor_dims = { + {0, 128}, + {0, 256}, {1, 16}, {16, 48}, {65, 96}, @@ -470,3 +559,69 @@ INSTANTIATE_TEST_SUITE_P( return name; } ); + +/***************************************************************************** + * Swizzled-scale dequantization tests + *****************************************************************************/ + +class DequantizeMXFP8SwizzledTestSuite : public ::testing::TestWithParam + , + std::pair, + transformer_engine::DType, + transformer_engine::DType>> {}; + +TEST_P(DequantizeMXFP8SwizzledTestSuite, TestDequantizeMXFP8Swizzled) +{ + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + using namespace transformer_engine; + using namespace test; + + const auto tensor_size = std::get<0>(GetParam()); + const auto block_size = std::get<1>(GetParam()); + const DType input_type = std::get<2>(GetParam()); + const DType output_type = std::get<3>(GetParam()); + + const bool rowwise = block_size.second != 1; + const bool colwise = block_size.first != 1; + + if (rowwise && colwise) { + GTEST_SKIP(); + } + + if (rowwise && tensor_size.second % 32 != 0) { + GTEST_SKIP(); + } + if (colwise && tensor_size.first % 32 != 0) { + GTEST_SKIP(); + } + + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY(input_type, InputType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(output_type, OutputType, + performTest_x1_swizzled( + tensor_size.first, tensor_size.second, rowwise, colwise); + ); + ); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + DequantizeMXFP8SwizzledTestSuite, + ::testing::Combine( + ::testing::ValuesIn(tensor_dims), + ::testing::ValuesIn(block_sizes), + ::testing::Values(DType::kFloat8E4M3, DType::kFloat8E5M2), + ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16)), + [](const testing::TestParamInfo& info) + { + std::string name = std::to_string(std::get<0>(info.param).first) + "X" + + std::to_string(std::get<0>(info.param).second) + "X" + + std::to_string(std::get<1>(info.param).first) + "X" + + std::to_string(std::get<1>(info.param).second) + "X" + + test::typeName(std::get<2>(info.param)) + "X" + + test::typeName(std::get<3>(info.param)); + return name; + } +); diff --git a/tests/cpp/operator/test_dequantize_nvfp4.cu b/tests/cpp/operator/test_dequantize_nvfp4.cu new file mode 100644 index 000000000..96e85cb5e --- /dev/null +++ b/tests/cpp/operator/test_dequantize_nvfp4.cu @@ -0,0 +1,297 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#if FP4_TYPE_SUPPORTED +#include +#endif + +#include +#include +#include "../test_common.h" +#include "transformer_engine/transformer_engine.h" + +using namespace transformer_engine; +using namespace test; + +#if FP4_TYPE_SUPPORTED + +namespace { + +float2 cvt_fp4x2_to_float2(fp4e2m1x2 fp4_pair) { + const __half2_raw raw = + __nv_cvt_fp4x2_to_halfraw2( + *reinterpret_cast<__nv_fp4x2_storage_t *>(&fp4_pair), __NV_E2M1); + const __half2 h2(raw); + return {static_cast(h2.x), static_cast(h2.y)}; +} + +template +void compute_ref_dequantize_nvfp4(const uint8_t *packed_data, + const fp8e4m3 *scales, + float amax, + OType *output, + size_t rows, + size_t cols, + size_t scale_stride) { + constexpr float factor_inv = 1.0f / (6.0f * 448.0f); + constexpr size_t BLOCK_SIZE = 16; + const size_t Mread = cols / BLOCK_SIZE; + const size_t bytes_per_block = BLOCK_SIZE / 2; + + for (size_t row = 0; row < rows; ++row) { + for (size_t block = 0; block < Mread; ++block) { + const fp8e4m3 scale = scales[row * scale_stride + block]; + const float final_scale = static_cast(scale) * amax * factor_inv; + + for (size_t pair_idx = 0; pair_idx < bytes_per_block; ++pair_idx) { + const size_t byte_idx = + (row * Mread + block) * bytes_per_block + pair_idx; + fp4e2m1x2 fp4_pair; + std::memcpy(&fp4_pair, &packed_data[byte_idx], 1); + const float2 values = cvt_fp4x2_to_float2(fp4_pair); + + const size_t col0 = block * BLOCK_SIZE + pair_idx * 2; + output[row * cols + col0] = + static_cast(values.x * final_scale); + output[row * cols + col0 + 1] = + static_cast(values.y * final_scale); + } + } + } +} + +template +float compute_amax(const test::Tensor &t, size_t rows, size_t cols) { + t.to_cpu(); + const auto *data = t.rowwise_cpu_dptr(); + float amax = 0.0f; + for (size_t i = 0; i < rows * cols; ++i) { + amax = std::max(amax, std::abs(static_cast(data[i]))); + } + return amax; +} + +// Quantize a high-precision input to NVFP4, then dequantize and compare +// against a CPU reference computed from the quantized data. +template +void performTest_dequantize_nvfp4(const size_t rows, const size_t cols) { + using namespace test; + DType otype = TypeInfo::dtype; + + Tensor input("input", std::vector{rows, cols}, otype); + fillCase(&input, InputsFillCase::uniform); + + Tensor quantized("quantized", std::vector{rows, cols}, + DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING); + if (rows > 0 && cols > 0) { + quantized.set_tensor_amax(compute_amax(input, rows, cols)); + } else { + quantized.set_tensor_amax(0.0f); + } + + if (rows > 0 && cols > 0) { + nvte_quantize(input.data(), quantized.data(), 0); + cudaDeviceSynchronize(); + } + + Tensor output("output", std::vector{rows, cols}, otype, true, false); + nvte_dequantize(quantized.data(), output.data(), 0); + cudaDeviceSynchronize(); + + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + if (rows > 0 && cols > 0) { + quantized.to_cpu(); + const uint8_t *fp4_data = + reinterpret_cast(quantized.rowwise_cpu_dptr()); + const fp8e4m3 *scales = quantized.rowwise_cpu_scale_inv_ptr(); + const float amax_val = quantized.amax(); + const NVTEShape scale_shape = quantized.rowwise_scale_inv_shape(); + const size_t scale_stride = scale_shape.data[scale_shape.ndim - 1]; + + std::unique_ptr ref_output = + std::make_unique(rows * cols); + compute_ref_dequantize_nvfp4( + fp4_data, scales, amax_val, ref_output.get(), + rows, cols, scale_stride); + + auto [atol, rtol] = getTolerances(otype); + compareResults("output_nvfp4", output, ref_output.get(), true, atol, rtol); + } +} + +// Dequantize NVFP4 with GEMM-swizzled scales and compare against compact path. +template +void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols) { + using namespace test; + DType otype = TypeInfo::dtype; + + Tensor input("input", std::vector{rows, cols}, otype); + fillCase(&input, InputsFillCase::uniform); + + Tensor quantized_compact("quantized_compact", std::vector{rows, cols}, + DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING); + if (rows > 0 && cols > 0) { + quantized_compact.set_tensor_amax(compute_amax(input, rows, cols)); + } else { + quantized_compact.set_tensor_amax(0.0f); + } + + if (rows > 0 && cols > 0) { + nvte_quantize(input.data(), quantized_compact.data(), 0); + cudaDeviceSynchronize(); + } + + // Dequantize with compact scales → reference output + Tensor output_compact("output_compact", std::vector{rows, cols}, otype, true, false); + nvte_dequantize(quantized_compact.data(), output_compact.data(), 0); + cudaDeviceSynchronize(); + + // Create tensor with same FP4 data but swizzled scales + Tensor quantized_swizzled("quantized_swizzled", std::vector{rows, cols}, + DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING); + quantized_swizzled.set_tensor_amax(0.0f); + quantized_swizzled.set_with_gemm_swizzled_scales(true); + + // Copy amax and scale from compact to swizzled before FP4 data, + // since from_cpu() uploads all CPU buffers (including zero-init data). + quantized_compact.to_cpu(); + quantized_swizzled.set_tensor_amax(quantized_compact.amax()); + + // Copy FP4 data after from_cpu() to avoid being overwritten + const size_t data_bytes = rows * cols / 2; + if (data_bytes > 0) { + cudaMemcpy(quantized_swizzled.rowwise_dptr(), quantized_compact.rowwise_dptr(), + data_bytes, cudaMemcpyDeviceToDevice); + } + + // Swizzle scales + if (data_bytes > 0) { + nvte_swizzle_scaling_factors(quantized_compact.data(), quantized_swizzled.data(), 0); + } + + // Dequantize with swizzled scales + Tensor output_swizzled("output_swizzled", std::vector{rows, cols}, otype, true, false); + nvte_dequantize(quantized_swizzled.data(), output_swizzled.data(), 0); + cudaDeviceSynchronize(); + + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); + + // Read compact output as reference + const size_t num_elems = rows * cols; + std::unique_ptr ref_output = std::make_unique(num_elems); + if (num_elems > 0) { + cudaMemcpy(ref_output.get(), output_compact.rowwise_dptr(), + num_elems * sizeof(OutputType), cudaMemcpyDeviceToHost); + } + + auto [atol, rtol] = getTolerances(otype); + if (num_elems > 0) { + compareResults("output_nvfp4_swizzled", output_swizzled, + ref_output.get(), true, atol, rtol); + } +} + +std::vector> nvfp4_tensor_dims = { + {0, 128}, + {0, 256}, + {32, 32}, + {32, 64}, + {64, 96}, + {128, 128}, + {128, 256}, + {256, 256}, + {256, 512}, + {512, 1024}, + {992, 512}, + {768, 1024}, +}; + +} // namespace + +class DequantizeNVFP4TestSuite : public ::testing::TestWithParam + , + transformer_engine::DType>> {}; + +TEST_P(DequantizeNVFP4TestSuite, TestDequantizeNVFP4) +{ + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + const auto tensor_size = std::get<0>(GetParam()); + const DType output_type = std::get<1>(GetParam()); + + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(output_type, OutputType, + performTest_dequantize_nvfp4( + tensor_size.first, tensor_size.second); + ); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + DequantizeNVFP4TestSuite, + ::testing::Combine( + ::testing::ValuesIn(nvfp4_tensor_dims), + ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16)), + [](const testing::TestParamInfo& info) + { + std::string name = std::to_string(std::get<0>(info.param).first) + "X" + + std::to_string(std::get<0>(info.param).second) + "X" + + test::typeName(std::get<1>(info.param)); + return name; + } +); + +class DequantizeNVFP4SwizzledTestSuite : public ::testing::TestWithParam + , + transformer_engine::DType>> {}; + +TEST_P(DequantizeNVFP4SwizzledTestSuite, TestDequantizeNVFP4Swizzled) +{ + if (getDeviceComputeCapability() < blackwellComputeCapability) { + GTEST_SKIP(); + } + + const auto tensor_size = std::get<0>(GetParam()); + const DType output_type = std::get<1>(GetParam()); + + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(output_type, OutputType, + performTest_dequantize_nvfp4_swizzled( + tensor_size.first, tensor_size.second); + ); +} + +INSTANTIATE_TEST_SUITE_P( + OperatorTest, + DequantizeNVFP4SwizzledTestSuite, + ::testing::Combine( + ::testing::ValuesIn(nvfp4_tensor_dims), + ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16)), + [](const testing::TestParamInfo& info) + { + std::string name = std::to_string(std::get<0>(info.param).first) + "X" + + std::to_string(std::get<0>(info.param).second) + "X" + + test::typeName(std::get<1>(info.param)) + "X" + + "Swizzled"; + return name; + } +); + +#endif // FP4_TYPE_SUPPORTED diff --git a/transformer_engine/common/cast/dispatch/dequantize.cuh b/transformer_engine/common/cast/dispatch/dequantize.cuh index 12787d609..63c1b046f 100644 --- a/transformer_engine/common/cast/dispatch/dequantize.cuh +++ b/transformer_engine/common/cast/dispatch/dequantize.cuh @@ -26,6 +26,10 @@ inline void dequantize_helper(const Tensor &input, Tensor *output, cudaStream_t CheckInputTensor(input, "cast_input"); CheckOutputTensor(*output, "cast_output"); + if (input.numel() == 0) { + return; + } + switch (input.scaling_mode) { case NVTE_DELAYED_TENSOR_SCALING: { NVTE_CHECK(is_fp8_dtype(input.dtype()), "Input must have FP8 type."); diff --git a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh index f8fecaa4e..6441a567a 100644 --- a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh @@ -20,6 +20,7 @@ #include "../../util/math.h" #include "../../util/ptx.cuh" #include "../../utils.cuh" +#include "swizzle.cuh" namespace transformer_engine { namespace dispatch { @@ -42,12 +43,13 @@ constexpr size_t THREADS_PER_CHUNK_X_COLWISE = CHUNK_DIM_X; constexpr size_t ITERATIONS = CHUNK_DIM_Y / BUFFER_DIM_Y; // 8 = 128 / 16 static_assert(ITERATIONS >= 1); -template +template __global__ void __launch_bounds__(THREADS_PER_CHUNK) dequantize_mxfp8_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, const e8m0_t *const scales_ptr, const size_t rows, const size_t cols, - const size_t scales_stride) { + const size_t scales_stride, const size_t num_scale_tiles_X) { #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) constexpr bool USE_ROWWISE_SCALING = SCALE_DIM_X > 1; @@ -158,7 +160,18 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) ? (scales_rowwise_chunk_offset_X + tid_rowwise_X / THREADS_PER_SCALE_X_ROWWISE) : (scales_colwise_chunk_offset_X + tid_colwise_X); - const int scale_idx = scale_offset_Y * scales_stride + scale_offset_X; + size_t scale_idx; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + if constexpr (USE_ROWWISE_SCALING) { + scale_idx = + swizzle::gemm_swizzled_scale_idx(scale_offset_Y, scale_offset_X, num_scale_tiles_X); + } else { + scale_idx = + swizzle::gemm_swizzled_scale_idx(scale_offset_X, scale_offset_Y, num_scale_tiles_X); + } + } else { + scale_idx = scale_offset_Y * scales_stride + scale_offset_X; + } const e8m0_t biased_exponent = scales_ptr[scale_idx]; const float block_scale = ptx::exp2f(biased_exponent); @@ -239,10 +252,11 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) NVTE_CHECK(is_fp8_dtype(input.columnwise_data.dtype), "Input must have FP8 type."); } - NVTE_CHECK(!input.with_gemm_swizzled_scales, "Input must have scales in compact format."); NVTE_CHECK(!is_fp8_dtype(output->data.dtype), "Output must be in higher precision."); NVTE_CHECK(output->shape() == input.shape(), "Input and output shapes need to match."); + const bool with_gemm_swizzled_scales = input.with_gemm_swizzled_scales; + // TODO: Make more general const size_t scale_dim_X_rowwise = use_rowwise_scaling ? 32 : 1; const size_t scale_dim_Y_colwise = use_colwise_scaling ? 32 : 1; @@ -276,6 +290,9 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) const size_t scales_stride = use_rowwise_scaling ? scales_X_rowwise : scales_X_colwise; + const size_t num_scale_tiles_X = use_rowwise_scaling ? DIVUP(cols, static_cast(128)) + : DIVUP(rows, static_cast(128)); + const SimpleTensor &input_data = use_rowwise_scaling ? input.data : input.columnwise_data; const dim3 block(THREADS_PER_CHUNK); @@ -289,21 +306,26 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) input.dtype(), IType, TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( output->dtype(), OType, - - alignas(64) CUtensorMap tensor_map_input{}; - alignas(64) CUtensorMap tensor_map_output{}; - - create_2D_tensor_map(tensor_map_input, input_data, rows, cols, SHMEM_DIM_Y, - SHMEM_DIM_X, cols, 0, typeToNumBits(input.dtype())); - create_2D_tensor_map(tensor_map_output, output->data, rows, cols, SHMEM_DIM_Y, - SHMEM_DIM_X, cols, 0, typeToNumBits(output->dtype())); - - dequantize_mxfp8_kernel - <<>>(tensor_map_input, tensor_map_output, scales_ptr, - rows, cols, scales_stride);); // NOLINT(*) - ); // NOLINT(*) - ); // NOLINT(*) - ); // NOLINT(*) + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, + + alignas(64) CUtensorMap tensor_map_input{}; + alignas(64) CUtensorMap tensor_map_output{}; + + create_2D_tensor_map(tensor_map_input, input_data, rows, cols, SHMEM_DIM_Y, + SHMEM_DIM_X, cols, 0, typeToNumBits(input.dtype())); + create_2D_tensor_map(tensor_map_output, output->data, rows, cols, SHMEM_DIM_Y, + SHMEM_DIM_X, cols, 0, typeToNumBits(output->dtype())); + + dequantize_mxfp8_kernel + <<>>(tensor_map_input, tensor_map_output, scales_ptr, + rows, cols, scales_stride, + num_scale_tiles_X);); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) + ); // NOLINT(*) NVTE_CHECK_CUDA(cudaGetLastError()); } } // namespace mxfp8 diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index ccdc4c93e..414320815 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -20,6 +20,7 @@ #include "../../util/math.h" #include "../../util/ptx.cuh" #include "../../utils.cuh" +#include "../mxfp8/swizzle.cuh" #if FP4_TYPE_SUPPORTED #include @@ -30,11 +31,11 @@ namespace dispatch { namespace nvfp4 { namespace dequantize_kernel { #if FP4_TYPE_SUPPORTED -template +template __global__ void __launch_bounds__(512) dequantize_fp4_kernel(const void *const input, OType *output, const fp8e4m3 *const scales, const float *const tensor_amax, const size_t N, const size_t M, - const size_t scale_stride) { + const size_t scale_stride, const size_t num_scale_tiles_X) { const size_t thread_idx = blockIdx.x * blockDim.x + threadIdx.x; const size_t x = thread_idx % M; const size_t y = thread_idx / M; @@ -52,7 +53,12 @@ __global__ void __launch_bounds__(512) OVec *output_vec = reinterpret_cast(output); const size_t my_index = x + y * M; - const size_t my_scale_index = x + y * scale_stride; + size_t my_scale_index; + if constexpr (WITH_GEMM_SWIZZLED_SCALES) { + my_scale_index = mxfp8::swizzle::gemm_swizzled_scale_idx(y, x, num_scale_tiles_X); + } else { + my_scale_index = x + y * scale_stride; + } const size_t my_output_index = (x + y * M) * 4; fp4vec value; value.vec = input_vectorized[my_index]; @@ -80,10 +86,11 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) CheckInputTensor(input, "input"); CheckOutputTensor(*output, "output"); NVTE_CHECK(input.data.dtype == DType::kFloat4E2M1, "Input must have FP4 type."); - NVTE_CHECK(!input.with_gemm_swizzled_scales, "Input must have scales in compact format."); NVTE_CHECK(is_high_precision_dtype(output->data.dtype), "Output must be in higher precision."); NVTE_CHECK(output->data.shape == input.data.shape, "Input and output shapes need to match."); + const bool with_gemm_swizzled_scales = input.with_gemm_swizzled_scales; + constexpr int FP4_BLOCK_SIZE = 16; const size_t N = input.flat_first_dim(); const size_t M = input.flat_last_dim(); @@ -95,15 +102,19 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) const size_t total = N * Mread; const size_t threads = 512; const size_t blocks = DIVUP(total, threads); + const size_t num_scale_tiles_X = DIVUP(Mread, static_cast(4)); TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( output->data.dtype, OType, - - dequantize_fp4_kernel<<>>( - input.data.dptr, reinterpret_cast(output->data.dptr), - reinterpret_cast(input.scale_inv.dptr), - reinterpret_cast(input.amax.dptr), N, Mread, - input.scale_inv.shape.back());); // NOLINT(*) + TRANSFORMER_ENGINE_SWITCH_CONDITION( + with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, + + dequantize_fp4_kernel<<>>( + input.data.dptr, reinterpret_cast(output->data.dptr), + reinterpret_cast(input.scale_inv.dptr), + reinterpret_cast(input.amax.dptr), N, Mread, input.scale_inv.shape.back(), + num_scale_tiles_X);); // NOLINT(*) + ); // NOLINT(*) NVTE_CHECK_CUDA(cudaGetLastError()); #else NVTE_ERROR("CUDA 12.8 or higher is needed for FP4 calculation!"); diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 720a27411..cab8abae1 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -518,25 +518,11 @@ def backward( ) elif ctx.backward_override == "dequantized": inputmats_dequant = [] - for m_split, inputmat in zip(ctx.m_splits, inputmats): + for inputmat in inputmats: if isinstance(inputmat, QuantizedTensorStorage): - if m_split == 0: - # Dequant kernels for some quantized storage formats - # (e.g. MXFP8/Float8BlockScaling) do not accept empty - # M-dimension inputs. For empty grouped splits, materialize - # an explicit empty high-precision matrix instead of invoking - # dequantize(). - inputmats_dequant.append( - torch.empty( - (0, ctx.weights_shape_1), - dtype=ctx.activation_dtype, - device=ctx.device, - ) - ) - else: - inputmats_dequant.append( - inputmat.dequantize(dtype=ctx.activation_dtype) - ) + inputmats_dequant.append( + inputmat.dequantize(dtype=ctx.activation_dtype) + ) else: inputmats_dequant.append(cast_if_needed(inputmat, ctx.activation_dtype)) inputmats = inputmats_dequant @@ -1331,15 +1317,6 @@ def _get_quantizers(self): for i in range(self.num_gemms): grad_output_quantizers[i].internal = True grad_output_quantizers[i].optimize_for_gemm = True - fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() - if fp8_recipe.backward_override == "dequantized" and ( - fp8_recipe.mxfp8() or fp8_recipe.nvfp4() - ): - for input_quantizer in input_quantizers: - input_quantizer.optimize_for_gemm = False - if torch.is_grad_enabled(): - for grad_output_quantizer in grad_output_quantizers: - grad_output_quantizer.optimize_for_gemm = False return ( input_quantizers, weight_quantizers, diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index d69e643c4..abfa6af03 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -1731,13 +1731,6 @@ def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): grad_output_quantizer.optimize_for_gemm = True if fp8_grad: grad_input_quantizer = self.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_INPUT1] - fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() - if fp8_recipe.backward_override == "dequantized" and ( - fp8_recipe.mxfp8() or fp8_recipe.nvfp4() - ): - input_quantizer.optimize_for_gemm = False - if grad_output_quantizer is not None: - grad_output_quantizer.optimize_for_gemm = False return ( input_quantizer, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 7498760af..2b14eaaf2 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -1742,13 +1742,6 @@ def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): grad_output_quantizer.optimize_for_gemm = True if fp8_grad: grad_input_quantizer = self.quantizers["scaling_bwd"][FP8BwdTensorIdx.GRAD_INPUT1] - fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() - if fp8_recipe.backward_override == "dequantized" and ( - fp8_recipe.mxfp8() or fp8_recipe.nvfp4() - ): - input_quantizer.optimize_for_gemm = False - if grad_output_quantizer is not None: - grad_output_quantizer.optimize_for_gemm = False return ( input_quantizer, weight_quantizer, diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index 19fcf62ce..46d52f7ff 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -356,15 +356,6 @@ def reset_recipe_state(self, *, recipe: Optional[Recipe]) -> None: grad_output_quantizer.internal = True if not (self.tensor_parallel_mode == "row" and self.sequence_parallel): grad_output_quantizer.optimize_for_gemm = True - if FP8GlobalStateManager.is_fp8_enabled(): - fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() - if fp8_recipe.backward_override is not None and ( - fp8_recipe.mxfp8() or fp8_recipe.nvfp4() - ): - if input_quantizer is not None: - input_quantizer.optimize_for_gemm = False - if grad_output_quantizer is not None: - grad_output_quantizer.optimize_for_gemm = False # Configure weight quantizer # Note: This function may be called in base class constructor, From 3e07f5df07797d01f5c9058bac91c7bab65fad42 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Sun, 3 May 2026 20:51:28 -0700 Subject: [PATCH 033/180] [JAX] Remove xla deterministic arg for MNIST test to not timeout L2_jax_unittest CI (#2952) remove xla deterministic arg to not timeout CI Signed-off-by: tdophung --- qa/L2_jax_unittest/test.sh | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/qa/L2_jax_unittest/test.sh b/qa/L2_jax_unittest/test.sh index 8441486e2..38cbc8ad3 100644 --- a/qa/L2_jax_unittest/test.sh +++ b/qa/L2_jax_unittest/test.sh @@ -31,11 +31,15 @@ mkdir -p "$XML_LOG_DIR" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_jax_not_distributed.xml $TE_PATH/tests/jax -k 'not distributed' || test_fail "tests/jax/*not_distributed_*" pip3 install -r $TE_PATH/examples/jax/mnist/requirements.txt || error_exit "Failed to install mnist requirements" -# Make mnist and encoder tests run-to-run deterministic for stable CI results -export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops" +# Note: mnist intentionally does NOT set --xla_gpu_deterministic_ops because it +# significantly slows down small conv/GEMM kernels and was causing CI timeouts. +# The mnist verify() already uses a tail-window min/max with relaxed thresholds +# to be robust to run-to-run numerical noise. NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_mnist.xml $TE_PATH/examples/jax/mnist || test_fail "mnist" pip3 install -r $TE_PATH/examples/jax/encoder/requirements.txt || error_exit "Failed to install encoder requirements" +# Make encoder tests to have run-to-run deterministic to have the stable CI results +export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops" NVTE_JAX_UNITTEST_LEVEL="L2" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_single_gpu_encoder.xml $TE_PATH/examples/jax/encoder/test_single_gpu_encoder.py || test_fail "test_single_gpu_encoder.py" # Test without custom calls export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops" From ad4b3fd19bb7b6e1536b819fc6ca6425a4d57ce2 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Mon, 4 May 2026 10:40:47 -0700 Subject: [PATCH 034/180] [PyTorch][Core] Fix CUBLAS GGEMM when weight dims are not divisible by 128 (#2954) * fix CUBLAS for GPT oss sizes Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review comments Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * ci fix Signed-off-by: Varun Thumbe * Fix test case dimensions in test_numerics.py Total dim should be divisible by 128 Signed-off-by: vthumbe1503 --------- Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/test_numerics.py | 69 ++++++++++++++----- .../common/cast/mxfp8/swizzle.cuh | 7 +- .../common/gemm/cublaslt_grouped_gemm.cu | 65 +++++++++++++++-- 3 files changed, 117 insertions(+), 24 deletions(-) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 5eef7f151..a718ea2a8 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -3084,7 +3084,10 @@ def _make_zero_tokens_grouped_tensor(logical_last_dim, is_a): weight_tensors = [torch.randn(n, k, dtype=dtype, device=device) for _ in range(z)] if use_mxfp8: grouped_A = _make_grouped_tensor_quantized_mxfp8( - weight_tensors, is_a=True, transposed=transa, device=device + weight_tensors, + rowwise=transa, + columnwise=not transa, + device=device, ) else: grouped_A = _make_grouped_tensor_uniform(z, n, k, device, dtype) @@ -3138,36 +3141,61 @@ def _make_zero_tokens_grouped_tensor(logical_last_dim, is_a): def _make_grouped_tensor_quantized_mxfp8( tensors: List[torch.Tensor], *, - is_a: bool, - transposed: bool, + rowwise: bool, + columnwise: bool, device: torch.device, - optimize_for_gemm: bool = True, + is_weight: bool = False, ) -> GroupedTensor: + """Create a quantized MXFP8 GroupedTensor from a list of per-expert tensors. + + For weights (uniform per-expert shape), we generally won't keep it swizzled since we + might need for future dequantize operations. Swizzling is done internally within + general_grouped_gemm_for_grouped_tensor call. + + For non-weight tensors (inputs / grad_outputs), we still pass + ``first_dims`` and keep ``optimize_for_gemm=True``; so the kernel must emit the + already-swizzled layout up front. + """ if not tensors: raise ValueError("Expected non-empty tensor list for grouped quantization.") - if is_a: - rowwise = transposed - columnwise = not transposed - else: - rowwise = not transposed - columnwise = transposed quantizer = MXFP8Quantizer( fp8_dtype=tex.DType.kFloat8E4M3, rowwise=rowwise, columnwise=columnwise, ) - quantizer.optimize_for_gemm = optimize_for_gemm + quantizer.optimize_for_gemm = not is_weight grouped_input = torch.cat(tensors, dim=0) - first_dims = torch.tensor([t.shape[0] for t in tensors], dtype=torch.int64, device=device) + if is_weight: + first_dims = None + else: + first_dims = torch.tensor([t.shape[0] for t in tensors], dtype=torch.int64, device=device) return tex.group_quantize(grouped_input, quantizer, len(tensors), first_dims) +def _per_tensor_quantize_mxfp8( + tensors: List[torch.Tensor], + *, + rowwise: bool, + columnwise: bool, +) -> List: + """Quantize each tensor individually with MXFP8. + Used to build reference discrete inputs for grouped GEMM. + """ + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + ) + return [quantizer(t) for t in tensors] + + @pytest.mark.parametrize( "shape", [ (1, 128, 128, 512), (8, 1024, 128, 512), (16, 4096, 128, 512), + (2, 256, 2880, 2880), ], ) @pytest.mark.parametrize("accumulate", [False, True]) @@ -3208,12 +3236,21 @@ def test_grouped_gemm_grouped_tensor_mxfp8( transa = layout[0] == "T" transb = layout[1] == "T" - grouped_A = _make_grouped_tensor_quantized_mxfp8(A, is_a=True, transposed=transa, device="cuda") + a_is_weight = all(t.shape == A[0].shape for t in A) + a_rowwise, a_columnwise = transa, not transa + b_rowwise, b_columnwise = not transb, transb + grouped_A = _make_grouped_tensor_quantized_mxfp8( + A, + rowwise=a_rowwise, + columnwise=a_columnwise, + device="cuda", + is_weight=a_is_weight, + ) grouped_B = _make_grouped_tensor_quantized_mxfp8( - B, is_a=False, transposed=transb, device="cuda" + B, rowwise=b_rowwise, columnwise=b_columnwise, device="cuda" ) - A_fp8 = grouped_A.split_into_quantized_tensors() - B_fp8 = grouped_B.split_into_quantized_tensors() + A_fp8 = _per_tensor_quantize_mxfp8(A, rowwise=a_rowwise, columnwise=a_columnwise) + B_fp8 = _per_tensor_quantize_mxfp8(B, rowwise=b_rowwise, columnwise=b_columnwise) general_grouped_gemm( A_fp8, diff --git a/transformer_engine/common/cast/mxfp8/swizzle.cuh b/transformer_engine/common/cast/mxfp8/swizzle.cuh index 7648e3f5c..e3876eb90 100644 --- a/transformer_engine/common/cast/mxfp8/swizzle.cuh +++ b/transformer_engine/common/cast/mxfp8/swizzle.cuh @@ -16,6 +16,9 @@ namespace dispatch { namespace mxfp8 { namespace swizzle { +constexpr size_t GEMM_SWIZZLED_SCALE_TILE_DIM_X = 4; +constexpr size_t GEMM_SWIZZLED_SCALE_TILE_DIM_Y = 128; + /*! \brief Convert compact scale indices into GEMM swizzled scale index * * MXFP8 GEMM expects scaling factors to be in a "swizzled" order @@ -25,8 +28,8 @@ namespace swizzle { * */ __device__ __forceinline__ size_t gemm_swizzled_scale_idx(size_t i, size_t j, size_t num_tiles_X) { - constexpr size_t TILE_DIM_X = 4; // Tile dim in scale buffer - constexpr size_t TILE_DIM_Y = 128; + constexpr size_t TILE_DIM_X = GEMM_SWIZZLED_SCALE_TILE_DIM_X; + constexpr size_t TILE_DIM_Y = GEMM_SWIZZLED_SCALE_TILE_DIM_Y; constexpr size_t TILE_SIZE = TILE_DIM_X * TILE_DIM_Y; const size_t tile_idx_X = j / TILE_DIM_X; const size_t tile_idx_Y = i / TILE_DIM_Y; diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index ed2275b44..6a7af158e 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -15,6 +15,7 @@ #include #include +#include "../cast/mxfp8/swizzle.cuh" #include "../common.h" #include "../util/cuda_runtime.h" #include "../util/handle_manager.h" @@ -330,6 +331,7 @@ struct GroupedOperandSelection { NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING; bool with_gemm_swizzled_scales = false; bool trans = false; + bool rowwise = true; }; constexpr int kMaxGroups = 64; @@ -613,6 +615,7 @@ inline GroupedOperandSelection select_grouped_operand(const transformer_engine:: sel.dptr = static_cast(t->columnwise_data.dptr); sel.scale_inv = t->columnwise_scale_inv.dptr; sel.dtype = col_dtype; + sel.rowwise = false; sel.shape = create_shape_info(t, swap_dims); }; @@ -621,6 +624,7 @@ inline GroupedOperandSelection select_grouped_operand(const transformer_engine:: sel.dptr = static_cast(t->data.dptr); sel.scale_inv = t->scale_inv.dptr; sel.dtype = row_dtype; + sel.rowwise = true; sel.shape = create_shape_info(t, /*swap_dims=*/false); }; @@ -846,6 +850,45 @@ __forceinline__ __device__ int64_t compute_grouped_tensor_offset(const TensorSha } } +__forceinline__ __device__ int64_t padded_mxfp8_scale_inv_bytes(int64_t first, int64_t last, + bool rowwise) { + namespace mxfp8_swizzle = transformer_engine::dispatch::mxfp8::swizzle; + constexpr int64_t kMxfp8BlockSize = 32; + // x is the dimension along which quantization is applied, y is other dimension + const int64_t scale_tile_y = static_cast(mxfp8_swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_Y); + const int64_t scale_tile_x = static_cast(mxfp8_swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_X); + // Padded byte size of the swizzled MXFP8 scale_inv for a single tensor with data + // shape (first, last). Rowwise scales use rows=first, cols=last; columnwise + // scales swap the orientation since they are stored in column-major order. + const int64_t scale_dim_y = rowwise ? first : last; + const int64_t padded_scale_dim_y = + ((scale_dim_y + scale_tile_y - 1) / scale_tile_y) * scale_tile_y; + const int64_t data_dim_x = rowwise ? last : first; + const int64_t scale_dim_x = (data_dim_x + kMxfp8BlockSize - 1) / kMxfp8BlockSize; + const int64_t padded_scale_dim_x = + ((scale_dim_x + scale_tile_x - 1) / scale_tile_x) * scale_tile_x; + // MXFP8 scales are E8M0 (1 byte per element), so element count == byte count. + return padded_scale_dim_y * padded_scale_dim_x; +} + +// Device helper: byte offset into a contiguous grouped MXFP8 scale_inv buffer for +// tensor `idx`. Each expert's scale_inv is expected to be padded +// to the 128x4 swizzled layout. +__forceinline__ __device__ int64_t compute_grouped_tensor_mxfp8_scale_inv_offset( + const TensorShapeInfo &meta, size_t idx, bool rowwise) { + if (meta.first_dims != nullptr || meta.last_dims != nullptr) { + int64_t cumsum = 0; + for (size_t i = 0; i < idx; i++) { + const int64_t f = meta.first_dims ? meta.first_dims[i] : meta.uniform_first; + const int64_t l = meta.last_dims ? meta.last_dims[i] : meta.uniform_last; + cumsum += padded_mxfp8_scale_inv_bytes(f, l, rowwise); + } + return cumsum; + } + return static_cast(idx) * + padded_mxfp8_scale_inv_bytes(meta.uniform_first, meta.uniform_last, rowwise); +} + // Linear scan to find which tensor contains the given row. // Returns the tensor index and writes the exclusive end-row of that tensor to *out_tensor_row_end. __forceinline__ __device__ int find_tensor_for_row(const int64_t *first_dims, int64_t uniform_first, @@ -977,7 +1020,8 @@ __global__ void setup_grouped_gemm_kernel( size_t b_elem_size, size_t c_elem_size, size_t d_elem_size, float *alpha_ptr, float *beta_ptr, // Scale inputs: for tensor scaling, pass float* and set mxfp8_base to nullptr // For MXFP8, pass nullptr for tensor_scale and set mxfp8_base - float *a_scale_base, float *b_scale_base, NVTEScalingMode scaling_mode, size_t num_tensors, + float *a_scale_base, float *b_scale_base, bool a_rowwise, bool b_rowwise, + NVTEScalingMode scaling_mode, size_t num_tensors, MultiTensorGroupGemmInputArgs a_multi_tensor_args, MultiTensorGroupGemmOutputArgs c_multi_tensor_args, MultiTensorGroupGemmOutputArgs d_multi_tensor_args) { @@ -1038,12 +1082,13 @@ __global__ void setup_grouped_gemm_kernel( // Fill scale pointers (per-matrix). // The interpretation of the scale buffers depends on the shared scaling recipe: - // NVTE_MXFP8_1D_SCALING : E8M0 byte stream; offset = data_offset / 32 elements // otherwise : one float per tensor, indexed by tensor index if (a_scale_base) { if (scaling_mode == NVTE_MXFP8_1D_SCALING) { + const int64_t a_scale_offset = + compute_grouped_tensor_mxfp8_scale_inv_offset(A_meta, idx, a_rowwise); a_scale_inv_ptrs[idx] = reinterpret_cast( - static_cast(static_cast(a_scale_base)) + a_offset / 32); + static_cast(static_cast(a_scale_base)) + a_scale_offset); } else { a_scale_inv_ptrs[idx] = static_cast(a_scale_base) + idx; } @@ -1052,8 +1097,10 @@ __global__ void setup_grouped_gemm_kernel( } if (b_scale_base) { if (scaling_mode == NVTE_MXFP8_1D_SCALING) { + const int64_t b_scale_offset = + compute_grouped_tensor_mxfp8_scale_inv_offset(B_meta, idx, b_rowwise); b_scale_inv_ptrs[idx] = reinterpret_cast( - static_cast(static_cast(b_scale_base)) + b_offset / 32); + static_cast(static_cast(b_scale_base)) + b_scale_offset); } else { b_scale_inv_ptrs[idx] = static_cast(b_scale_base) + idx; } @@ -1116,14 +1163,19 @@ inline void launch_grouped_gemm_setup( // A and B share the same scaling recipe (validated in validate_grouped_gemm_inputs). // Pass scale buffers as void* and let the kernel interpret them via scaling_mode. + + // Scale rowwise flag for MXFP8/NVFP4: to calculate scale_inv padding based offsets + // within kernel. Ignored for tensor scaling. + const bool a_rowwise = A_sel.rowwise; + const bool b_rowwise = B_sel.rowwise; setup_grouped_gemm_kernel<<>>( ws.A_ptrs, ws.B_ptrs, ws.C_ptrs, ws.D_ptrs, ws.a_rows, ws.a_cols, ws.b_rows, ws.b_cols, ws.d_rows, ws.d_cols, ws.alpha_ptrs, ws.beta_ptrs, ws.a_scale_inv_ptrs, ws.b_scale_inv_ptrs, A_sel.dptr, B_sel.dptr, c_base, d_base, A_meta, B_meta, C_meta, D_meta, a_elem_size, b_elem_size, c_elem_size, d_elem_size, static_cast(alpha_tensor->data.dptr), static_cast(beta_tensor->data.dptr), reinterpret_cast(A_sel.scale_inv), - reinterpret_cast(B_sel.scale_inv), A_sel.scaling_mode, num_tensors, - a_multi_tensor_args, c_multi_tensor_args, d_multi_tensor_args); + reinterpret_cast(B_sel.scale_inv), a_rowwise, b_rowwise, A_sel.scaling_mode, + num_tensors, a_multi_tensor_args, c_multi_tensor_args, d_multi_tensor_args); NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -1276,6 +1328,7 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num choose_grouped_operand_storage(static_cast(transa), /*is_A=*/true, mxfp8, is_fp8, non_tn_fp8_ok, A_list_info.all_row, A_list_info.all_col, "A"); A_sel.trans = choice.trans; + A_sel.rowwise = choice.use_rowwise; if (choice.use_rowwise) { NVTE_CHECK(A_list_info.all_row, "Grouped GEMM: A_list is missing row-wise data"); A_sel.dtype = A_list_info.row_dtype; From 528f16c5067a50c5a4ec2b8f4c466d3372536323 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Mon, 4 May 2026 18:06:37 -0400 Subject: [PATCH 035/180] [PyTorch] Guard/document single parameter feature for grouped linear (#2955) * Better documentation for single param and envvar guard Signed-off-by: Kirthi Shankar Sivamani * fix doc Signed-off-by: ksivamani * Fix test envvar Signed-off-by: Kirthi Shankar Sivamani --------- Signed-off-by: Kirthi Shankar Sivamani Signed-off-by: ksivamani --- qa/L0_pytorch_debug_unittest/test.sh | 2 +- qa/L0_pytorch_unittest/test.sh | 6 ++-- .../pytorch/module/grouped_linear.py | 12 +++++-- .../pytorch/ops/basic/grouped_linear.py | 10 ++++++ transformer_engine/pytorch/utils.py | 31 +++++++++++++++++++ 5 files changed, 55 insertions(+), 6 deletions(-) diff --git a/qa/L0_pytorch_debug_unittest/test.sh b/qa/L0_pytorch_debug_unittest/test.sh index ce65bc430..3efa46262 100644 --- a/qa/L0_pytorch_debug_unittest/test.sh +++ b/qa/L0_pytorch_debug_unittest/test.sh @@ -36,7 +36,7 @@ NVTE_TORCH_COMPILE=0 pytest -v -s --junitxml=$XML_LOG_DIR/test_api_features.xml pytest -v -s --junitxml=$XML_LOG_DIR/test_perf.xml $TE_PATH/tests/pytorch/debug/test_perf.py --feature_dirs=$NVTE_TEST_NVINSPECT_FEATURE_DIRS --configs_dir=$NVTE_TEST_NVINSPECT_CONFIGS_DIR || test_fail "test_perf.py" # standard sanity and numerics tests with initialized debug -NVTE_TEST_NVINSPECT_ENABLED=1 NVTE_TEST_NVINSPECT_CONFIG_FILE=$NVTE_TEST_NVINSPECT_DUMMY_CONFIG_FILE NVTE_TEST_NVINSPECT_FEATURE_DIRS=$NVTE_TEST_NVINSPECT_FEATURE_DIRS PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 pytest -v -s --junitxml=$XML_LOG_DIR/test_sanity_2.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "debug test_sanity.py" +NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 NVTE_TEST_NVINSPECT_ENABLED=1 NVTE_TEST_NVINSPECT_CONFIG_FILE=$NVTE_TEST_NVINSPECT_DUMMY_CONFIG_FILE NVTE_TEST_NVINSPECT_FEATURE_DIRS=$NVTE_TEST_NVINSPECT_FEATURE_DIRS PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 pytest -v -s --junitxml=$XML_LOG_DIR/test_sanity_2.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "debug test_sanity.py" NVTE_TEST_NVINSPECT_ENABLED=1 NVTE_TEST_NVINSPECT_CONFIG_FILE=$NVTE_TEST_NVINSPECT_DUMMY_CONFIG_FILE NVTE_TEST_NVINSPECT_FEATURE_DIRS=$NVTE_TEST_NVINSPECT_FEATURE_DIRS PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 pytest -v -s --junitxml=$XML_LOG_DIR/test_numerics_2.xml $TE_PATH/tests/pytorch/test_numerics.py || test_fail "debug test_numerics.py" if [ "$RET" -ne 0 ]; then diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index a8f8cf875..22636828f 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -24,7 +24,7 @@ mkdir -p "$XML_LOG_DIR" pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" -python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py" +NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_recipe.xml $TE_PATH/tests/pytorch/test_recipe.py || test_fail "test_recipe.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_deferred_init.xml $TE_PATH/tests/pytorch/test_deferred_init.py || test_fail "test_deferred_init.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/test_numerics.py || test_fail "test_numerics.py" @@ -37,11 +37,11 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_quantized_tensor python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8blockwisetensor.xml $TE_PATH/tests/pytorch/test_float8blockwisetensor.py || test_fail "test_float8blockwisetensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_scaling_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_scaling_exact.py || test_fail "test_float8_blockwise_scaling_exact.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_float8_blockwise_gemm_exact.xml $TE_PATH/tests/pytorch/test_float8_blockwise_gemm_exact.py || test_fail "test_float8_blockwise_gemm_exact.py" -python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/test_grouped_tensor.xml $TE_PATH/tests/pytorch/test_grouped_tensor.py || test_fail "test_grouped_tensor.py" +NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/test_grouped_tensor.xml $TE_PATH/tests/pytorch/test_grouped_tensor.py || test_fail "test_grouped_tensor.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_gqa.xml $TE_PATH/tests/pytorch/test_gqa.py || test_fail "test_gqa.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_optimizer.xml $TE_PATH/tests/pytorch/test_fused_optimizer.py || test_fail "test_fused_optimizer.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_multi_tensor.xml $TE_PATH/tests/pytorch/test_multi_tensor.py || test_fail "test_multi_tensor.py" -NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/test_fusible_ops.py || test_fail "test_fusible_ops.py" +NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 NVTE_CUTEDSL_FUSED_GROUPED_MLP=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_PATH/tests/pytorch/test_fusible_ops.py || test_fail "test_fusible_ops.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_backward_override.xml $TE_PATH/tests/pytorch/test_backward_override.py || test_fail "test_backward_override.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_permutation.xml $TE_PATH/tests/pytorch/test_permutation.py || test_fail "test_permutation.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_parallel_cross_entropy.xml $TE_PATH/tests/pytorch/test_parallel_cross_entropy.py || test_fail "test_parallel_cross_entropy.py" diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index cab8abae1..4ae7b47b9 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -31,6 +31,7 @@ clear_tensor_data, init_method_constant, requires_grad, + resolve_grouped_linear_single_param_flags, get_nvtx_range_context, ) from ..distributed import ( @@ -659,11 +660,15 @@ class GroupedLinear(TransformerEngineBaseModule): single_grouped_weight : bool, default = False If set to ``True``, grouped weights are stored as a single grouped parameter instead of one parameter per GEMM. - EXPERIMENTAL and subject to change. + EXPERIMENTAL and subject to change. Gated by the + ``NVTE_GROUPED_LINEAR_SINGLE_PARAM`` environment variable: if the env var + is not set this argument is forced to ``False`` with a warning. single_grouped_bias : bool, default = False If set to ``True``, grouped biases are stored as a single grouped bias instead of one bias per GEMM. - EXPERIMENTAL and subject to change. + EXPERIMENTAL and subject to change. Gated by the + ``NVTE_GROUPED_LINEAR_SINGLE_PARAM`` environment variable: if the env var + is not set this argument is forced to ``False`` with a warning. Notes ----- @@ -712,6 +717,9 @@ def __init__( self.ub_overlap_ag = ub_overlap_ag self.ub_name = ub_name self.save_original_input = save_original_input + single_grouped_weight, single_grouped_bias = resolve_grouped_linear_single_param_flags( + single_grouped_weight, single_grouped_bias + ) self.single_grouped_weight = single_grouped_weight self.single_grouped_bias = single_grouped_bias if ub_overlap_rs or ub_overlap_ag: diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index b503cb186..a86abb132 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -30,6 +30,7 @@ canonicalize_dtype, clear_tensor_data, devices_match, + resolve_grouped_linear_single_param_flags, round_up_to_nearest_multiple, ) from .._common import is_quantized_tensor, maybe_dequantize @@ -78,11 +79,17 @@ class GroupedLinear(BasicOperation): ``main_grad`` instead of accumulating. single_grouped_weight : bool, default = ``False`` Store all expert weights as one ``GroupedTensor`` parameter ``weight``. + EXPERIMENTAL and subject to change. Gated by the + ``NVTE_GROUPED_LINEAR_SINGLE_PARAM`` environment variable: if the env var + is not set this argument is forced to ``False`` with a warning. delay_wgrad_compute : bool, default = ``False`` Whether to delay weight gradient computation single_grouped_bias : bool, default = ``False`` If ``True`` (and ``bias=True``), store all expert biases as one ``GroupedTensor`` parameter named ``bias`` instead of ``bias0``..``bias{N-1}``. + EXPERIMENTAL and subject to change. Gated by the + ``NVTE_GROUPED_LINEAR_SINGLE_PARAM`` environment variable: if the env var + is not set this argument is forced to ``False`` with a warning. scale_bias : bool, default = ``False`` If ``True`` (and ``bias=True``), expects a probability tensor as an additional extra input and adds ``bias * scales`` instead of ``bias`` @@ -123,6 +130,9 @@ def __init__( self.num_groups: int = num_groups self.in_features: int = in_features self.out_features: int = out_features + single_grouped_weight, single_grouped_bias = resolve_grouped_linear_single_param_flags( + single_grouped_weight, single_grouped_bias + ) self.single_grouped_weight: bool = single_grouped_weight self.single_grouped_bias: bool = single_grouped_bias self.use_bias: bool = bias diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index a76f205ac..250daec67 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -7,6 +7,7 @@ import functools import math import os +import warnings from typing import Any, Callable, List, Optional, Sequence, Tuple, Union from contextlib import nullcontext import numpy as np @@ -81,6 +82,36 @@ def get_device_compute_capability() -> Tuple[int, int]: return _get_device_compute_capability(torch.cuda.current_device()) +def resolve_grouped_linear_single_param_flags( + single_grouped_weight: bool, + single_grouped_bias: bool, +) -> Tuple[bool, bool]: + """Gate ``single_grouped_weight`` / ``single_grouped_bias`` on ``NVTE_GROUPED_LINEAR_SINGLE_PARAM``.""" + if not (single_grouped_weight or single_grouped_bias): + return single_grouped_weight, single_grouped_bias + + env_enabled = int(os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0")) > 0 + if not env_enabled: + warnings.warn( + f"GroupedLinear was constructed with single_grouped_weight={single_grouped_weight} " + f"and single_grouped_bias={single_grouped_bias}, but the " + "NVTE_GROUPED_LINEAR_SINGLE_PARAM environment variable is not set. " + "Disabling single grouped weight/bias and falling back to per-expert parameters.", + UserWarning, + stacklevel=3, + ) + return False, False + + warnings.warn( + "GroupedLinear is using single_grouped_weight/single_grouped_bias. " + "This feature is experimental, may change in future " + "releases, and is known to be non-deterministic in certain cases.", + UserWarning, + stacklevel=3, + ) + return single_grouped_weight, single_grouped_bias + + def attention_mask_func( attention_scores: torch.Tensor, attention_mask: torch.Tensor ) -> torch.Tensor: From 3ded616899b1af4ae1be9ceaffa8012d4373567d Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Tue, 5 May 2026 16:22:13 -0700 Subject: [PATCH 036/180] Graph Safe support for TE Grouped linear Op (#2923) * starting effort Signed-off-by: Varun Thumbe * all tests seem to be working Signed-off-by: Varun Thumbe * cuda graph test + clean ups Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * clean up Signed-off-by: Varun Thumbe * cleanup Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * clean up main grad business Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * clean up a but Signed-off-by: Varun Thumbe * fix on l40/hopper to skip Signed-off-by: Varun Thumbe * address review comments + save activation in backward + common context savings for fused/unfused paths Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * cleanup Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review comments Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix linting errors Signed-off-by: Varun Thumbe --------- Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/test_fusible_ops.py | 337 +++++-- tests/pytorch/utils.py | 72 ++ transformer_engine/pytorch/ops/_common.py | 94 ++ .../pytorch/ops/basic/basic_linear.py | 32 +- .../pytorch/ops/basic/grouped_linear.py | 874 +++++++++++++++--- .../pytorch/ops/fused/backward_grouped_mlp.py | 193 ++-- .../pytorch/ops/fused/backward_linear_add.py | 29 +- .../ops/fused/backward_linear_scale.py | 29 +- .../pytorch/ops/fused/forward_grouped_mlp.py | 79 +- .../ops/fused/userbuffers_backward_linear.py | 32 +- .../tensor/storage/grouped_tensor_storage.py | 45 + 11 files changed, 1336 insertions(+), 480 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 47507dc38..7691582f9 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -7,6 +7,7 @@ from collections.abc import Iterable, Sequence import functools import io +import os import math import random from typing import Optional @@ -42,7 +43,6 @@ ) from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor from transformer_engine.pytorch.cpp_extensions.gemm import general_grouped_gemm_for_grouped_tensor -from transformer_engine.pytorch.module.base import get_dummy_wgrad import transformer_engine_torch as tex # Import utility functions @@ -51,6 +51,7 @@ assert_close_grads, dtype_tols, make_recipe, + MegatronTrainingHelper, quantization_tols, reset_rng_states, ) @@ -212,76 +213,6 @@ def to_cpu(tensor: Optional[torch.Tensor]) -> Optional[torch.Tensor]: return out -class MegatronTrainingHelper: - """Test-side stand-in for the Megatron-Core DDP / MegatronFSDP wrapper. - Megatron's DDP wrapper (and MegatronFSDP) owns the per-parameter - ``main_grad`` buffer and the ``overwrite_main_grad`` / - ``grad_added_to_main_grad`` attributes that coordinate - ``fuse_wgrad_accumulation`` with TE modules. These helpers reproduce the - relevant slice of that protocol so TE tests can exercise the - accumulate-into-``main_grad`` code path without pulling in the full - Megatron-Core dependency. - """ - - @staticmethod - def init_main_grad_buffers( - weight_params: Iterable[torch.nn.Parameter], - *, - fill_value: float, - overwrite_main_grad: bool, - zero_out_wgrad: bool = False, - dtype: torch.dtype = torch.float32, - ) -> None: - """Allocate ``main_grad`` and stamp the wrapper attributes on each - param, mirroring what the Megatron DDP/FSDP wrapper does before - backward.""" - for wp in weight_params: - wp.main_grad = torch.full(wp.size(), fill_value, device=wp.device, dtype=dtype) - wp.overwrite_main_grad = overwrite_main_grad - wp.zero_out_wgrad = zero_out_wgrad - wp.grad_added_to_main_grad = False - - @staticmethod - def verify_main_grad_accumulation( - weight_params: Iterable[torch.nn.Parameter], - *, - expected_main_grads: Iterable[torch.Tensor], - rtol: float = 0.0, - atol: float = 0.0, - ) -> None: - """Check that backward produced what the Megatron wrapper expects: - each ``main_grad`` matches ``expected_main_grads``, - ``grad_added_to_main_grad`` was flipped to ``True`` so the wrapper's - post-backward hooks won't double-accumulate, and ``param.grad`` was - replaced by the cached dummy tensor (so a wrapper hook that did - ``main_grad += grad`` would be a no-op rather than double-counting). - """ - for wp, expected in zip(weight_params, expected_main_grads): - torch.testing.assert_close(wp.main_grad.to(expected), expected, rtol=rtol, atol=atol) - - assert wp.grad_added_to_main_grad is True, ( - "weight.grad_added_to_main_grad was not flipped to True; " - "the Megatron DDP/FSDP wrapper hook will double-accumulate." - ) - - # ``.grad`` should be the cached dummy tensor returned by - # ``get_dummy_wgrad`` -- shared storage, not the real wgrad. - expected_dummy = get_dummy_wgrad(list(wp.size()), wp.dtype) - assert ( - wp.grad is not None - ), "weight.grad is None; the Megatron protocol expects a dummy tensor stand-in here." - assert wp.grad.data_ptr() == expected_dummy.data_ptr(), ( - "weight.grad does not share storage with the cached dummy " - "wgrad; downstream wrapper hooks risk double-accumulating." - ) - if getattr(wp, "zero_out_wgrad", False): - assert torch.all(wp.grad == 0), ( - "weight.zero_out_wgrad=True but the dummy weight.grad " - "was not zeroed; downstream hooks reading .grad would " - "see stale bytes from the previous step." - ) - - class TestSequentialContainer: """Tests for sequential container""" @@ -2098,6 +2029,8 @@ def test_dropout( @pytest.mark.parametrize("input_requires_grad", (False, True)) @pytest.mark.parametrize("weight_requires_grad", (False, True)) @pytest.mark.parametrize("delay_wgrad_compute", (False, True)) + @pytest.mark.parametrize("single_grouped_weight", (False, True)) + @pytest.mark.parametrize("single_grouped_bias", (False, True)) def test_grouped_linear( self, *, @@ -2113,9 +2046,17 @@ def test_grouped_linear( input_requires_grad: bool, weight_requires_grad: bool, delay_wgrad_compute: bool, + single_grouped_weight: bool, + single_grouped_bias: bool, ) -> None: """Grouped GEMM""" - + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and ( + single_grouped_weight or single_grouped_bias + ): + pytest.skip( + "single_grouped_weight/single_grouped_bias requires" + " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1" + ) # Split sizes split_sizes = [split_alignment * i for i in range(group_size)] random.shuffle(split_sizes) @@ -2136,6 +2077,18 @@ def test_grouped_linear( if quantization is not None and dtype not in (torch.bfloat16, torch.float16): pytest.skip("Quantized group GEMM is only supported with BF16/FP16") + if single_grouped_bias and not bias: + pytest.skip("single_grouped_bias requires bias=True") + if ( + single_grouped_weight + and quantized_weight + and quantization in ("fp8_delayed_scaling", "fp8_current_scaling") + ): + pytest.skip( + "single_grouped_weight does not support FP8 delayed/current scaling " + "with quantized_model_init" + ) + # Random data x_ref, x_test = make_reference_and_test_tensors( in_shape, @@ -2194,12 +2147,26 @@ def test_grouped_linear( device=device, dtype=dtype, delay_wgrad_compute=delay_wgrad_compute, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, ) with torch.no_grad(): + if single_grouped_weight: + op_weights = op.weight.quantized_tensors + if op_weights is None: + op_weights = op.weight.split_into_quantized_tensors() + if single_grouped_bias: + op_bias_parts = op.bias.split_into_quantized_tensors() for group_idx in range(group_size): - getattr(op, f"weight{group_idx}").copy_(ws_test[group_idx]) + if single_grouped_weight: + op_weights[group_idx].copy_(ws_test[group_idx]) + else: + getattr(op, f"weight{group_idx}").copy_(ws_test[group_idx]) if bias: - getattr(op, f"bias{group_idx}").copy_(bs_test[group_idx]) + if single_grouped_bias: + op_bias_parts[group_idx].reshape(-1).copy_(bs_test[group_idx]) + else: + getattr(op, f"bias{group_idx}").copy_(bs_test[group_idx]) del ws_test, bs_test for param in op.parameters(): param.requires_grad_(requires_grad=weight_requires_grad) @@ -2227,20 +2194,222 @@ def test_grouped_linear( torch.testing.assert_close(dx_test, x_ref.grad, **tols) else: assert x_test.grad is None - for group_idx in range(group_size): - w_test = getattr(op, f"weight{group_idx}") + if single_grouped_weight: if weight_requires_grad: - dw_test = w_test.grad.to(dtype=torch.float64, device="cpu") - torch.testing.assert_close(dw_test, ws_ref[group_idx].grad, **tols) + dw_test_all = op.weight.grad.to(dtype=torch.float64, device="cpu") + w_ref_grad = torch.stack([w.grad for w in ws_ref], dim=0) + torch.testing.assert_close(dw_test_all, w_ref_grad, **tols) else: - assert w_test.grad is None - if bias: - b_test = getattr(op, f"bias{group_idx}") + assert op.weight.grad is None + else: + for group_idx in range(group_size): + w_test = getattr(op, f"weight{group_idx}") if weight_requires_grad: - db_test = b_test.grad.to(dtype=torch.float64, device="cpu") - torch.testing.assert_close(db_test, bs_ref[group_idx].grad, **tols) + dw_test = w_test.grad.to(dtype=torch.float64, device="cpu") + torch.testing.assert_close(dw_test, ws_ref[group_idx].grad, **tols) else: - assert b_test.grad is None + assert w_test.grad is None + if bias: + if single_grouped_bias: + if weight_requires_grad: + db_test_all = op.bias.grad.to(dtype=torch.float64, device="cpu") + b_ref_grad = torch.stack([b.grad for b in bs_ref], dim=0) + torch.testing.assert_close(db_test_all, b_ref_grad, **tols) + else: + assert op.bias.grad is None + else: + for group_idx in range(group_size): + b_test = getattr(op, f"bias{group_idx}") + if weight_requires_grad: + db_test = b_test.grad.to(dtype=torch.float64, device="cpu") + torch.testing.assert_close(db_test, bs_ref[group_idx].grad, **tols) + else: + assert b_test.grad is None + + @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16)) + @pytest.mark.parametrize( + "quantization", + [None] + (["mxfp8"] if mxfp8_available else []), + ) + @pytest.mark.parametrize("quantized_weight", (False, True)) + @pytest.mark.parametrize("bias", (False, True)) + @pytest.mark.parametrize("single_grouped_weight", (False, True)) + @pytest.mark.parametrize("single_grouped_bias", (False, True)) + @pytest.mark.parametrize("accumulate_into_main_grad", (False, True)) + def test_grouped_linear_cuda_graph_safe( + self, + *, + dtype: torch.dtype, + quantization: Optional[str], + quantized_weight: bool, + bias: bool, + single_grouped_weight: bool, + single_grouped_bias: bool, + accumulate_into_main_grad: bool, + device: torch.device = "cuda", + group_size: int = 4, + in_features: int = 128, + out_features: int = 128, + split_alignment: int = 128, + token_padding: int = 256, + ) -> None: + """GroupedLinear forward+backward should be CUDA graph capturable. + + Exercises the grouped-tensor / cublas-grouped-gemm path which uses + GPU-resident split offsets and is the only flow safe to capture. + """ + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and ( + single_grouped_weight or single_grouped_bias + ): + pytest.skip( + "single_grouped_weight/single_grouped_bias requires" + " NVTE_GROUPED_LINEAR_SINGLE_PARAM=1" + ) + if torch.cuda.get_device_capability() < (10, 0): + pytest.skip("Grouped GEMM CUDA-graph-safe path requires SM100+ (Blackwell)") + # Skip invalid configurations + if quantization is None and quantized_weight: + pytest.skip("quantized_weight requires a quantization recipe") + if single_grouped_bias and not bias: + pytest.skip("single_grouped_bias requires bias=True") + + # Split sizes (statically pinned for graph capture) + split_sizes = [split_alignment * (i + 1) for i in range(group_size)] + random.shuffle(split_sizes) + split_sizes = torch.tensor(split_sizes, dtype=torch.int, device=device) + # Pad input tokens to validate the sync-free flow + in_shape = (split_sizes.sum().item() + token_padding, in_features) + out_shape = (in_shape[0], out_features) + + recipe = make_recipe(quantization) + with te.quantized_model_init(enabled=quantized_weight, recipe=recipe): + op = te_ops.GroupedLinear( + group_size, + in_features, + out_features, + bias=bias, + device=device, + dtype=dtype, + accumulate_into_main_grad=accumulate_into_main_grad, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, + ) + + def _weight_params() -> list[torch.nn.Parameter]: + if single_grouped_weight: + return [op.weight] + return [getattr(op, f"weight{i}") for i in range(group_size)] + + def _bias_params() -> list[torch.nn.Parameter]: + if not bias: + return [] + if single_grouped_bias: + return [op.bias] + return [getattr(op, f"bias{i}") for i in range(group_size)] + + def _init_main_grads(value: float = 0.0) -> None: + if not accumulate_into_main_grad: + return + with torch.no_grad(): + for w in _weight_params(): + if getattr(w, "main_grad", None) is None: + w.main_grad = torch.empty(w.size(), device=device, dtype=torch.float32) + w.main_grad.fill_(value) + + def _collect_main_grads() -> list[torch.Tensor]: + return [w.main_grad.detach().clone() for w in _weight_params()] + + def _zero_param_grads() -> None: + for param in op.parameters(): + if param.grad is None: + param.grad = torch.zeros_like(param) + else: + param.grad.zero_() + + static_split_sizes = split_sizes.clone() + + def train_step( + x: torch.Tensor, + dy: torch.Tensor, + out_buf: torch.Tensor, + *, + use_graphed: bool, + ) -> torch.Tensor: + with te.autocast(enabled=quantization is not None, recipe=recipe): + out = ( + graphed_module(x, static_split_sizes) + if use_graphed + else op(x, static_split_sizes) + ) + out.backward(dy) + out_buf.copy_(out) + return out_buf + + _init_main_grads(0.0) + + static_x = torch.randn(in_shape, device=device, dtype=dtype, requires_grad=True) + static_dy = torch.randn(out_shape, device=device, dtype=dtype) + static_out_buf = torch.empty(out_shape, device=device, dtype=dtype) + + graphed_module = te.make_graphed_callables( + op, + (static_x, static_split_sizes), + num_warmup_iters=3, + enabled=quantization is not None, + recipe=recipe, + ) + + # Replace static buffers with fresh data (graph captures must replay + # against new inputs without re-recording). + fresh_x = torch.randn_like(static_x) + fresh_dy = torch.randn_like(static_dy) + with torch.no_grad(): + static_x.copy_(fresh_x) + static_dy.copy_(fresh_dy) + + # Reset grads & main_grads so the captured iteration starts fresh. + _zero_param_grads() + _init_main_grads(0.5) + if static_x.grad is not None: + static_x.grad.zero_() + + # Replay the graph + graph_out = ( + train_step(static_x, static_dy, static_out_buf, use_graphed=True).detach().clone() + ) + torch.cuda.synchronize() + graph_dx = static_x.grad.detach().clone() + if accumulate_into_main_grad: + graph_main_grads = _collect_main_grads() + graph_param_grads: list[torch.Tensor] = [] + else: + graph_main_grads = [] + graph_param_grads = [param.grad.detach().clone() for param in op.parameters()] + + # Reference: same op invoked eagerly with the same fresh inputs and + # the same starting grad/main_grad state. + _zero_param_grads() + _init_main_grads(0.5) + static_x.grad.zero_() + + expected_x = fresh_x.detach().clone().requires_grad_(True) + expected_dy = fresh_dy.detach().clone() + with te.autocast(enabled=quantization is not None, recipe=recipe): + expected_out = op(expected_x, static_split_sizes) + expected_out.backward(expected_dy) + + tols = dtype_tols(dtype) + if quantization is not None: + tols = quantization_tols(quantization) + + assert_close(graph_out, expected_out, **tols) + assert_close(graph_dx, expected_x.grad, **tols) + if accumulate_into_main_grad: + for g, w in zip(graph_main_grads, _weight_params()): + assert_close(g, w.main_grad, **tols) + else: + for g, param in zip(graph_param_grads, op.parameters()): + assert_close(g, param.grad, **tols) @pytest.mark.parametrize("in_shape", ((71, 192), (5, 7, 128))) @pytest.mark.parametrize("input_requires_grad", (False, True)) diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 8f8852edc..c7cbe78a6 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -8,6 +8,7 @@ import os import random import subprocess +from collections.abc import Iterable from contextlib import contextmanager from typing import Optional, Sequence, Tuple, Dict, Any, List from packaging.version import Version as PkgVersion @@ -27,6 +28,7 @@ check_set_window_size, ) from transformer_engine.pytorch.cpp_extensions.fused_attn import FusedAttnBackend +from transformer_engine.pytorch.module.base import get_dummy_wgrad def str_to_dtype(dtype: str | torch.dtype) -> torch.dtype: @@ -477,3 +479,73 @@ def run_distributed( msg += f"\n--- stderr ---\n{stderr_tail}" raise AssertionError(msg) return result + + +class MegatronTrainingHelper: + """Test-side stand-in for the Megatron-Core DDP / MegatronFSDP wrapper. + Megatron's DDP wrapper (and MegatronFSDP) owns the per-parameter + ``main_grad`` buffer and the ``overwrite_main_grad`` / + ``grad_added_to_main_grad`` attributes that coordinate + ``fuse_wgrad_accumulation`` with TE modules. These helpers reproduce the + relevant slice of that protocol so TE tests can exercise the + accumulate-into-``main_grad`` code path without pulling in the full + Megatron-Core dependency. + """ + + @staticmethod + def init_main_grad_buffers( + weight_params: Iterable[torch.nn.Parameter], + *, + fill_value: float, + overwrite_main_grad: bool, + zero_out_wgrad: bool = False, + dtype: torch.dtype = torch.float32, + ) -> None: + """Allocate ``main_grad`` and stamp the wrapper attributes on each + param, mirroring what the Megatron DDP/FSDP wrapper does before + backward.""" + for wp in weight_params: + wp.main_grad = torch.full(wp.size(), fill_value, device=wp.device, dtype=dtype) + wp.overwrite_main_grad = overwrite_main_grad + wp.zero_out_wgrad = zero_out_wgrad + wp.grad_added_to_main_grad = False + + @staticmethod + def verify_main_grad_accumulation( + weight_params: Iterable[torch.nn.Parameter], + *, + expected_main_grads: Iterable[torch.Tensor], + rtol: float = 0.0, + atol: float = 0.0, + ) -> None: + """Check that backward produced what the Megatron wrapper expects: + each ``main_grad`` matches ``expected_main_grads``, + ``grad_added_to_main_grad`` was flipped to ``True`` so the wrapper's + post-backward hooks won't double-accumulate, and ``param.grad`` was + replaced by the cached dummy tensor (so a wrapper hook that did + ``main_grad += grad`` would be a no-op rather than double-counting). + """ + for wp, expected in zip(weight_params, expected_main_grads): + torch.testing.assert_close(wp.main_grad.to(expected), expected, rtol=rtol, atol=atol) + + assert wp.grad_added_to_main_grad is True, ( + "weight.grad_added_to_main_grad was not flipped to True; " + "the Megatron DDP/FSDP wrapper hook will double-accumulate." + ) + + # ``.grad`` should be the cached dummy tensor returned by + # ``get_dummy_wgrad`` -- shared storage, not the real wgrad. + expected_dummy = get_dummy_wgrad(list(wp.size()), wp.dtype) + assert ( + wp.grad is not None + ), "weight.grad is None; the Megatron protocol expects a dummy tensor stand-in here." + assert wp.grad.data_ptr() == expected_dummy.data_ptr(), ( + "weight.grad does not share storage with the cached dummy " + "wgrad; downstream wrapper hooks risk double-accumulating." + ) + if getattr(wp, "zero_out_wgrad", False): + assert torch.all(wp.grad == 0), ( + "weight.zero_out_wgrad=True but the dummy weight.grad " + "was not zeroed; downstream hooks reading .grad would " + "see stale bytes from the previous step." + ) diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index beef6fe52..9325d87ae 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -6,6 +6,7 @@ from __future__ import annotations import functools +import math from importlib.metadata import PackageNotFoundError, version as get_pkg_version from typing import Optional @@ -88,6 +89,99 @@ def get_fp8_meta_from_fp8_tensor(tensor: Float8Tensor) -> tuple[FP8TensorMeta, i return fp8_meta, 0 +def get_main_grad_from_param( + weight_param: torch.nn.Parameter, + *, + op_label: str = "", +) -> torch.Tensor: + """Refresh ``main_grad`` from FSDP (if applicable) and return it. + Used by Megatron-LM-style wgrad fusion paths + (``accumulate_into_main_grad=True``) to obtain the buffer the wgrad GEMM + will write into. + Raises if the parameter does not have a ``main_grad`` attribute or if it + is ``None``. + """ + if hasattr(weight_param, "__fsdp_param__"): + weight_param.main_grad = weight_param.get_main_grad() + if not hasattr(weight_param, "main_grad") or weight_param.main_grad is None: + prefix = f"{op_label} " if op_label else "" + raise RuntimeError( + f"{prefix}operation is configured with accumulate_into_main_grad=True, " + "but weight parameter does not have a valid main_grad attribute" + ) + return weight_param.main_grad + + +def get_accumulate_flag_in_param(weight_param: torch.nn.Parameter) -> bool: + """Return whether the wgrad GEMM should accumulate into ``main_grad``. + + Returns ``False`` (i.e. overwrite) when the parameter has + ``overwrite_main_grad=True`` (used in Megatron-FSDP), and ``True`` + otherwise. + """ + return not getattr(weight_param, "overwrite_main_grad", False) + + +def view_main_grad_as_grouped_buffer( + main_grad: torch.Tensor, + num_groups: int, + weight_shape: tuple[int, ...], + *, + label: str = "", +) -> torch.Tensor: + """Return ``main_grad`` viewed as ``(num_groups, *weight_shape)`` without copy. + Raises if the numel doesn't match or if the existing stride pattern does + not allow a zero-copy view to the grouped layout. + """ + grouped_shape = (num_groups, *weight_shape) + if tuple(main_grad.shape) == grouped_shape: + return main_grad + prefix = f"{label} " if label else "Grouped weight " + if main_grad.numel() != math.prod(grouped_shape): + raise RuntimeError( + f"{prefix}main_grad expected shape {grouped_shape} or matching numel, " + f"but got shape {tuple(main_grad.shape)}" + ) + try: + return main_grad.view(grouped_shape) + except RuntimeError as e: + raise RuntimeError( + f"{prefix}main_grad must be viewable as {grouped_shape} without copy, " + f"but got shape {tuple(main_grad.shape)} and stride " + f"{tuple(main_grad.stride())}" + ) from e + + +def get_dummy_wgrads_for_params( + weight_params: list[torch.nn.Parameter], +) -> list[Optional[torch.Tensor]]: + """Build dummy ``.grad`` placeholders for Megatron-LM wgrad-fusion params. + + For each parameter that exposes ``grad_added_to_main_grad``, set the flag + to ``True`` and return a dummy wgrad tensor (zeroed if + ``zero_out_wgrad`` is also set on the parameter). For parameters without + the flag, the corresponding entry is ``None``. + + The returned list has the same length and order as ``weight_params``. + """ + from ..module.base import get_dummy_wgrad # pylint: disable=import-outside-toplevel + + out: list[Optional[torch.Tensor]] = [] + for wp in weight_params: + if hasattr(wp, "grad_added_to_main_grad"): + wp.grad_added_to_main_grad = True + out.append( + get_dummy_wgrad( + list(wp.size()), + wp.dtype, + zero=getattr(wp, "zero_out_wgrad", False), + ) + ) + else: + out.append(None) + return out + + def validate_grouped_mlp_dims(fc1, glu_op, fc2) -> None: """Validate FC1 / scaled GLU / FC2 dimensions for fused grouped MLP.""" diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index 46d52f7ff..41f0855f1 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -24,7 +24,6 @@ _2X_ACC_FPROP, _2X_ACC_DGRAD, _2X_ACC_WGRAD, - get_dummy_wgrad, ) from ...tensor import Quantizer from ...tensor.float8_tensor import Float8Quantizer @@ -36,7 +35,13 @@ devices_match, ) from ..op import BasicOperation, OperationContext -from .._common import maybe_dequantize, is_quantized_tensor +from .._common import ( + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, + is_quantized_tensor, + maybe_dequantize, +) def _wait_async(handle: Optional[Any]) -> None: @@ -1060,16 +1065,9 @@ def op_backward( grad_weight = None if ctx.weight_requires_grad and accumulate_into_main_grad: weight_param = self.weight - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - accumulate_into_main_grad = not getattr(weight_param, "overwrite_main_grad", False) - if not hasattr(weight_param, "main_grad"): - raise RuntimeError( - "BasicLinear op is configured with " - "accumulate_into_main_grad=True, " - "but weight parameter does not have main_grad attribute" - ) - grad_weight = weight_param.main_grad.detach() + main_grad = get_main_grad_from_param(weight_param, op_label="BasicLinear") + accumulate_into_main_grad = get_accumulate_flag_in_param(weight_param) + grad_weight = main_grad.detach() else: accumulate_into_main_grad = False @@ -1099,14 +1097,6 @@ def op_backward( # Megatron-LM wgrad fusion # Note: Return dummy tensor for grad weight if needed. if accumulate_into_main_grad: - grad_weight = None - weight_param = self.weight - if hasattr(weight_param, "grad_added_to_main_grad"): - weight_param.grad_added_to_main_grad = True - grad_weight = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) + grad_weight = get_dummy_wgrads_for_params([self.weight])[0] return grad_input, [grad_weight] diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index a86abb132..e698c2697 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -14,28 +14,36 @@ import torch import transformer_engine_torch as tex -from ...cpp_extensions import general_grouped_gemm +from ...cpp_extensions import general_grouped_gemm, general_grouped_gemm_for_grouped_tensor from ...distributed import CudaRNGStatesTracker from ...module._common import WeightGradStore from ...module.base import ( _2X_ACC_FPROP, _2X_ACC_DGRAD, _2X_ACC_WGRAD, - get_dummy_wgrad, ) from ...quantization import FP8GlobalStateManager, Recipe +from ...quantized_tensor import QuantizedTensorStorage from ...tensor import MXFP8Quantizer, MXFP8Tensor, Quantizer from ...utils import ( canonicalize_device, canonicalize_dtype, clear_tensor_data, devices_match, + get_device_compute_capability, resolve_grouped_linear_single_param_flags, round_up_to_nearest_multiple, ) -from .._common import is_quantized_tensor, maybe_dequantize +from .._common import ( + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, + is_quantized_tensor, + maybe_dequantize, + view_main_grad_as_grouped_buffer, +) from ..op import BasicOperation, OperationContext -from ...tensor import GroupedTensor +from ...tensor import GroupedTensor, GroupedTensorStorage from ...triton.grouped_dbias_dscales import ( compute_grouped_dbias, compute_grouped_dbias_dscales, @@ -242,7 +250,7 @@ def backward_dw(self) -> None: else: # Fused MXFP8 grouped MLP saves `GroupedTensor` activations for wgrad. clear_tensor_data( - activations.data, + activations.rowwise_data, activations.columnwise_data, activations.scale_inv, activations.columnwise_scale_inv, @@ -724,6 +732,151 @@ def op_backward(self, *args, **kwargs): "It overrides `fuser_backward` instead of `op_backward`." ) + @staticmethod + def _is_graph_safe_path_supported( + *, + with_quantized_compute: bool, + input_quantizers: Sequence[Optional[Quantizer]], + dtype: torch.dtype, + ) -> bool: + """Whether the graph-safe grouped-tensor flow can be used. + + * The graph-safe path dispatches to ``general_grouped_gemm_for_grouped_tensor``, + which is backed by ``nvte_grouped_gemm_with_discrete_inputA`` in the common + library. That kernel requires Blackwell (SM100) or newer with cuBLAS 13.3+. + * Quantized compute is currently MXFP8-only; every other quantization + recipe (fp8 delayed / current scaling, fp8 block scaling, NVFP4, ...) + falls back to the legacy flow. + * Unquantized compute supports BF16/FP16 only -- FP32 is excluded + because the cublasLt grouped GEMM doesn't support it. + """ + if get_device_compute_capability() < (10, 0): + return False + if with_quantized_compute: + return all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) + return dtype in (torch.bfloat16, torch.float16) + + def _get_grouped_weight_for_gemm( + self, + weight_param: GroupedTensor, + weight_quantizers: list[Optional[Quantizer]], + columnwise_usage: bool, + with_quantized_compute: bool, + dtype: torch.dtype, + ) -> GroupedTensor: + """Prepare weights for ``general_grouped_gemm_for_grouped_tensor``. + Supports MXFP8/BF16/FP16 compute paths. + """ + num_groups = self.num_groups + is_weight_quantized = weight_param.quantizer is not None + if is_weight_quantized and with_quantized_compute: + # GGEMM can use it as it is + return weight_param + if is_weight_quantized and not with_quantized_compute: + # This use-case isnt optimized yet. Involves a per-group + # dequantize loop and a torch.stack copy. + weight_parts = weight_param.quantized_tensors + if weight_parts is None: + weight_parts = weight_param.split_into_quantized_tensors() + dequantized = [maybe_dequantize(w, dtype) for w in weight_parts] + weight_data = torch.stack(dequantized, dim=0).contiguous() + return GroupedTensor( + shape=(num_groups * self.out_features, self.in_features), + dtype=dtype, + num_tensors=num_groups, + shapes=[(self.out_features, self.in_features)] * num_groups, + quantizer=None, + data=weight_data.reshape(-1), + ) + if not with_quantized_compute: + # Make sure that weight param is the correct dtype, + # otherwise cast it to the correct dtype. + if weight_param.rowwise_data.dtype == dtype: + return weight_param + weight_data = weight_param.rowwise_data.to(dtype=dtype) + return GroupedTensor( + shape=(num_groups * self.out_features, self.in_features), + dtype=dtype, + num_tensors=num_groups, + shapes=[(self.out_features, self.in_features)] * num_groups, + quantizer=None, + data=weight_data.reshape(-1), + ) + # Quantized compute path, use the fused group quantize kernel. + weight_quantizer = weight_quantizers[0] + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + return tex.group_quantize( + weight_param.rowwise_data.view(weight_param.logical_shape), + weight_quantizer, + num_groups, + None, + ) + + def _get_discrete_weights_for_gemm( + self, + weight_params: Optional[GroupedTensor] | list[torch.Tensor], + weight_quantizers: list[Optional[Quantizer]], + columnwise_usage: bool, + with_quantized_compute: bool, + dtype: torch.dtype, + ) -> list[torch.Tensor]: + """Prepare weights for ``general_grouped_gemm_for_grouped_tensor``. + Returns a Python list, which dispatches the GEMM to ``discrete_in`` mode. + """ + out: list[torch.Tensor] = [] + for w, quantizer in zip(weight_params, weight_quantizers): + if not with_quantized_compute: + w = maybe_dequantize(w, dtype) + elif with_quantized_compute and not is_quantized_tensor(w): + quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + w = quantizer(w) + out.append(w) + return out + + def _get_weight_tensors(self) -> list[torch.nn.Parameter]: + """Return the weight parameters in registration order. + + Length is 1 when ``single_grouped_weight=True`` (one + ``GroupedTensor`` parameter), otherwise ``num_groups``. + """ + if self.single_grouped_weight: + return [self.weight] + return [getattr(self, f"weight{idx}") for idx in range(self.num_groups)] + + def _get_grouped_bias_for_gemm( + self, + dtype: torch.dtype, + ) -> Optional[torch.Tensor]: + """Build a uniform GroupedTensor of per-group biases for the cublas + grouped GEMM. + + Each group expects a (1, out_features) bias vector. Returns ``None`` + when no additive bias is configured. + """ + if not self.has_bias: + return None + num_groups = self.num_groups + + if self.single_grouped_bias: + # Already a contiguous (num_groups * out_features) buffer. + bias_data = self.bias.rowwise_data + if bias_data.dtype != dtype: + bias_data = bias_data.to(dtype=dtype) + else: + bias_list = [ + maybe_dequantize(getattr(self, f"bias{idx}"), dtype) for idx in range(num_groups) + ] + bias_data = torch.stack(bias_list, dim=0).contiguous() + + return GroupedTensor( + shape=(num_groups, self.out_features), + dtype=dtype, + num_tensors=num_groups, + shapes=[(1, self.out_features)] * num_groups, + quantizer=None, + data=bias_data.reshape(-1), + ) + def fuser_forward( self, basic_op_ctxs: list[OperationContext], @@ -735,7 +888,6 @@ def fuser_forward( basic_op_kwargs: list[dict[str, Any]], ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: num_groups = self.num_groups - has_bias = self.has_bias weight_param = self.weight if self.single_grouped_weight else self.weight0 device = weight_param.device @@ -753,13 +905,11 @@ def fuser_forward( # Quantizers input_quantizers = [None] * num_groups weight_quantizers = [None] * num_groups - grad_output_quantizers = [None] * num_groups with_quantized_compute = FP8GlobalStateManager.is_fp8_enabled() if with_quantized_compute: for group_idx in range(num_groups): input_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx) weight_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx + 1) - grad_output_quantizers[group_idx] = self.get_quantizer("backward", group_idx) # Get autocast dtype if needed if torch.is_autocast_enabled(): @@ -767,17 +917,151 @@ def fuser_forward( else: dtype = weight_param.dtype - # Extract split sizes from extra input + # Extract split sizes from extra input. Keep on GPU for graph safety. split_sizes = basic_op_extra_inputs[0][0] - split_sizes_int = [int(s) for s in split_sizes.tolist()] - if len(split_sizes_int) != num_groups: - raise ValueError(f"Expected {num_groups} splits, but got {len(split_sizes_int)}.") + if int(split_sizes.numel()) != num_groups: + raise ValueError(f"Expected {num_groups} splits, but got {int(split_sizes.numel())}.") + if split_sizes.dtype != torch.int64: + split_sizes = split_sizes.to(dtype=torch.int64) + if split_sizes.device != device: + split_sizes = split_sizes.to(device=device) # Extract scales tensor for bias scaling scales = None if self._scale_bias: scales = basic_op_extra_inputs[0][1] + # Dispatch: graph-safe GroupedTensor flow whenever it can be used. + # See ``_is_graph_safe_path_supported`` for the gating rationale -- + # in short it requires Blackwell (SM100+) plus a supported dtype / + # quantization recipe. Otherwise we fall back to the legacy + # ``tex.split_quantize`` + ``general_grouped_gemm`` flow. + use_grouped_tensor_path = self._is_graph_safe_path_supported( + with_quantized_compute=with_quantized_compute, + input_quantizers=input_quantizers, + dtype=dtype, + ) + + if use_grouped_tensor_path: + out, tensors_to_save = self._fuser_forward_grouped_tensor( + input_=input_, + split_sizes=split_sizes, + scales=scales, + with_quantized_compute=with_quantized_compute, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + dtype=dtype, + input_requires_grad=input_requires_grad, + weight_requires_grad=weight_requires_grad, + device=device, + ) + else: + out, tensors_to_save = self._fuser_forward_split_quantize( + input_=input_, + split_sizes=split_sizes, + scales=scales, + with_quantized_compute=with_quantized_compute, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + dtype=dtype, + input_requires_grad=input_requires_grad, + weight_requires_grad=weight_requires_grad, + device=device, + ) + + # Save tensors and autograd metadata on the basic-op context. + self.fuser_forward_save_ctx( + basic_op_ctxs=basic_op_ctxs, + input_=input_, + tensors_to_save=[tensors_to_save], + requires_grad=[ctx.requires_grad], + basic_op_extra_inputs=basic_op_extra_inputs, + prev_op_grad_output_quantizer=prev_op_grad_output_quantizer, + next_op_input_quantizer=next_op_input_quantizer, + basic_op_kwargs=basic_op_kwargs, + use_grouped_tensor_path=use_grouped_tensor_path, + ) + + return out, [()] + + def fuser_forward_save_ctx( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, # pylint: disable=unused-argument + tensors_to_save: list[ + tuple[Optional[torch.Tensor | QuantizedTensorStorage | GroupedTensorStorage], ...] + ], + *, + requires_grad: list[bool], + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], # pylint: disable=unused-argument + prev_op_grad_output_quantizer: Optional[Quantizer], # pylint: disable=unused-argument + next_op_input_quantizer: Optional[Quantizer], # pylint: disable=unused-argument + basic_op_kwargs: list[dict[str, Any]], # pylint: disable=unused-argument + use_grouped_tensor_path: bool, + ) -> None: + """ + Save tensors and autograd metadata in context. + """ + if not requires_grad[0]: + return + + ctx = basic_op_ctxs[0] + ctx.save_for_backward(*tensors_to_save[0]) + + num_groups = self.num_groups + weight_param = self.weight if self.single_grouped_weight else self.weight0 + + with_quantized_compute = FP8GlobalStateManager.is_fp8_enabled() + input_quantizers = [None] * num_groups + weight_quantizers = [None] * num_groups + grad_output_quantizers = [None] * num_groups + if with_quantized_compute: + for group_idx in range(num_groups): + input_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx) + weight_quantizers[group_idx] = self.get_quantizer("forward", 2 * group_idx + 1) + grad_output_quantizers[group_idx] = self.get_quantizer("backward", group_idx) + + ctx.use_grouped_tensor_path = use_grouped_tensor_path + ctx.with_quantized_compute = with_quantized_compute + ctx.input_quantizers = input_quantizers + ctx.weight_quantizers = weight_quantizers + ctx.grad_output_quantizers = grad_output_quantizers + ctx.grad_input_quantizers = None + # ``split_sizes`` and ``base_split_offsets`` are routed through + # ``save_for_backward`` (see ``_fuser_forward_split_quantize`` and + # ``_fuser_forward_grouped_tensor`` for the saved-tensor layout). + if torch.is_autocast_enabled(): + ctx.dtype = torch.get_autocast_dtype("cuda") + else: + ctx.dtype = weight_param.dtype + ctx.input_requires_grad = requires_grad[0] + ctx.weight_requires_grad = requires_grad[0] and weight_param.requires_grad + + # ================================================================== + # Legacy `tex.split_quantize` + `general_grouped_gemm` flow. + # ``m_splits`` is needed on CPU here, so this flow is NOT cuda-graphable. + # ================================================================== + def _fuser_forward_split_quantize( + self, + *, + input_: torch.Tensor, + split_sizes: torch.Tensor, + scales: Optional[torch.Tensor], + with_quantized_compute: bool, + input_quantizers: list[Optional[Quantizer]], + weight_quantizers: list[Optional[Quantizer]], + dtype: torch.dtype, + input_requires_grad: bool, + weight_requires_grad: bool, + device: torch.device, + ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: + """Legacy ``tex.split_quantize`` + ``general_grouped_gemm`` flow.""" + num_groups = self.num_groups + has_bias = self.has_bias + + # Need CPU split sizes for split_quantize / general_grouped_gemm. + split_sizes_int = [int(s) for s in split_sizes.tolist()] + # Extract params if self.single_grouped_weight: weights = self.weight.quantized_tensors @@ -787,26 +1071,15 @@ def fuser_forward( weights = [getattr(self, f"weight{idx}") for idx in range(num_groups)] bs = None if has_bias: - if self.single_grouped_bias: - bias_parts = self.bias.quantized_tensors - if bias_parts is None: - bias_parts = self.bias.split_into_quantized_tensors() - bs = [maybe_dequantize(p.reshape(-1), dtype) for p in bias_parts] - else: - bs = [ - maybe_dequantize(getattr(self, f"bias{idx}"), dtype) - for idx in range(num_groups) - ] - - # Convert weight dtype if needed - ws = [] - for w, quantizer in zip(weights, weight_quantizers): - if not with_quantized_compute: - w = maybe_dequantize(w, dtype) - elif with_quantized_compute and not is_quantized_tensor(w): - quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - w = quantizer(w) - ws.append(w) + bs = self._get_bias_tensors(dtype) + + ws = self._get_discrete_weights_for_gemm( + weights, + weight_quantizers, + columnwise_usage=input_requires_grad, + with_quantized_compute=with_quantized_compute, + dtype=dtype, + ) # Split input tensor and convert dtypes if needed x = maybe_dequantize(input_, dtype) @@ -839,9 +1112,6 @@ def fuser_forward( ) # Add bias * scales when scale_bias is enabled - # TODO(vthumbe): Need to use GroupedBiasAdd kernel here. - # Would be done as part of larger refactor for GroupedLinear + GroupedTensor - # integration. if self._scale_bias and has_bias: scales_splits = torch.split(scales, split_sizes_int) out_splits = torch.split(out, split_sizes_int) @@ -863,24 +1133,147 @@ def fuser_forward( for x in xs: x.update_usage(rowwise_usage=False, columnwise_usage=True) - # Save state for backward pass - if ctx.requires_grad: - saved = [split_sizes] + # Build the tuple of tensors to save for backward. Layout: + # [split_sizes, base_split_offsets, split_points, + # (scales if scale_bias), *xs, *ws] + # ``base_split_offsets`` and ``split_points`` are unused on the + # split-quantize backward path but are included as ``None`` so the + # saved-tensor layout matches the graph-safe + # ``_fuser_forward_grouped_tensor`` path (and the fused MLP forward). + saved: list[Optional[torch.Tensor]] = [split_sizes, None, None] + if self._scale_bias: + saved.append(scales) + saved.extend(xs) + saved.extend(ws) + return out, tuple(saved) + + def _fuser_forward_grouped_tensor( + self, + *, + input_: torch.Tensor, + split_sizes: torch.Tensor, + scales: Optional[torch.Tensor], + with_quantized_compute: bool, + input_quantizers: list[Optional[Quantizer]], + weight_quantizers: list[Optional[Quantizer]], + dtype: torch.dtype, + input_requires_grad: bool, + weight_requires_grad: bool, + device: torch.device, + ) -> tuple[torch.Tensor, tuple[Optional[torch.Tensor], ...]]: + """Graph-safe GroupedTensor forward path (pure compute). + Returns ``(output, tensors_to_save)``. ``split_sizes``, + ``base_split_offsets`` and ``split_points`` are returned so that + ``fuser_forward_save_ctx`` can call ``save_for_backward`` on them. + """ + num_groups = self.num_groups + has_bias = self.has_bias + + base_split_offsets = tex.splits_to_offsets(split_sizes, 1) + split_points = base_split_offsets[1:].to(dtype=torch.int) + + # Flatten to 2D so the first dim is the total token count. + original_shape = list(input_.size()) + x = maybe_dequantize(input_, dtype).reshape(-1, self.in_features) + total_tokens = x.size(0) + + # Build the input GroupedTensor. + if with_quantized_compute: + input_quantizer = input_quantizers[0] + input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + input_quantizer.optimize_for_gemm = True + grouped_x = tex.group_quantize(x, input_quantizer, num_groups, split_sizes) + else: + # No quantize: wrap the contiguous high-precision buffer. + grouped_x = GroupedTensor( + shape=(total_tokens, self.in_features), + dtype=dtype, + num_tensors=num_groups, + quantizer=None, + data=x.reshape(-1), + first_dims=split_sizes, + tensor_offsets=base_split_offsets * self.in_features, + ) + + # Build the weight GroupedTensor / list. + if self.single_grouped_weight: + # GroupedTensor + grouped_weights = self._get_grouped_weight_for_gemm( + self.weight, + weight_quantizers, + columnwise_usage=input_requires_grad, + with_quantized_compute=with_quantized_compute, + dtype=dtype, + ) + else: + # Discrete weights + grouped_weights = self._get_discrete_weights_for_gemm( + [getattr(self, f"weight{idx}") for idx in range(num_groups)], + weight_quantizers, + columnwise_usage=input_requires_grad, + with_quantized_compute=with_quantized_compute, + dtype=dtype, + ) + + # Allocate output buffer and wrap as a GroupedTensor view. + out_shape = original_shape[:-1] + [self.out_features] + out = torch.empty(out_shape, dtype=dtype, device=device) + grouped_out = GroupedTensor( + shape=(total_tokens, self.out_features), + dtype=dtype, + num_tensors=num_groups, + quantizer=None, + data=out.reshape(-1), + first_dims=split_sizes, + tensor_offsets=base_split_offsets * self.out_features, + ) + + # Bias: hand off to the grouped GEMM (graph-safe, fused). Plain bias + # uses ``bias=``; scaled bias also passes per-token ``bias_scale=``. + grouped_bias = None + bias_scale: Optional[torch.Tensor] = None + if has_bias: + # Bias always needs to be passed as a GroupedTensor for the grouped GEMM. + grouped_bias = self._get_grouped_bias_for_gemm(dtype) if self._scale_bias: - saved.append(scales) - saved.extend(xs) - saved.extend(ws) - ctx.save_for_backward(*saved) - ctx.with_quantized_compute = with_quantized_compute - ctx.input_quantizers = input_quantizers - ctx.weight_quantizers = weight_quantizers - ctx.grad_output_quantizers = grad_output_quantizers - ctx.grad_input_quantizers = None - ctx.dtype = dtype - ctx.input_requires_grad = input_requires_grad - ctx.weight_requires_grad = weight_requires_grad + bias_scale = scales.reshape(-1) + if bias_scale.dtype != torch.float32: + bias_scale = bias_scale.to(dtype=torch.float32) + + # Forward grouped GEMM. + general_grouped_gemm_for_grouped_tensor( + grouped_weights, + grouped_x, + grouped_out, + layout="TN", + use_split_accumulator=_2X_ACC_FPROP, + bias=grouped_bias, + bias_scale=bias_scale, + ) - return out, [()] + if not input_requires_grad: + grouped_weights = None if self.single_grouped_weight else [None] * num_groups + + if not weight_requires_grad: + grouped_x = None + + # Build the tuple of tensors to save for backward. Layout: + # [split_sizes, base_split_offsets, split_points, + # (scales if _scale_bias), grouped_x, *weights] + if grouped_x is not None: + if with_quantized_compute: + # only columnwise data is needed for wgrad + grouped_x.rowwise_data = None + grouped_x.scale_inv = None + saved: list[Optional[torch.Tensor]] = [split_sizes, base_split_offsets, split_points] + if self._scale_bias: + saved.append(scales) + saved.append(grouped_x) + if self.single_grouped_weight: + saved.append(grouped_weights) + else: + saved.extend(grouped_weights) + return out, tuple(saved) def fuser_backward( self, @@ -892,16 +1285,43 @@ def fuser_backward( torch.Tensor, Iterable[Iterable[Optional[torch.Tensor]]], Iterable[Iterable[Optional[torch.Tensor]]], + ]: + ctx = basic_op_ctxs[0] + # Dispatch to the path used in forward (saved as ``ctx.use_grouped_tensor_path``). + if getattr(ctx, "use_grouped_tensor_path", False): + return self._fuser_backward_grouped_tensor( + ctx=ctx, + grad_output=grad_output, + ) + return self._fuser_backward_split_quantize( + ctx=ctx, + grad_output=grad_output, + ) + + def _fuser_backward_split_quantize( + self, + *, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], ]: num_groups = self.num_groups has_bias = self.has_bias - weight_param = self.weight if self.single_grouped_weight else self.weight0 - device = weight_param.device + weights = self._get_weight_tensors() + device = weights[0].device - # Saved tensors from forward pass - ctx = basic_op_ctxs[0] + # Saved tensors from forward pass. Layout: + # [split_sizes, base_split_offsets, split_points, + # (scales if _scale_bias), *xs, *ws] + # ``base_split_offsets`` and ``split_points`` are unused on this path + # but are present so the saved-tensor layout matches the graph-safe + # path (and the fused MLP forward). saved_tensors = ctx.saved_tensors - split_sizes, saved_tensors = saved_tensors[0], saved_tensors[1:] + split_sizes = saved_tensors[0] + saved_tensors = saved_tensors[3:] scales = None if self._scale_bias: scales, saved_tensors = saved_tensors[0], saved_tensors[1:] @@ -941,58 +1361,43 @@ def fuser_backward( dbias_packed = compute_grouped_dbias(dy_2d, offsets, num_groups) grad_biases = [dbias_packed[idx].to(dtype=ctx.dtype) for idx in range(num_groups)] - # Initialize grad weight buffers + # Initialize grad weight buffers. accumulate_into_main_grad = self._accumulate_into_main_grad grad_weights = [None] * num_groups + final_weight_grads: list[Optional[torch.Tensor]] = ( + [None] if self.single_grouped_weight else [None] * num_groups + ) if ctx.weight_requires_grad: - if accumulate_into_main_grad: - # Megatron-LM wgrad fusion - # Note: Get grad tensors from params so we can - # accumulate directly into it. - if self.single_grouped_weight: - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - main_grad = weight_param.main_grad - if isinstance(main_grad, GroupedTensor): - grad_weights = main_grad.quantized_tensors - if grad_weights is None: - grad_weights = main_grad.split_into_quantized_tensors() - else: - # main_grad may be [num_groups, out, in] or a flat buffer. - # Canonicalize to grouped layout before slicing per-group views. - weight_shape = (self.out_features, self.in_features) - grouped_shape = (num_groups, *weight_shape) - if main_grad.shape != grouped_shape: - if main_grad.numel() != math.prod(grouped_shape): - raise RuntimeError( - "GroupedLinear expected grouped weight main_grad to have " - f"shape {grouped_shape} or matching numel, " - f"but got shape {tuple(main_grad.shape)}" - ) - main_grad = main_grad.reshape(grouped_shape) - grad_weights = [main_grad[idx] for idx in range(num_groups)] - accumulate_into_main_grad = not getattr( - weight_param, "overwrite_main_grad", False + weight_shape = (self.out_features, self.in_features) + grouped_shape = (num_groups, *weight_shape) + if self.single_grouped_weight: + if accumulate_into_main_grad: + # Megatron-LM wgrad fusion: GEMM accumulates into the + # parameter's ``main_grad`` directly. + main_grad = get_main_grad_from_param(weights[0], op_label="GroupedLinear") + main_grad = view_main_grad_as_grouped_buffer( + main_grad, num_groups, weight_shape, label="GroupedLinear weight" ) + final_weight_grads[0] = main_grad + grad_weights = [main_grad[idx] for idx in range(num_groups)] + accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: - for group_idx in range(num_groups): - weight_param = getattr(self, f"weight{group_idx}") - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - grad_weights[group_idx] = weight_param.main_grad - accumulate_into_main_grad = not getattr( - self.weight0, "overwrite_main_grad", False + final_weight_grads[0] = torch.empty( + grouped_shape, dtype=ctx.dtype, device=device ) + grad_weights = [final_weight_grads[0][idx] for idx in range(num_groups)] else: - weight_shape = (self.out_features, self.in_features) - for group_idx in range(num_groups): - grad_weights[group_idx] = torch.empty( - weight_shape, - dtype=ctx.dtype, - device=device, - ) - else: - accumulate_into_main_grad = False + if accumulate_into_main_grad: + grad_weights = [ + get_main_grad_from_param(w, op_label="GroupedLinear") for w in weights + ] + accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) + else: + grad_weights = [ + torch.empty(weight_shape, dtype=ctx.dtype, device=device) + for _ in range(num_groups) + ] + final_weight_grads = list(grad_weights) # Perform dgrad GEMMs grad_input = None @@ -1050,54 +1455,14 @@ def fuser_backward( if not delay_wgrad: clear_tensor_data(*xs) - # Megatron-LM wgrad fusion - # Note: Return dummy tensor for grad weight if needed. - if accumulate_into_main_grad: - grad_weights = [None] * num_groups - if self.single_grouped_weight: - if hasattr(weight_param, "grad_added_to_main_grad"): - weight_param.grad_added_to_main_grad = True - grad_weight = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) - else: - grad_weight = None - # Be mindful of param registration order. - if has_bias: - if self.single_grouped_bias: - final_bias_grads = torch.stack(grad_biases, dim=0).to(ctx.dtype) - grad_params = [grad_weight, final_bias_grads] - else: - grad_params = grad_biases + [grad_weight] - else: - grad_params = [grad_weight] - grad_extra = (None, grad_scales) if self._scale_bias else (None,) - return grad_input, [grad_params], [grad_extra] - for group_idx in range(num_groups): - weight_param = getattr(self, f"weight{group_idx}") - if hasattr(weight_param, "grad_added_to_main_grad"): - weight_param.grad_added_to_main_grad = True - grad_weights[group_idx] = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) - - if self.single_grouped_weight: - grad_weight = None - if ctx.weight_requires_grad: - if delay_wgrad: - grad_weight = None - else: - grad_weight = torch.stack(grad_weights, dim=0) - final_weight_grads = [grad_weight] - else: - if delay_wgrad and ctx.weight_requires_grad and not accumulate_into_main_grad: - final_weight_grads = [None] * num_groups - else: - final_weight_grads = grad_weights + # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, + # signal that ``main_grad`` already carries the wgrad and replace + # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into + # ``main_grad`` again. + if ctx.weight_requires_grad and self._accumulate_into_main_grad: + final_weight_grads = get_dummy_wgrads_for_params(weights) + elif ctx.weight_requires_grad and delay_wgrad: + final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups if not has_bias: grad_params = list(final_weight_grads) @@ -1112,3 +1477,214 @@ def fuser_backward( grad_extra = (None, grad_scales) if self._scale_bias else (None,) return grad_input, [grad_params], [grad_extra] + + # ================================================================== + # Graph-safe backward: counterpart of `_fuser_forward_grouped_tensor`. + # ================================================================== + def _fuser_backward_grouped_tensor( + self, + *, + ctx: OperationContext, + grad_output: torch.Tensor, + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + num_groups = self.num_groups + has_bias = self.has_bias + weights = self._get_weight_tensors() + device = weights[0].device + dtype = ctx.dtype + + with_quantized_compute = bool(getattr(ctx, "with_quantized_compute", False)) + + # Saved tensors from forward pass + # Layout: [split_sizes, base_split_offsets, split_points, + # (scales if _scale_bias), grouped_x, *weights] + # ``split_points`` is unused on this path but is present so the + # saved-tensor layout matches the fused MLP forward (which needs it + # for the cuDNN grouped GEMM kernel). + saved_tensors = ctx.saved_tensors + split_sizes = saved_tensors[0] + base_split_offsets = saved_tensors[1] + saved_tensors = saved_tensors[3:] + scales = None + if self._scale_bias: + scales, saved_tensors = saved_tensors[0], saved_tensors[1:] + grouped_x, saved_tensors = saved_tensors[0], saved_tensors[1:] + if self.single_grouped_weight: + ws, saved_tensors = saved_tensors[0], saved_tensors[1:] + else: + ws, saved_tensors = saved_tensors[:num_groups], saved_tensors[num_groups:] + + # Flatten grad_output to 2D (total_tokens, out_features) + # to figure out total tokens. + dy_2d = grad_output.reshape(-1, self.out_features) + total_tokens = dy_2d.size(0) + + # Build the grad_output GroupedTensor. + # Optionally get dbias is fusion available with bgrad_group_quantize + dbias_packed = None + if with_quantized_compute: + grad_output_quantizer = ctx.grad_output_quantizers[0] + grad_output_quantizer.set_usage( + rowwise=ctx.input_requires_grad, columnwise=ctx.weight_requires_grad + ) + grad_output_quantizer.optimize_for_gemm = True + + if has_bias and not self._scale_bias: + grouped_dy, dbias_packed = tex.bgrad_group_quantize( + dy_2d, grad_output_quantizer, num_groups, split_sizes + ) + else: + grouped_dy = tex.group_quantize( + dy_2d, grad_output_quantizer, num_groups, split_sizes + ) + else: + dy_2d = maybe_dequantize(dy_2d, dtype) + # Wrap BF16/FP16 buffer as a GroupedTensor for grouped gemm + grouped_dy = GroupedTensor( + shape=(total_tokens, self.out_features), + dtype=dtype, + num_tensors=num_groups, + quantizer=None, + data=dy_2d.reshape(-1), + first_dims=split_sizes, + tensor_offsets=base_split_offsets * self.out_features, + ) + + # Bias Grads compute if not already computed in bgrad_group_quantize + final_bias_grads: Optional[torch.Tensor] = None + grad_scales: Optional[torch.Tensor] = None + if has_bias: + if self._scale_bias: + bias_packed = torch.stack(self._get_bias_tensors(dtype)) + scales_f32 = scales.to(dtype=torch.float32) + dbias_packed, grad_scales = compute_grouped_dbias_dscales( + dy_2d, + scales_f32, + bias_packed, + offsets=base_split_offsets, + ) + elif dbias_packed is None: + # BF16/FP16 path + dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, num_groups) + if self.single_grouped_bias: + final_bias_grads = [dbias_packed.to(dtype=dtype)] + else: + final_bias_grads = [dbias_packed[idx].to(dtype=dtype) for idx in range(num_groups)] + + # ---- dgrad GEMM ---------------------------------------------------- + grad_input = None + if ctx.input_requires_grad: + grad_input_shape = list(grad_output.size())[:-1] + [self.in_features] + grad_input = torch.empty(grad_input_shape, dtype=dtype, device=device) + grouped_grad_input = GroupedTensor( + shape=(total_tokens, self.in_features), + dtype=dtype, + num_tensors=num_groups, + quantizer=None, + data=grad_input.reshape(-1), + first_dims=split_sizes, + tensor_offsets=base_split_offsets * self.in_features, + ) + general_grouped_gemm_for_grouped_tensor( + ws, + grouped_dy, + grouped_grad_input, + layout="NN", + use_split_accumulator=_2X_ACC_DGRAD, + ) + + # params init for wgrad GEMM + accumulate_into_main_grad = False + weight_shape = (self.out_features, self.in_features) + wgrad_output: Any = None + grouped_wgrad: Optional[GroupedTensor] = None + final_weight_grads: list[Optional[torch.Tensor]] = ( + [None] if self.single_grouped_weight else [None] * num_groups + ) + + # Get the right wgrad buffers for grouped gemm. + # Can be a GroupedTensor or list of tensors based on single_grouped_weight. + if ctx.weight_requires_grad: + if self.single_grouped_weight: + if self._accumulate_into_main_grad: + # Main-grad fusion: GEMM writes directly into ``main_grad``. + # ``overwrite_main_grad`` only flips the GEMM's + # ``accumulate`` flag. + main_grad = get_main_grad_from_param(weights[0], op_label="GroupedLinear") + main_grad = view_main_grad_as_grouped_buffer( + main_grad, num_groups, weight_shape, label="GroupedLinear weight" + ) + grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( + num_tensors=num_groups, + tensor_shape=weight_shape, + rowwise_data=main_grad.view(-1), + dtype=main_grad.dtype, + ) + accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) + else: + grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes( + num_tensors=num_groups, + shapes=[weight_shape] * num_groups, + quantizer=None, + device=device, + dtype=dtype, + ) + final_weight_grads[0] = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) + wgrad_output = grouped_wgrad + else: + if self._accumulate_into_main_grad: + final_weight_grads = [ + get_main_grad_from_param(w, op_label="GroupedLinear") for w in weights + ] + accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) + else: + final_weight_grads = [ + torch.empty(weight_shape, dtype=dtype, device=device) + for _ in range(num_groups) + ] + wgrad_output = final_weight_grads + + # wgrad GEMM + delay_wgrad = ( + ctx.weight_requires_grad + and self.wgrad_store is not None + and self.wgrad_store.delay_wgrad_compute() + ) + if ctx.weight_requires_grad: + wgrad_gemm = functools.partial( + general_grouped_gemm_for_grouped_tensor, + layout="NT", + accumulate=accumulate_into_main_grad, + use_split_accumulator=_2X_ACC_WGRAD, + ) + if delay_wgrad: + self.wgrad_store.put([grouped_x, grouped_dy, wgrad_output], wgrad_gemm) + else: + wgrad_gemm(grouped_x, grouped_dy, wgrad_output) + + # Megatron-LM wgrad fusion: regardless of overwrite vs. accumulate, + # signal that ``main_grad`` already carries the wgrad and replace + # ``.grad`` with a dummy so DDP/FSDP hooks won't add ``.grad`` into + # ``main_grad`` again. + if ctx.weight_requires_grad and self._accumulate_into_main_grad: + final_weight_grads = get_dummy_wgrads_for_params(weights) + elif ctx.weight_requires_grad and delay_wgrad: + final_weight_grads = [None] if self.single_grouped_weight else [None] * num_groups + + # Assemble grad params in parameter registration order and return. + if not has_bias: + grad_params = final_weight_grads + elif self.single_grouped_bias: + grad_params = final_weight_grads + final_bias_grads + else: + if self.single_grouped_weight: + grad_params = final_bias_grads + final_weight_grads + else: + grad_params = final_weight_grads + final_bias_grads + + grad_extra = (None, grad_scales) if self._scale_bias else (None,) + return grad_input, [grad_params], [grad_extra] diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index b07ebb73e..320c7c39e 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -7,7 +7,6 @@ from __future__ import annotations from collections.abc import Callable import functools -import math import os from typing import Optional @@ -25,11 +24,15 @@ from .._common import ( _cudnn_frontend_version_supported, fuse_grouped_mlp_ops, + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, maybe_dequantize, + view_main_grad_as_grouped_buffer, validate_grouped_mlp_dims, ) from ...cpp_extensions import general_grouped_gemm_for_grouped_tensor -from ...module.base import _2X_ACC_WGRAD, get_dummy_wgrad +from ...module.base import _2X_ACC_WGRAD from ...triton.grouped_dbias_dscales import compute_grouped_dbias_dscales @@ -149,44 +152,32 @@ def _compute_grad_params( Returns the grad_params list in parameter registration order. """ - # Allocate grad buffers, determine accumulate flag + # Allocate grad buffers, determine accumulate flag. accumulate_into_main_grad = False grouped_wgrad = None wgrad_output = None + op_label = f"Grouped MLP fused backward ({label})" if label else "Grouped MLP fused backward" + weights = fc_op._get_weight_tensors() if fc_op.single_grouped_weight: w_list = [None] if ctx.weight_requires_grad: - weight_param = fc_op.weight if fc_op._accumulate_into_main_grad: - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - main_grad = weight_param.main_grad - grouped_shape = (num_groups, *weight_shape) - if main_grad.shape != grouped_shape: - if main_grad.numel() != math.prod(grouped_shape): - raise RuntimeError( - f"Grouped MLP fused backward expected {label} main_grad to have " - f"shape {grouped_shape} or matching numel, " - f"but got shape {tuple(main_grad.shape)}" - ) - try: - main_grad = main_grad.view(grouped_shape) - except RuntimeError as e: - raise RuntimeError( - f"Grouped MLP fused backward requires {label} main_grad to be " - f"viewable as {grouped_shape} without copy, but got shape" - f" {tuple(main_grad.shape)} and stride" - f" {tuple(main_grad.stride())}" - ) from e - accumulate_into_main_grad = not getattr(weight_param, "overwrite_main_grad", False) + # Main-grad fusion: GEMM writes directly into ``main_grad``. + # ``overwrite_main_grad`` only flips the GEMM's ``accumulate`` + # flag (overwrite vs. accumulate); it does not change the + # output buffer. + main_grad = get_main_grad_from_param(weights[0], op_label=op_label) + main_grad = view_main_grad_as_grouped_buffer( + main_grad, num_groups, weight_shape, label=f"{op_label} weight" + ) grouped_wgrad = GroupedTensor.make_grouped_tensor_from_rowwise_data( num_tensors=num_groups, tensor_shape=weight_shape, rowwise_data=main_grad, dtype=main_grad.dtype, ) - - if grouped_wgrad is None: + accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) + else: grouped_wgrad = GroupedTensor.make_grouped_tensor_with_shapes( num_tensors=num_groups, shapes=[weight_shape] * num_groups, @@ -195,19 +186,17 @@ def _compute_grad_params( dtype=dtype, ) wgrad_output = grouped_wgrad + w_list = [grouped_wgrad.rowwise_data.view(num_groups, *weight_shape)] else: w_list = [None] * num_groups if ctx.weight_requires_grad: if fc_op._accumulate_into_main_grad: - for idx in range(num_groups): - wp = getattr(fc_op, f"weight{idx}") - if hasattr(wp, "__fsdp_param__"): - wp.main_grad = wp.get_main_grad() - w_list[idx] = wp.main_grad - accumulate_into_main_grad = not getattr(fc_op.weight0, "overwrite_main_grad", False) + w_list = [get_main_grad_from_param(w, op_label=op_label) for w in weights] + accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: - for idx in range(num_groups): - w_list[idx] = torch.empty(weight_shape, dtype=dtype, device=device) + w_list = [ + torch.empty(weight_shape, dtype=dtype, device=device) for _ in range(num_groups) + ] wgrad_output = w_list if ctx.weight_requires_grad: @@ -237,34 +226,11 @@ def _compute_grad_params( else: gemm_fn(grouped_x, grouped_dy, wgrad_output) - # Extract results, mark accumulated if needed - if fc_op.single_grouped_weight: - packed_wgrad = None - if not delay_wgrad: - packed_wgrad = grouped_wgrad.rowwise_data.view(num_groups, *weight_shape) - if fc_op._accumulate_into_main_grad and hasattr( - weight_param, "grad_added_to_main_grad" - ): - weight_param.grad_added_to_main_grad = True - packed_wgrad = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) - w_list = [packed_wgrad] - else: - if delay_wgrad or fc_op._accumulate_into_main_grad: - w_list = [None] * num_groups - if fc_op._accumulate_into_main_grad: - for idx in range(num_groups): - wp = getattr(fc_op, f"weight{idx}") - if hasattr(wp, "grad_added_to_main_grad"): - wp.grad_added_to_main_grad = True - w_list[idx] = get_dummy_wgrad( - list(wp.size()), - wp.dtype, - zero=getattr(wp, "zero_out_wgrad", False), - ) + # Need to return dummy wgrads for Megatron-LM wgrad fusion if grad is already added + if fc_op._accumulate_into_main_grad: + w_list = get_dummy_wgrads_for_params(weights) + elif delay_wgrad: + w_list = [None] if fc_op.single_grouped_weight else [None] * num_groups # Assemble grad_params in parameter registration order. if not fc_op.has_bias: @@ -372,18 +338,15 @@ def fuser_backward( grad_output = grad_output.reshape(-1, fc2_weight_shape[0]) out_shape = list(grad_output.size()) num_groups = fc1_op.num_groups - fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 - device = fc1_weight_param.device + device = fc1_op._get_weight_tensors()[0].device dtype = fc1_ctx.dtype - # Saved tensors from FC1 forward + # Saved tensors from FC1 forward. + # Layout: [split_sizes, base_split_offsets, split_points, + # grouped_fc1_x, *fc1_weights] saved_tensors = fc1_ctx.saved_tensors - split_sizes, split_points, saved_tensors = ( - saved_tensors[0], - saved_tensors[1], - saved_tensors[2:], - ) - + split_sizes, base_split_offsets, split_points = saved_tensors[:3] + grouped_fc1_x, saved_tensors = saved_tensors[3], saved_tensors[4:] if fc1_op.single_grouped_weight: grouped_fc1_weight, saved_tensors = saved_tensors[0], saved_tensors[1:] else: @@ -392,21 +355,21 @@ def fuser_backward( saved_tensors[num_groups:], ) - ( - fc1_x_col_data, - fc1_x_col_scale, - fc1_x_tensor_offsets, - ), saved_tensors = ( - saved_tensors[:3], - saved_tensors[3:], - ) - # Saved tensors from scaled SwiGLU forward swiglu_in, scales = swiglu_ctx.saved_tensors - # Saved tensors from FC2 forward - saved_tensors = fc2_ctx.saved_tensors - _, saved_tensors = saved_tensors[0], saved_tensors[1:] # Assume same split sizes as FC1 + # Saved tensors from FC2 forward. + # Layout: [split_sizes, base_split_offsets, split_points, + # (fc2_scales if _scale_bias), + # grouped_fc2_x, *fc2_weights] + scale_bias = fc2_op._scale_bias and fc2_op.has_bias + saved_tensors = fc2_ctx.saved_tensors[3:] + if fc2_op._scale_bias: + # Saved for the unfused backward path, which reads its own + # per-op scales here. The fused backward below currently reuses + # the SwiGLU ``scales``. + saved_tensors = saved_tensors[1:] + grouped_fc2_x, saved_tensors = saved_tensors[0], saved_tensors[1:] if fc2_op.single_grouped_weight: grouped_fc2_weight, saved_tensors = saved_tensors[0], saved_tensors[1:] else: @@ -415,53 +378,19 @@ def fuser_backward( saved_tensors[num_groups:], ) - ( - fc2_x_col_data, - fc2_x_col_scale, - fc2_x_tensor_offsets, - ), saved_tensors = ( - saved_tensors[:3], - saved_tensors[3:], - ) - # Group splits if int(split_sizes.numel()) != num_groups: raise ValueError(f"Expected {num_groups} splits, but got {int(split_sizes.numel())}.") - scale_bias = fc2_op._scale_bias and fc2_op.has_bias - grouped_fc1_x = None - if fc1_ctx.weight_requires_grad: - grouped_fc1_x = GroupedTensor( - shape=(out_shape[0], fc1_weight_shape[1]), - dtype=dtype, - num_tensors=num_groups, - quantizer=fc1_ctx.input_quantizer, - columnwise_data=fc1_x_col_data, - columnwise_scale_inv=fc1_x_col_scale, - first_dims=split_sizes, - tensor_offsets=fc1_x_tensor_offsets, - with_gemm_swizzled_scales=True, - ) - - grouped_fc2_x = None - if fc2_ctx.weight_requires_grad: - grouped_fc2_x = GroupedTensor( - shape=(out_shape[0], fc2_weight_shape[1]), - dtype=dtype, - num_tensors=num_groups, - quantizer=fc2_ctx.input_quantizer, - columnwise_data=fc2_x_col_data, - columnwise_scale_inv=fc2_x_col_scale, - first_dims=split_sizes, - tensor_offsets=fc2_x_tensor_offsets, - with_gemm_swizzled_scales=True, - ) + if not fc1_ctx.weight_requires_grad: + grouped_fc1_x = None + if not fc2_ctx.weight_requires_grad: + grouped_fc2_x = None # Split grad output tensor and convert dtypes if needed - fc2_ctx.grad_output_quantizer.set_usage( - rowwise=True, columnwise=fc2_ctx.weight_requires_grad - ) - fc2_ctx.grad_output_quantizer.optimize_for_gemm = True + fc2_grad_output_quantizer = fc2_ctx.grad_output_quantizers[0] + fc2_grad_output_quantizer.set_usage(rowwise=True, columnwise=fc2_ctx.weight_requires_grad) + fc2_grad_output_quantizer.optimize_for_gemm = True output_fc2_dbias = fc2_op.has_bias fc2_dbias_packed = None fc2_dy = None @@ -476,14 +405,14 @@ def fuser_backward( if output_fc2_dbias and not scale_bias: grouped_fc2_dy, fc2_dbias_packed = tex.bgrad_group_quantize( fc2_dy, - fc2_ctx.grad_output_quantizer, + fc2_grad_output_quantizer, num_groups, split_sizes, ) else: grouped_fc2_dy = tex.group_quantize( fc2_dy, - fc2_ctx.grad_output_quantizer, + fc2_grad_output_quantizer, num_groups, split_sizes, ) @@ -600,7 +529,7 @@ def fuser_backward( fc2_dy, scales_f32, bias_packed, - offsets=fc1_ctx.base_split_offsets, + offsets=base_split_offsets, dscales=grad_scales, ) fc2_dbias_packed_result = fc2_dbias_packed_result.to(dtype=dtype) @@ -629,12 +558,12 @@ def fuser_backward( fc1_bias_grads = [dbias_2d[group_idx] for group_idx in range(num_groups)] # FC1 grad output for dgrad and wgrad GEMMs - fc1_dy_tensor_offsets = fc1_ctx.base_split_offsets * fc1_weight_shape[0] + fc1_dy_tensor_offsets = base_split_offsets * fc1_weight_shape[0] grouped_fc1_dy = GroupedTensor( shape=(out_shape[0], fc1_weight_shape[0]), dtype=dtype, num_tensors=num_groups, - quantizer=fc1_ctx.grad_output_quantizer, + quantizer=fc1_ctx.grad_output_quantizers[0], data=fc1_dy_row_data, columnwise_data=fc1_dy_col_data, scale_inv=fc1_dy_row_scale, @@ -668,7 +597,7 @@ def fuser_backward( and fc2_op.wgrad_store.delay_wgrad_compute() ): clear_tensor_data( - grouped_fc2_x.data, + grouped_fc2_x.rowwise_data, grouped_fc2_x.columnwise_data, grouped_fc2_x.scale_inv, grouped_fc2_x.columnwise_scale_inv, @@ -764,7 +693,7 @@ def fuser_backward( and fc1_op.wgrad_store.delay_wgrad_compute() ): clear_tensor_data( - grouped_fc1_x.data, + grouped_fc1_x.rowwise_data, grouped_fc1_x.columnwise_data, grouped_fc1_x.scale_inv, grouped_fc1_x.columnwise_scale_inv, diff --git a/transformer_engine/pytorch/ops/fused/backward_linear_add.py b/transformer_engine/pytorch/ops/fused/backward_linear_add.py index c06e212e8..382fecfd0 100644 --- a/transformer_engine/pytorch/ops/fused/backward_linear_add.py +++ b/transformer_engine/pytorch/ops/fused/backward_linear_add.py @@ -9,9 +9,13 @@ import torch -from ...module.base import get_dummy_wgrad from ...utils import clear_tensor_data from ..basic import BasicLinear, MakeExtraOutput +from .._common import ( + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, +) from ..op import FusedOperation, FusibleOperation, OperationContext @@ -57,16 +61,9 @@ def fuser_backward( grad_weight = None if linear_op_ctx.weight_requires_grad and accumulate_into_main_grad: weight_param = linear_op.weight - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - accumulate_into_main_grad = not getattr(weight_param, "overwrite_main_grad", False) - if not hasattr(weight_param, "main_grad"): - raise RuntimeError( - "BasicLinear op is configured with " - "accumulate_into_main_grad=True, " - "but weight parameter does not have main_grad attribute" - ) - grad_weight = weight_param.main_grad.detach() + main_grad = get_main_grad_from_param(weight_param, op_label="BasicLinear") + accumulate_into_main_grad = get_accumulate_flag_in_param(weight_param) + grad_weight = main_grad.detach() else: accumulate_into_main_grad = False @@ -99,15 +96,7 @@ def fuser_backward( # Megatron-LM wgrad fusion # Note: Return dummy tensor for grad weight if needed. if accumulate_into_main_grad: - grad_weight = None - weight_param = linear_op.weight - if hasattr(weight_param, "grad_added_to_main_grad"): - weight_param.grad_added_to_main_grad = True - grad_weight = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) + grad_weight = get_dummy_wgrads_for_params([linear_op.weight])[0] return grad_input, [(), (grad_weight,)], [(), ()] diff --git a/transformer_engine/pytorch/ops/fused/backward_linear_scale.py b/transformer_engine/pytorch/ops/fused/backward_linear_scale.py index 709073e6f..b48c2e6d5 100644 --- a/transformer_engine/pytorch/ops/fused/backward_linear_scale.py +++ b/transformer_engine/pytorch/ops/fused/backward_linear_scale.py @@ -9,9 +9,13 @@ import torch -from ...module.base import get_dummy_wgrad from ...utils import clear_tensor_data from ..basic import BasicLinear, ConstantScale +from .._common import ( + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, +) from ..op import FusedOperation, FusibleOperation, OperationContext @@ -58,16 +62,9 @@ def fuser_backward( grad_weight = None if linear_op_ctx.weight_requires_grad and accumulate_into_main_grad: weight_param = linear_op.weight - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - accumulate_into_main_grad = not getattr(weight_param, "overwrite_main_grad", False) - if not hasattr(weight_param, "main_grad"): - raise RuntimeError( - "BasicLinear op is configured with " - "accumulate_into_main_grad=True, " - "but weight parameter does not have main_grad attribute" - ) - grad_weight = weight_param.main_grad.detach() + main_grad = get_main_grad_from_param(weight_param, op_label="BasicLinear") + accumulate_into_main_grad = get_accumulate_flag_in_param(weight_param) + grad_weight = main_grad.detach() else: accumulate_into_main_grad = False @@ -99,15 +96,7 @@ def fuser_backward( # Megatron-LM wgrad fusion # Note: Return dummy tensor for grad weight if needed. if accumulate_into_main_grad: - grad_weight = None - weight_param = linear_op.weight - if hasattr(weight_param, "grad_added_to_main_grad"): - weight_param.grad_added_to_main_grad = True - grad_weight = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) + grad_weight = get_dummy_wgrads_for_params([linear_op.weight])[0] return grad_input, [(grad_weight,), ()], [(), ()] diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index 599e5f96a..91db2ff9b 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -160,10 +160,9 @@ def fuser_forward( if int(split_sizes.numel()) != num_groups: raise ValueError(f"Expected {num_groups} splits, but got {int(split_sizes.numel())}.") split_sizes = split_sizes.to(dtype=torch.int64, device=device) - base_offsets = tex.splits_to_offsets(split_sizes, 1) - split_points = base_offsets[1:].to(dtype=torch.int) - fc1_x_tensor_offsets = base_offsets * fc1_weight_shape[1] - fc2_x_tensor_offsets = base_offsets * fc2_weight_shape[1] + base_split_offsets = tex.splits_to_offsets(split_sizes, 1) + split_points = base_split_offsets[1:].to(dtype=torch.int) + fc2_x_tensor_offsets = base_split_offsets * fc2_weight_shape[1] # Extract post-scales from extra input scales = basic_op_extra_inputs[1][0] @@ -452,27 +451,35 @@ def fuser_forward( # Save state for backward pass if requires_grad: mark_grouped_tensor(grouped_fc1_x, swiglu_in, scales, grouped_fc2_x) - fc1_input_tensors = ( - grouped_fc1_x.columnwise_data, - grouped_fc1_x.columnwise_scale_inv, - fc1_x_tensor_offsets, - ) - # FC1 + + # Save the input ``GroupedTensor``s themselves for the activations. + for grouped_fc_x in (grouped_fc1_x, grouped_fc2_x): + if grouped_fc_x is not None: + grouped_fc_x.rowwise_data = None + grouped_fc_x.scale_inv = None + + # FC1 saved-tensor layout. + # [split_sizes, base_split_offsets, split_points, + # grouped_fc1_x, *fc1_weight_tensors] fc1_weight_tensors = ( [grouped_fc1_weight] if fc1_op.single_grouped_weight else grouped_fc1_weight ) fc1_ctx.save_for_backward( - split_sizes, split_points, *fc1_weight_tensors, *fc1_input_tensors + split_sizes, + base_split_offsets, + split_points, + grouped_fc1_x, + *fc1_weight_tensors, ) + fc1_ctx.use_grouped_tensor_path = True fc1_ctx.with_quantized_compute = True - fc1_ctx.input_quantizer = fc1_input_quantizer - fc1_ctx.weight_quantizer = fc1_weight_quantizer - fc1_ctx.grad_output_quantizer = fc1_grad_output_quantizer + fc1_ctx.input_quantizers = [fc1_input_quantizer] + fc1_ctx.weight_quantizers = [fc1_weight_quantizer] + fc1_ctx.grad_output_quantizers = [fc1_grad_output_quantizer] fc1_ctx.grad_input_quantizers = None fc1_ctx.dtype = dtype fc1_ctx.input_requires_grad = input_requires_grad fc1_ctx.weight_requires_grad = weight_requires_grad - fc1_ctx.base_split_offsets = base_offsets # Scaled SwiGLU swiglu_ctx.save_for_backward(swiglu_in, scales) @@ -480,25 +487,31 @@ def fuser_forward( swiglu_ctx.extra_input_requires_grad = True swiglu_ctx.dtype = dtype - # FC2 state - if grouped_fc2_x is not None: - fc2_input_tensors = ( - grouped_fc2_x.columnwise_data, - grouped_fc2_x.columnwise_scale_inv, - fc2_x_tensor_offsets, - ) - else: - fc2_input_tensors = (None, None, None) - - if fc2_op.single_grouped_weight: - fc2_ctx.save_for_backward(split_sizes, grouped_fc2_weight, *fc2_input_tensors) - else: - fc2_ctx.save_for_backward(split_sizes, *grouped_fc2_weight, *fc2_input_tensors) - + # FC2 saved-tensor layout. Matches the unfused + # ``GroupedLinear._fuser_forward_grouped_tensor`` layout so the + # unfused backward (basic/grouped_linear.py) can consume the same + # ctx when the fused backward is unavailable. + # [split_sizes, base_split_offsets, split_points, + # (fc2_scales if _scale_bias), + # grouped_fc2_x, *fc2_weight_tensors] + fc2_weight_tensors = ( + [grouped_fc2_weight] if fc2_op.single_grouped_weight else grouped_fc2_weight + ) + fc2_saved: list[Optional[torch.Tensor]] = [ + split_sizes, + base_split_offsets, + split_points, + ] + if fc2_op._scale_bias: + fc2_saved.append(fc2_scales) + fc2_saved.append(grouped_fc2_x) + fc2_saved.extend(fc2_weight_tensors) + fc2_ctx.save_for_backward(*fc2_saved) + fc2_ctx.use_grouped_tensor_path = True fc2_ctx.with_quantized_compute = True - fc2_ctx.input_quantizer = fc2_input_quantizer - fc2_ctx.weight_quantizer = fc2_weight_quantizer - fc2_ctx.grad_output_quantizer = fc2_grad_output_quantizer + fc2_ctx.input_quantizers = [fc2_input_quantizer] + fc2_ctx.weight_quantizers = [fc2_weight_quantizer] + fc2_ctx.grad_output_quantizers = [fc2_grad_output_quantizer] fc2_ctx.grad_input_quantizers = None fc2_ctx.dtype = dtype fc2_ctx.input_requires_grad = input_requires_grad diff --git a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py index fbaf69d75..7d67815f9 100644 --- a/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py +++ b/transformer_engine/pytorch/ops/fused/userbuffers_backward_linear.py @@ -17,14 +17,19 @@ _2X_ACC_DGRAD, _2X_ACC_WGRAD, fill_userbuffers_buffer_for_all_gather, - get_dummy_wgrad, get_ub, ) from ...quantized_tensor import Quantizer from ...tensor.mxfp8_tensor import MXFP8Quantizer from ...utils import canonicalize_device, canonicalize_dtype, clear_tensor_data from ..basic import BasicLinear, Bias, ReduceScatter -from .._common import maybe_dequantize, is_quantized_tensor +from .._common import ( + get_accumulate_flag_in_param, + get_dummy_wgrads_for_params, + get_main_grad_from_param, + is_quantized_tensor, + maybe_dequantize, +) from ..op import FusedOperation, FusibleOperation, OperationContext @@ -519,16 +524,9 @@ def fuser_backward( grad_weight = None if linear_op_ctx.weight_requires_grad and accumulate_into_main_grad: weight_param = linear_op.weight - if hasattr(weight_param, "__fsdp_param__"): - weight_param.main_grad = weight_param.get_main_grad() - accumulate_into_main_grad = not getattr(weight_param, "overwrite_main_grad", False) - if not hasattr(weight_param, "main_grad"): - raise RuntimeError( - "BasicLinear op is configured with " - "accumulate_into_main_grad=True, " - "but weight parameter does not have main_grad attribute" - ) - grad_weight = weight_param.main_grad.detach() + main_grad = get_main_grad_from_param(weight_param, op_label="UserbuffersBackwardLinear") + accumulate_into_main_grad = get_accumulate_flag_in_param(weight_param) + grad_weight = main_grad.detach() else: accumulate_into_main_grad = False @@ -563,15 +561,7 @@ def fuser_backward( # Megatron-LM wgrad fusion # Note: Return dummy tensor for grad weight if needed. if accumulate_into_main_grad: - grad_weight = None - weight_param = linear_op.weight - if hasattr(weight_param, "grad_added_to_main_grad"): - weight_param.grad_added_to_main_grad = True - grad_weight = get_dummy_wgrad( - list(weight_param.size()), - weight_param.dtype, - zero=getattr(weight_param, "zero_out_wgrad", False), - ) + grad_weight = get_dummy_wgrads_for_params([linear_op.weight])[0] # Return gradients grad_params = [() for _ in range(len(self.basic_ops))] diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 5f12c3ed8..485b32328 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -303,6 +303,51 @@ def get_dtype(self) -> torch.dtype: return self.fake_dtype + def prepare_for_saving( + self, + ) -> Tuple[list[Optional[torch.Tensor]], "GroupedTensorStorage"]: + """Prepare the tensor base for saving for backward.""" + tensors = [ + self.rowwise_data, + self.columnwise_data, + self.scale_inv, + self.columnwise_scale_inv, + self.amax, + self.columnwise_amax, + self.scale, + self.first_dims, + self.last_dims, + self.tensor_offsets, + ] + self.rowwise_data = None + self.columnwise_data = None + self.scale_inv = None + self.columnwise_scale_inv = None + self.amax = None + self.columnwise_amax = None + self.scale = None + self.first_dims = None + self.last_dims = None + self.tensor_offsets = None + self.quantized_tensors = None + return tensors, self + + def restore_from_saved( + self, tensors: list[Optional[torch.Tensor]] + ) -> list[Optional[torch.Tensor]]: + """Restore the tensor base data from the saved tensors list.""" + self.rowwise_data = tensors[0] + self.columnwise_data = tensors[1] + self.scale_inv = tensors[2] + self.columnwise_scale_inv = tensors[3] + self.amax = tensors[4] + self.columnwise_amax = tensors[5] + self.scale = tensors[6] + self.first_dims = tensors[7] + self.last_dims = tensors[8] + self.tensor_offsets = tensors[9] + return tensors[10:] + def clear(self) -> None: """ Reset tensor data and clear all buffers. From 3c89426637254a098b27b8bc3094b65c2fa66a43 Mon Sep 17 00:00:00 2001 From: vcherepanov-nv Date: Tue, 5 May 2026 21:59:28 -0700 Subject: [PATCH 037/180] [Common] Always define cuBLASMp comm GEMM API (#2963) * Always define cuBLASMp comm GEMM API Signed-off-by: Vladimir Cherepanov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Vladimir Cherepanov Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/common/CMakeLists.txt | 6 +-- .../common/comm_gemm/comm_gemm.cpp | 47 ++++++++++++++++++- 2 files changed, 47 insertions(+), 6 deletions(-) diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 781fe4881..734941595 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -159,6 +159,7 @@ list(APPEND transformer_engine_cpp_sources util/cuda_runtime.cpp util/multi_stream.cpp util/rtc.cpp + comm_gemm/comm_gemm.cpp comm_gemm_overlap/userbuffers/ipcsocket.cc comm_gemm_overlap/userbuffers/userbuffers-host.cpp comm_gemm_overlap/comm_gemm_overlap.cpp @@ -280,11 +281,6 @@ foreach(cuda_source IN LISTS transformer_engine_cuda_arch_specific_sources) endif() endforeach() -if (NVTE_WITH_CUBLASMP) -list(APPEND transformer_engine_SOURCES - comm_gemm/comm_gemm.cpp) -endif() - add_library(transformer_engine SHARED ${transformer_engine_SOURCES}) # Disable CMake's automatic architecture flag injection. # All architectures are handled explicitly via per-source COMPILE_OPTIONS diff --git a/transformer_engine/common/comm_gemm/comm_gemm.cpp b/transformer_engine/common/comm_gemm/comm_gemm.cpp index a7d78f7ac..ce389c200 100644 --- a/transformer_engine/common/comm_gemm/comm_gemm.cpp +++ b/transformer_engine/common/comm_gemm/comm_gemm.cpp @@ -6,7 +6,6 @@ #include "transformer_engine/comm_gemm.h" -#include #include #include @@ -21,6 +20,10 @@ #include "../common.h" #include "../util/logging.h" +#ifdef NVTE_WITH_CUBLASMP + +#include + using namespace transformer_engine; namespace { @@ -530,3 +533,45 @@ int64_t nvte_comm_gemm_numroc(NVTECommGemmCtx* ctx, int64_t global_size) { NVTE_API_CALL(nvte_comm_gemm_numroc); return cublasMpNumroc(global_size, block_size(ctx, global_size), ctx->rank, 0, ctx->nranks); } + +#else // NVTE_WITH_CUBLASMP + +struct NVTECommGemmCtx {}; + +NVTECommGemmCtx* nvte_comm_gemm_ctx_create(ncclComm_t comm, int nranks, int rank) { + NVTE_ERROR("Transformer Engine has not been built with cuBLASMp support."); +} + +void nvte_comm_gemm_ctx_destroy(NVTECommGemmCtx* ctx) { + NVTE_ERROR("Transformer Engine has not been built with cuBLASMp support."); +} + +void nvte_all_gather_gemm(NVTECommGemmCtx* ctx, int64_t m, int64_t n, int64_t k, const NVTETensor a, + const NVTETensor b, const NVTETensor d, const NVTETensor bias, + const NVTETensor pre_act_out, bool transa, bool transb, bool grad, + bool accumulate, int comm_sm_count, cudaStream_t main_stream, + NVTECommGemmAlgoType algo) { + NVTE_ERROR("Transformer Engine has not been built with cuBLASMp support."); +} + +void nvte_gemm_reduce_scatter(NVTECommGemmCtx* ctx, int64_t m, int64_t n, int64_t k, + const NVTETensor a, const NVTETensor b, const NVTETensor d, + const NVTETensor bias, const NVTETensor pre_act_out, bool transa, + bool transb, bool grad, bool accumulate, int comm_sm_count, + cudaStream_t main_stream, NVTECommGemmAlgoType algo) { + NVTE_ERROR("Transformer Engine has not been built with cuBLASMp support."); +} + +void nvte_gemm_all_reduce(NVTECommGemmCtx* ctx, int64_t m, int64_t n, int64_t k, const NVTETensor a, + const NVTETensor b, const NVTETensor d, const NVTETensor bias, + const NVTETensor pre_act_out, bool transa, bool transb, bool grad, + bool accumulate, int comm_sm_count, cudaStream_t main_stream, + NVTECommGemmAlgoType algo) { + NVTE_ERROR("Transformer Engine has not been built with cuBLASMp support."); +} + +int64_t nvte_comm_gemm_numroc(NVTECommGemmCtx* ctx, int64_t global_size) { + NVTE_ERROR("Transformer Engine has not been built with cuBLASMp support."); +} + +#endif // NVTE_WITH_CUBLASMP From 4b6923dd1141272c668a96e9486a5482ec6ada40 Mon Sep 17 00:00:00 2001 From: Kshitij Lakhani <33047503+KshitijLakhani@users.noreply.github.com> Date: Wed, 6 May 2026 11:08:17 -0700 Subject: [PATCH 038/180] [JAX][Common] Enable cuDNN fused attn backend for NO_MASK + bidirectional SWA (#2961) * Enable right side of sliding window for cuDNN fused attn backend Signed-off-by: Kshitij Lakhani * Add window size in the warning string when falling back to unfused attn Signed-off-by: Kshitij Lakhani * Add a test for bidirectional asymmetric SWA testing in fused attn. Also add a helper to pick window based on cuDNN version support in fused_attn.cpp Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * nit: Code clean up Signed-off-by: Kshitij Lakhani * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Kshitij Lakhani Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/jax/test_fused_attn.py | 32 +++++++++++++++---- .../common/fused_attn/fused_attn.cpp | 1 + transformer_engine/jax/flax/transformer.py | 2 +- 3 files changed, 28 insertions(+), 7 deletions(-) diff --git a/tests/jax/test_fused_attn.py b/tests/jax/test_fused_attn.py index 8b727b1d4..1fb010806 100644 --- a/tests/jax/test_fused_attn.py +++ b/tests/jax/test_fused_attn.py @@ -1060,6 +1060,30 @@ def check_dqkv(primitive, reference, pad, idx): assert_equal_collectives(target_hlo, self.coll_count_ref) +def _get_swa_window_size_for_test(s_kv: int, attn_mask_type: AttnMaskType) -> Tuple[int, int]: + """Pick a sliding-window size for SWA tests, gated on cuDNN version. + + cuDNN < 9.2: skip (no SWA support). + cuDNN >= 9.2: left-only window (s_kv // 10, 0). + cuDNN >= 9.6: bidirectional window (s_kv // 10, s_kv // 10 + 5) for the mask types whose + bidirectional fused dispatch is meaningful here (NO_MASK, PADDING_MASK). + Other mask types keep the left-only window: causal-family masks would + collapse (W, W) -> (W, 0), hence not tested here. + """ + cudnn_version = get_cudnn_version() + if cudnn_version < 90200: + pytest.skip("Sliding window attention requires cuDNN >= 9.2") + left_window_size = s_kv // 10 + # choose asymmetric window size for testing + right_window_size = left_window_size + 5 + if cudnn_version >= 90600 and attn_mask_type in ( + AttnMaskType.NO_MASK, + AttnMaskType.PADDING_MASK, + ): + return (left_window_size, right_window_size) + return (left_window_size, 0) + + @pytest.mark.parametrize( "attn_mask_type", [ @@ -1330,9 +1354,7 @@ def _test_forward( This test is not intended to run automatically during CI as it is time-consuming It is kept for development and debugging """ - window_size = None - if swa: - window_size = (s_kv // 10, 0) + window_size = _get_swa_window_size_for_test(s_kv, attn_mask_type) if swa else None runner = FusedAttnRunner( b, s_q, @@ -1383,9 +1405,7 @@ def test_backward( """ Test backward with parameterized configs """ - window_size = None - if swa: - window_size = (s_kv // 10, 0) + window_size = _get_swa_window_size_for_test(s_kv, attn_mask_type) if swa else None runner = FusedAttnRunner( b, s_q, diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index 141767b80..ae8ddbed6 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -469,6 +469,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( (sm_arch_ < 100 || (sm_arch_ >= 100 && ((max_seqlen_q == max_seqlen_kv && cudnn_runtime_version <= 90700) || cudnn_runtime_version > 90700)))) || + attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK || (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_BOTTOM_RIGHT_MASK && diff --git a/transformer_engine/jax/flax/transformer.py b/transformer_engine/jax/flax/transformer.py index 513677e4a..a2e792084 100644 --- a/transformer_engine/jax/flax/transformer.py +++ b/transformer_engine/jax/flax/transformer.py @@ -788,7 +788,7 @@ def __call__( "Fall back to the unfused attention.\n" "Please try to update the cuDNN and TE to the latest version.\n" f"{qkv_layout=}\n{attn_bias_type=}\n{attn_mask_type=}\n" - f"{self.attention_dropout=}\n{self.num_attention_heads=}\n" + f"{self.attention_dropout=}\n{self.num_attention_heads=}\n{self.window_size=}\n" f"{self.num_gqa_groups=}\n{seqlen_q=}\n{seqlen_kv=}\n{head_dim_qk=}\n{head_dim_v=}\n" ) From 2f3eda406546aca2de3cb1e6c815a1c7220a7a43 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 7 May 2026 16:34:53 -0700 Subject: [PATCH 039/180] [All] Remove legacy max512 backend (#2949) * remove max512 subbackend Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * minor tweak for docstring Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * revert fp8 t3hd changes Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove redudant test comparison 0 vs 1 subbackend Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove sub-backend 0 from header docstring Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --------- Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/envvars.rst | 4 +- tests/pytorch/attention/test_attention.py | 49 +- tests/pytorch/utils.py | 4 +- transformer_engine/common/CMakeLists.txt | 1 - .../common/fused_attn/fused_attn.cpp | 61 +- .../fused_attn_f16_max512_seqlen.cu | 1343 ----------------- .../fused_attn/fused_attn_f16_max512_seqlen.h | 41 - .../include/transformer_engine/fused_attn.h | 10 +- .../common/util/pybind_helper.h | 1 - .../jax/cpp_extensions/attention.py | 5 +- .../jax/csrc/extensions/attention.cpp | 11 - .../jax/csrc/extensions/pybind.cpp | 1 - .../dot_product_attention/backends.py | 31 +- .../attention/dot_product_attention/utils.py | 30 - .../pytorch/cpp_extensions/fused_attn.py | 38 +- .../pytorch/csrc/extensions/attention.cpp | 1 - 16 files changed, 40 insertions(+), 1591 deletions(-) delete mode 100644 transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu delete mode 100644 transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.h diff --git a/docs/envvars.rst b/docs/envvars.rst index 1e040b4c3..29ca49814 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -142,9 +142,9 @@ Attention Backend Selection .. envvar:: NVTE_FUSED_ATTN_BACKEND - :Type: ``int`` (0, 1, or 2) + :Type: ``int`` (1 or 2) :Default: Auto-selected - :Description: Force a specific FusedAttention backend. ``0`` = F16_max512_seqlen (cuDNN, ≤512 seq len), ``1`` = F16_arbitrary_seqlen (cuDNN, any seq len), ``2`` = FP8 backend. If not set, the backend is automatically selected based on the input configuration. + :Description: Force a specific FusedAttention backend. ``1`` = F16_arbitrary_seqlen (cuDNN, any seq len), ``2`` = FP8 backend. If not set, the backend is automatically selected based on the input configuration. .. envvar:: NVTE_FUSED_ATTN_FORCE_WORKSPACE_OPT diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index c9ea79144..894137c84 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -227,40 +227,16 @@ def test_dot_product_attention( # FusedAttention backend if fused_attn_supported: - if len(fused_attn_backends) == 1: - fused_attn_fwd, fused_max_logit, fused_attn_bwd = _run_dot_product_attention( - dtype, - config, - "FusedAttention", - ckpt_attn, - qkv_layout, - workspace_opt, - pad_between_seqs, - is_training, - ) - if len(fused_attn_backends) == 2: - os.environ["NVTE_FUSED_ATTN_BACKEND"] = "0" - fused_attn_fwd, _, fused_attn_bwd = _run_dot_product_attention( - dtype, - config, - "FusedAttention", - ckpt_attn, - qkv_layout, - workspace_opt, - pad_between_seqs, - is_training, - ) - os.environ["NVTE_FUSED_ATTN_BACKEND"] = "1" - fused_attn_fwd_1, _, fused_attn_bwd_1 = _run_dot_product_attention( - dtype, - config, - "FusedAttention", - ckpt_attn, - qkv_layout, - workspace_opt, - pad_between_seqs, - is_training, - ) + fused_attn_fwd, fused_max_logit, fused_attn_bwd = _run_dot_product_attention( + dtype, + config, + "FusedAttention", + ckpt_attn, + qkv_layout, + workspace_opt, + pad_between_seqs, + is_training, + ) # FlashAttention backend if flash_attn_supported: @@ -294,11 +270,6 @@ def test_dot_product_attention( torch.testing.assert_close(fused_attn_fwd, flash_attn_fwd, **tols) for i, _ in enumerate(flash_attn_bwd): torch.testing.assert_close(fused_attn_bwd[i], flash_attn_bwd[i], **tols) - if fused_attn_supported and len(fused_attn_backends) == 2: - logging.info("[test_dot_product_attention]: fused attn backend 0 vs 1") - torch.testing.assert_close(fused_attn_fwd, fused_attn_fwd_1, **tols) - for i, _ in enumerate(fused_attn_bwd): - torch.testing.assert_close(fused_attn_bwd[i], fused_attn_bwd_1[i], **tols) @pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index c7cbe78a6..3b2e50be3 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -392,11 +392,11 @@ def test(): _attention_backends["backend_selection_requires_update"] = False return available_backends, flash_attention_backend, fused_attention_backend - backends = {0: "F16_max512_seqlen", 1: "F16_arbitrary_seqlen", 2: "FP8"} + backends = {1: "F16_arbitrary_seqlen", 2: "FP8"} if AttentionLogging._is_logging_setup is False: AttentionLogging.setup_logging() - for i in range(3): + for i in backends: os.environ["NVTE_FUSED_ATTN_BACKEND"] = str(i) _attention_backends["backend_selection_requires_update"] = True available_backends, flash_attention_backend, fused_attention_backend = test() diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 734941595..030023d94 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -182,7 +182,6 @@ list(APPEND transformer_engine_cuda_sources dropout/dropout.cu fused_attn/context_parallel.cu fused_attn/kv_cache.cu - fused_attn/fused_attn_f16_max512_seqlen.cu fused_attn/fused_attn_f16_arbitrary_seqlen.cu fused_attn/fused_attn_fp8.cu fused_attn/utils.cu diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index ae8ddbed6..f18a006fc 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -11,7 +11,6 @@ #include "../util/cuda_runtime.h" #include "../util/system.h" #include "fused_attn_f16_arbitrary_seqlen.h" -#include "fused_attn_f16_max512_seqlen.h" #include "fused_attn_fp8.h" #include "utils.h" @@ -304,28 +303,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( << std::endl; } } else if ((q_dtype == NVTEDType::kNVTEFloat16) || (q_dtype == NVTEDType::kNVTEBFloat16)) { - bool flag_m512 = false; bool flag_arb = false; - if ((sm_arch_ == 80 || sm_arch_ == 90) && (max_seqlen_q <= 512 && max_seqlen_q % 64 == 0) && - (max_seqlen_kv <= 512 && max_seqlen_kv % 64 == 0) && (head_dim_qk == 64) && - (head_dim_v == 64) && (num_attn_heads == num_gqa_groups) && - ((bias_type == NVTE_Bias_Type::NVTE_NO_BIAS) || - (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS)) && - ((attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK) || - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - (attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK && - max_seqlen_q == max_seqlen_kv) || - (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK)) && - ((qkv_layout == NVTE_QKV_Layout::NVTE_SB3HD) || - (qkv_layout == NVTE_QKV_Layout::NVTE_SBHD_SB2HD) || - (qkv_layout == NVTE_QKV_Layout::NVTE_BS3HD) || - (qkv_layout == NVTE_QKV_Layout::NVTE_BSHD_BS2HD) || - (qkv_layout == NVTE_QKV_Layout::NVTE_BSHD_BSHD_BSHD)) && - ((window_size_left == -1) && (window_size_right == -1 || window_size_right == 0)) && - !requires_64bit_ragged_offset && - (softmax_type == NVTE_Softmax_Type::NVTE_VANILLA_SOFTMAX) && !return_max_logit) { - flag_m512 = true; - } if ( // TODO(cyang): replace with cudnn-frontend check_support for cleaner logic and better error messaging // architecture @@ -499,31 +477,9 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( dropout == 0.0 && bias_type == NVTE_Bias_Type::NVTE_NO_BIAS))))) { flag_arb = true; } - if (((max_seqlen_q > 512) || (max_seqlen_kv > 512)) && (flag_arb == true)) { + if (flag_arb) { backend = NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; } - if ((max_seqlen_q <= 512) && (max_seqlen_kv <= 512)) { - if (flag_arb == true) { - backend = NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen; - } else if ((flag_arb == false) && (flag_m512 == true)) { - backend = NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen; - } - int env_backend = static_cast(backend); - env_backend = transformer_engine::getenv("NVTE_FUSED_ATTN_BACKEND", env_backend); - if (((env_backend == static_cast(NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen)) && - flag_m512) || - ((env_backend == static_cast(NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen)) && - flag_arb)) { - backend = static_cast(env_backend); - } - } - if (cudnn_runtime_version < 8901 && - backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: FP16/BF16 fused attention is supported by cuDNN 8.9.1+." - " Please upgrade your cuDNN version if possible." - << std::endl; - } if (cudnn_runtime_version < 8900 && backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; @@ -668,12 +624,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right, return_max_logit, cuda_graph, false); - if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { - fused_attn_max_512_fwd(b, h_q, max_seqlen_q, max_seqlen_kv, d_qk, is_training, attn_scale, - dropout, qkv_layout, bias_type, attn_mask_type, input_Q, input_K, - input_V, input_Bias, output_O, Aux_CTX_Tensors, input_cu_seqlens_q, - input_cu_seqlens_kv, input_rng_state, wkspace, stream, handle); - } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { + if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { fused_attn_arbitrary_seqlen_fwd( b, h_q, h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, t_q, t_kv, num_pages_k, num_pages_v, page_size_k, page_size_v, max_pages_per_seq_k, max_pages_per_seq_v, is_training, @@ -754,13 +705,7 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso h_kv, max_seqlen_q, max_seqlen_kv, d_qk, d_v, window_size_left, window_size_right, false, cuda_graph, deterministic); - if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[0]); - fused_attn_max_512_bwd(b, h_q, max_seqlen_q, max_seqlen_kv, d_qk, attn_scale, dropout, - qkv_layout, bias_type, attn_mask_type, input_Q, input_K, input_V, - input_dO, output_S, output_dQ, output_dK, output_dV, output_dBias, - input_cu_seqlens_q, input_cu_seqlens_kv, wkspace, stream, handle); - } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { + if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) { size_t i = 0; Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); Tensor *input_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu b/transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu deleted file mode 100644 index d5151a51f..000000000 --- a/transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.cu +++ /dev/null @@ -1,1343 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -#include -#include -#include - -#include -#include - -#include "../common.h" -#include "../cudnn_utils.h" -#include "fused_attn_f16_max512_seqlen.h" -#include "utils.h" - -#define Q_ID 1 -#define K_ID 2 -#define V_ID 3 -#define O_ID 4 -#define S_ID 5 -#define B_ID 6 -#define DROPOUT_CONST_ID 7 -#define S_CONST_ID 8 -#define Q_SEQLEN_ID 9 -#define K_SEQLEN_ID 10 -#define dQ_ID 11 -#define dK_ID 12 -#define dV_ID 13 -#define dO_ID 14 -#define MASK_VAL_ID 15 -#define dS_ID 16 -#define dBias_ID 17 -#define DROPOUT_SEED_ID 18 -#define DROPOUT_OFFSET_ID 19 - -#define VIRTUAL_ID 20 - -namespace transformer_engine { -namespace fused_attn { - -static void createScale(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, - NVTE_QKV_Layout layout, cudnnDataType_t tensorType, - std::vector &ops) { - // scale - int64_t scale_dim[4] = {1, 1, 1, 1}; - int64_t scale_stride[4] = {1, 1, 1, 1}; - - int64_t k_dim[4] = {b, h, d, s_kv}; - int64_t k_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, k_stride, layout, - NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose); - - auto scaleTensor = - tensor_create(tensorType, S_CONST_ID, scale_dim, scale_stride, false, true); // is by value - auto kTensor = tensor_create(tensorType, K_ID, k_dim, k_stride, false, false); - auto afterScaleKTensor = - tensor_create(tensorType, VIRTUAL_ID, k_dim, k_stride, true, false); // is virtual - - // Define the scale descriptor - auto scaleDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // Create a Scale Node. - auto scale_op = binary_pw_op_create(kTensor, scaleTensor, afterScaleKTensor, scaleDesc); - - ops.push_back(std::move(scale_op)); -} - -static cudnn_frontend::Tensor createBMM1(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, - NVTE_QKV_Layout layout, cudnnDataType_t tensorType, - bool zero_s, std::vector &ops) { - // Creates the necessary tensor descriptors - int64_t q_dim[4] = {b, h, s_q, d}; - int64_t q_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, q_stride, layout, NVTE_QKV_Matrix::NVTE_Q_Matrix); - - int64_t k_dim[4] = {b, h, d, s_kv}; - int64_t k_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, k_stride, layout, - NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose); - - int64_t p_dim[4] = {b, h, s_q, s_kv}; - int64_t p_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, p_stride, layout, NVTE_QKV_Matrix::NVTE_S_Matrix); - - int64_t seqlen_dim[4] = {b, 1, 1, 1}; - int64_t seqlen_stride[4] = {1, 1, 1, 1}; - - auto qTensor = tensor_create(tensorType, Q_ID, q_dim, q_stride, false, false); - auto afterScaleKTensor = - tensor_create(tensorType, VIRTUAL_ID, k_dim, k_stride, true, false); // is virtual - // first GEMM output - auto pTensor = tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 1, p_dim, p_stride, true, - false); // is virtual - - auto seqlenQTensor = - tensor_create(CUDNN_DATA_INT32, Q_SEQLEN_ID, seqlen_dim, seqlen_stride, false, false); - auto seqlenKTensor = - tensor_create(CUDNN_DATA_INT32, K_SEQLEN_ID, seqlen_dim, seqlen_stride, false, false); - - // Define the matmul 1 desc - // set padding value optionally to 0 for writing zeros to S tensor (if not set, old behaviour) - auto matmul_1_Desc = cudnn_frontend::MatMulDescBuilder().setComputeType(CUDNN_DATA_FLOAT).build(); - - if (zero_s) { - matmul_1_Desc = cudnn_frontend::MatMulDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setPaddingValue(0.0f) - .build(); - } - - // Create a matmul 1 Node - auto matmul_op1 = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_MATMUL_DESCRIPTOR) - .setaMatDesc(qTensor) - .setbMatDesc(afterScaleKTensor) - .setcMatDesc(pTensor) - .setmOverrideDesc(seqlenQTensor) - .setnOverrideDesc(seqlenKTensor) - .setmatmulDesc(matmul_1_Desc) - .build(); - - ops.push_back(std::move(matmul_op1)); - - return pTensor; -} - -static cudnn_frontend::Tensor createBias(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, - NVTE_QKV_Layout layout, cudnnDataType_t tensorType, - std::vector &ops, - cudnn_frontend::Tensor const &prevBlockOutputTensor) { - NVTE_CHECK(ops.size() != 0, "Bias op constructed incorrectly as the first one."); - - int64_t b_dim[4] = {1, h, s_q, s_kv}; - int64_t b_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1}; - - int64_t afterBias_dim[4] = {b, h, s_q, s_kv}; - int64_t afterBias_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, afterBias_stride, layout, - NVTE_QKV_Matrix::NVTE_S_Matrix); - - // bias - auto bTensor = tensor_create(tensorType, B_ID, b_dim, b_stride, false, false); - // output - auto afterBiasTensor = tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 50, afterBias_dim, - afterBias_stride, true, false); // is virtual - - // Define the bias descriptor - auto biasDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_ADD); - - // Create a Bias Node. - auto bias_op = binary_pw_op_create(prevBlockOutputTensor, bTensor, afterBiasTensor, biasDesc); - - ops.push_back(std::move(bias_op)); - - return afterBiasTensor; -} - -static cudnn_frontend::Tensor createMask(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, - NVTE_QKV_Layout layout, NVTE_Mask_Type mask_type, - cudnnDataType_t tensorType, - std::vector &ops, - cudnn_frontend::Tensor const &prevBlockOutputTensor, - bool is_bprop) { - NVTE_CHECK(ops.size() != 0, "Padding mask constructed incorrectly as the first one."); - - // subtraction output - int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv}; - int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1}; - - int64_t seqlen_dim[4] = {b, 1, 1, 1}; - int64_t seqlen_stride[4] = {1, 1, 1, 1}; - - int64_t maskVal_dim[4] = {1, 1, 1, 1}; - int64_t maskVal_stride[4] = {1, 1, 1, 1}; - - // mask value to put in the masked pixels - auto maskValTensor = tensor_create(CUDNN_DATA_FLOAT, MASK_VAL_ID, maskVal_dim, maskVal_stride, - false, true); // is by value - - auto seqlenQTensor = - tensor_create(CUDNN_DATA_INT32, Q_SEQLEN_ID, seqlen_dim, seqlen_stride, false, false); - auto seqlenKTensor = - tensor_create(CUDNN_DATA_INT32, K_SEQLEN_ID, seqlen_dim, seqlen_stride, false, false); - // gen index row output - auto rowIndexTensor = tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 100, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - // gen index column output - auto columnIndexTensor = tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 101, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - // less than row output - auto lessThanRowTensor = - tensor_create(CUDNN_DATA_BOOLEAN, VIRTUAL_ID + 102, afterBMM1_dim, afterBMM1_stride, true, - false); // is virtual - // less than column output - auto lessThanColTensor = tensor_create(CUDNN_DATA_BOOLEAN, VIRTUAL_ID + 103, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - // padding mask (lessthanRow && lessthanCol) - auto paddingMaskTensor = tensor_create(CUDNN_DATA_BOOLEAN, VIRTUAL_ID + 104, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - // row >= col check for causal mask - auto rowGreaterColTensor = tensor_create(CUDNN_DATA_BOOLEAN, VIRTUAL_ID + 105, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - // create causal mask (padding && row >= col) - auto causalMaskTensor = tensor_create(CUDNN_DATA_BOOLEAN, VIRTUAL_ID + 106, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - - // output after masking - int64_t maskOutputTensor_id = VIRTUAL_ID + 107; - int64_t maskOutputTensor_virtual = true; - cudnnDataType_t maskOutputTensor_dataType = CUDNN_DATA_FLOAT; - auto maskOutputTensor_reorderType = cudnn_frontend::TensorReordering_t::NONE; - - if (is_bprop) { - maskOutputTensor_id = dS_ID; - maskOutputTensor_virtual = false; - maskOutputTensor_dataType = tensorType; - maskOutputTensor_reorderType = cudnn_frontend::TensorReordering_t::F16x16; - } - - auto maskOutputTensor = - cudnn_frontend::TensorBuilder() - .setDim(4, afterBMM1_dim) - .setStride(4, afterBMM1_stride) - .setAlignment(16) // 16B alignment is needed to run a tensor core engine - .setByValue(false) - .setDataType(maskOutputTensor_dataType) - .setVirtual(maskOutputTensor_virtual) - .setId(maskOutputTensor_id) - .setReorderType(maskOutputTensor_reorderType) - .build(); - - // Define the gen index for row descriptor - auto genIndexRowDesc = cudnn_frontend::PointWiseDescBuilder() - .setMode(CUDNN_POINTWISE_GEN_INDEX) - .setAxis(2) - .setComputeType(CUDNN_DATA_FLOAT) - .build(); - - // Create a gen index Node. - auto genIndexRow_op = unary_pw_op_create(prevBlockOutputTensor, rowIndexTensor, genIndexRowDesc); - - // Define the gen index for row descriptor - auto genIndexColumnDesc = cudnn_frontend::PointWiseDescBuilder() - .setMode(CUDNN_POINTWISE_GEN_INDEX) - .setAxis(3) - .setComputeType(CUDNN_DATA_FLOAT) - .build(); - - // Create a gen index Node. - auto genIndexColumn_op = - unary_pw_op_create(prevBlockOutputTensor, columnIndexTensor, genIndexColumnDesc); - - // Define the less than comparison for row descriptor - auto lessThanRowDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_CMP_LT); - - // Create a less than comparison for row Node. - auto lessThanRow_op = - binary_pw_op_create(rowIndexTensor, seqlenQTensor, lessThanRowTensor, lessThanRowDesc); - - // Define the less than comparison for column descriptor - auto lessThanColDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_CMP_LT); - - // Create a less than comparison for col Node. - auto lessThanCol_op = - binary_pw_op_create(columnIndexTensor, seqlenKTensor, lessThanColTensor, lessThanColDesc); - - // Define the less than comparison for column descriptor - auto paddingMaskAndDesc = pw_desc_create(CUDNN_DATA_BOOLEAN, CUDNN_POINTWISE_LOGICAL_AND); - - // Create a and node for combining lessThanRow and lessThanCol - auto paddingMaskAnd_op = binary_pw_op_create(lessThanRowTensor, lessThanColTensor, - paddingMaskTensor, paddingMaskAndDesc); - - // Define the greater than equal to comparison descriptor - auto rowGreaterColDesc = pw_desc_create(CUDNN_DATA_BOOLEAN, CUDNN_POINTWISE_CMP_GE); - - // Create a greater than equal to Node. - auto rowGreaterCol_op = binary_pw_op_create(rowIndexTensor, columnIndexTensor, - rowGreaterColTensor, rowGreaterColDesc); - - // Define the and to create causal mask descriptor - auto causalMaskAndDesc = pw_desc_create(CUDNN_DATA_BOOLEAN, CUDNN_POINTWISE_LOGICAL_AND); - - // Create a causal Mask Node. - auto causalMaskAnd_op = binary_pw_op_create(paddingMaskTensor, rowGreaterColTensor, - causalMaskTensor, causalMaskAndDesc); - - /////////////////// Apply the mask ////////////////////////// - - auto maskTensor = (mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) - ? std::move(causalMaskTensor) - : std::move(paddingMaskTensor); - - // Define the binary select to perform masking descriptor - auto maskDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_BINARY_SELECT); - - // Create a binary select Node. - auto mask_op = ternary_pw_op_create(prevBlockOutputTensor, maskValTensor, maskTensor, - maskOutputTensor, maskDesc); - - ops.push_back(std::move(genIndexRow_op)); - ops.push_back(std::move(genIndexColumn_op)); - ops.push_back(std::move(lessThanRow_op)); - ops.push_back(std::move(lessThanCol_op)); - ops.push_back(std::move(paddingMaskAnd_op)); - if (mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK) { - ops.push_back(std::move(rowGreaterCol_op)); - ops.push_back(std::move(causalMaskAnd_op)); - } - ops.push_back(std::move(mask_op)); - - return maskOutputTensor; -} - -static cudnn_frontend::Tensor createSoftmaxForward( - int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, NVTE_QKV_Layout layout, - bool enable_dropout, bool softmax_output_virtual, cudnnDataType_t tensorType, - std::vector &ops, - cudnn_frontend::Tensor const &prevBlockOutputTensor) { - int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv}; - int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1}; - - int64_t afterReduction_dim[4] = {b, h, s_q, 1}; - int64_t afterReduction_stride[4] = {h * s_q, s_q, 1, 1}; - - cudnnDataType_t softmaxOutputType = enable_dropout ? CUDNN_DATA_FLOAT : tensorType; - uint64_t softmaxOutputName = softmax_output_virtual ? VIRTUAL_ID + 154 : S_ID; - - // max (x) - auto afterMaxReductionTensor = - tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 150, afterReduction_dim, afterReduction_stride, - true, false); // is virtual - // x - max(x) - auto afterSubtractionTensor = tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 151, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - // e^(x - max(x)) - auto afterExponentTensor = tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 152, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual; - // sum (e^(x - max(x))) - auto afterAddReductionTensor = - tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 153, afterReduction_dim, afterReduction_stride, - true, false); // is virtual - // divide (e/ sum(e)) - - auto reorder_type = cudnn_frontend::TensorReordering_t::F16x16; - - auto afterDivisionTensor = - cudnn_frontend::TensorBuilder() - .setDim(4, afterBMM1_dim) - .setStride(4, afterBMM1_stride) - .setId(softmaxOutputName) - .setAlignment(16) // 16B alignment is needed to run a tensor core engine - .setDataType(softmaxOutputType) - .setVirtual(softmax_output_virtual) - .setByValue(false) - .setReorderType(reorder_type) - .build(); - - // Define the reduction descriptor - auto reductionMaxDesc = cudnn_frontend::ReductionDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setReductionOp(CUDNN_REDUCE_TENSOR_MAX) - .build(); - - // Create a reduction max Node. - auto reductionMax_op = - cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_REDUCTION_DESCRIPTOR) - .setxDesc(prevBlockOutputTensor) - .setyDesc(afterMaxReductionTensor) - .setreductionDesc(reductionMaxDesc) - .build(); - - // Define the subtract descriptor - auto subtractDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_SUB); - - // Create a subtract Node. - auto subtract_op = binary_pw_op_create(prevBlockOutputTensor, afterMaxReductionTensor, - afterSubtractionTensor, subtractDesc); - - // Define the exponent descriptor - auto exponentDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_EXP); - - // Create a exponent Node. - auto exponent_op = unary_pw_op_create(afterSubtractionTensor, afterExponentTensor, exponentDesc); - - // Define the reduction descriptor - auto reductionAddDesc = cudnn_frontend::ReductionDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setReductionOp(CUDNN_REDUCE_TENSOR_ADD) - .build(); - - // Create a reduction add Node. - auto reductionAdd_op = - cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_REDUCTION_DESCRIPTOR) - .setxDesc(afterExponentTensor) - .setyDesc(afterAddReductionTensor) - .setreductionDesc(reductionAddDesc) - .build(); - - // Define the division descriptor - auto divisionDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_DIV); - - // Create a subtract Node. - auto division_op = binary_pw_op_create(afterExponentTensor, afterAddReductionTensor, - afterDivisionTensor, divisionDesc); - - ops.push_back(std::move(reductionMax_op)); - ops.push_back(std::move(subtract_op)); - ops.push_back(std::move(exponent_op)); - ops.push_back(std::move(reductionAdd_op)); - ops.push_back(std::move(division_op)); - - return afterDivisionTensor; -} - -static cudnn_frontend::Tensor createDropout(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, - int64_t d, double probability, - cudnnDataType_t tensorType, - std::vector &ops, - cudnn_frontend::Tensor const &prevBlockOutputTensor) { - NVTE_CHECK(ops.size() != 0, "Dropout DAG constructed incorrectly as the first one"); - - int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv}; - int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1}; - - int64_t scale_dim[4] = {1, 1, 1, 1}; - int64_t scale_stride[4] = {1, 1, 1, 1}; - - // mask for the dropout - auto dropoutMaskTensor = tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 200, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - - auto reorder_type = cudnn_frontend::TensorReordering_t::F16x16; - - // after dropout tensor - auto afterDropoutTensor = - cudnn_frontend::TensorBuilder() - .setDim(4, afterBMM1_dim) - .setStride(4, afterBMM1_stride) - .setId(S_ID) - .setAlignment(16) // 16B alignment is needed to run a tensor core engine - .setDataType(tensorType) - .setVirtual(false) - .setByValue(false) - .setReorderType(reorder_type) - .build(); - // scale after dropout - auto scaleDropoutTensor = - tensor_create(tensorType, DROPOUT_CONST_ID, scale_dim, scale_stride, false, - true); // is by value - // after Scale - auto afterScaleTensor = tensor_create(tensorType, VIRTUAL_ID + 201, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - - // Define the reduction descriptor - auto rngDesc = cudnn_frontend::RngDescBuilder() - .setRngDistribution(CUDNN_RNG_DISTRIBUTION_BERNOULLI) - .setBernoulliDistProbability(1.0 - probability) - .build(); - - auto dropoutSeed = - tensor_create(CUDNN_DATA_INT64, DROPOUT_SEED_ID, scale_dim, scale_stride, false, false); - auto dropoutOffset = - tensor_create(CUDNN_DATA_INT64, DROPOUT_OFFSET_ID, scale_dim, scale_stride, false, false); - - // Create a rng Node. - auto rng_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_RNG_DESCRIPTOR) - .setyDesc(dropoutMaskTensor) - .setSeedDesc(dropoutSeed) - .setOffsetDesc(dropoutOffset) - .setRngDesc(rngDesc) - .build(); - - // Define the multiply mask descriptor - auto maskMulDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // Create a multiply mask Node. - auto maskMul_op = binary_pw_op_create(prevBlockOutputTensor, dropoutMaskTensor, - afterDropoutTensor, maskMulDesc); - - // Define the multiply scale descriptor - auto scaleMulDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // Create a multiply mask Node. - auto scaleMul_op = - binary_pw_op_create(afterDropoutTensor, scaleDropoutTensor, afterScaleTensor, scaleMulDesc); - - ops.push_back(std::move(rng_op)); - ops.push_back(std::move(maskMul_op)); - ops.push_back(std::move(scaleMul_op)); - - return afterScaleTensor; -} - -static void createBMM2(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, - NVTE_QKV_Layout layout, cudnnDataType_t tensorType, - std::vector &ops, - cudnn_frontend::Tensor const &prevBlockOutputTensor) { - NVTE_CHECK(ops.size() != 0, "BMM2 op constructed incorrectly as the first one"); - - int64_t seqlen_dim[4] = {b, 1, 1, 1}; - int64_t seqlen_stride[4] = {1, 1, 1, 1}; - - int64_t v_dim[4] = {b, h, s_kv, d}; - int64_t v_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, v_stride, layout, NVTE_QKV_Matrix::NVTE_V_Matrix); - - int64_t o_dim[4] = {b, h, s_q, d}; - int64_t o_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, o_stride, layout, NVTE_QKV_Matrix::NVTE_O_Matrix); - - auto seqlenQTensor = - tensor_create(CUDNN_DATA_INT32, Q_SEQLEN_ID, seqlen_dim, seqlen_stride, false, false); - auto seqlenKTensor = - tensor_create(CUDNN_DATA_INT32, K_SEQLEN_ID, seqlen_dim, seqlen_stride, false, false); - auto vTensor = tensor_create(tensorType, V_ID, v_dim, v_stride, false, false); - // second GEMM output - auto oTensor = tensor_create(tensorType, O_ID, o_dim, o_stride, false, false); - - // Define the matmul 2 desc - // set padding value optionally to 0 for writing zeros to O tensor (if not set, old behaviour) - auto matmul_2_Desc = cudnn_frontend::MatMulDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setPaddingValue(0.0f) - .build(); - - // Create a matmul 2 Node - auto matmul_op2 = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_MATMUL_DESCRIPTOR) - .setaMatDesc(prevBlockOutputTensor) - .setbMatDesc(vTensor) - .setcMatDesc(oTensor) - .setmOverrideDesc(seqlenQTensor) - .setkOverrideDesc(seqlenKTensor) - .setmatmulDesc(matmul_2_Desc) - .build(); - - ops.push_back(std::move(matmul_op2)); -} - -static cudnn_frontend::Tensor createSoftmaxBackward(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, - int64_t d, NVTE_QKV_Layout layout, - cudnnDataType_t tensorType, - std::vector &ops, - cudnn_frontend::Tensor const &yTensor, - cudnn_frontend::Tensor const &dyTensor) { - NVTE_CHECK(ops.size() != 0, "Softmax backward constructed incorrectly as the first one"); - - int64_t p_dim[4] = {b, h, s_q, s_kv}; - int64_t p_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, p_stride, layout, NVTE_QKV_Matrix::NVTE_S_Matrix); - - int64_t p_reduction_dim[4] = {b, h, s_q, 1}; - int64_t p_reduction_stride[4]; - - p_reduction_stride[3] = 1; - p_reduction_stride[2] = 1; - p_reduction_stride[1] = s_q; - p_reduction_stride[0] = s_q * h; - - int64_t const_dim[4] = {1, 1, 1, 1}; - int64_t const_stride[4] = {1, 1, 1, 1}; - - // creating all tensors - auto softmaxScaleTensor = - tensor_create(CUDNN_DATA_FLOAT, S_CONST_ID, const_dim, const_stride, false, true); - auto dyMulYTensor = - tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 250, p_dim, p_stride, true, false); - auto dxAfterReductionTensor = tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 251, p_reduction_dim, - p_reduction_stride, true, false); - auto dxAfterSubtractionTensor = - tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 252, p_dim, p_stride, true, false); - auto dxUnscaleTensor = - tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 253, p_dim, p_stride, true, false); - auto dxTensor = tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 254, p_dim, p_stride, true, false); - - // creating all ops - // mul (y * dy) - auto mul_1_desc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - auto mul_1_op = binary_pw_op_create(yTensor, dyTensor, dyMulYTensor, mul_1_desc); - - // reduction add sum (y * dy) - auto reductionAddDesc = cudnn_frontend::ReductionDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setReductionOp(CUDNN_REDUCE_TENSOR_ADD) - .build(); - - auto reductionAdd_op = - cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_REDUCTION_DESCRIPTOR) - .setxDesc(dyMulYTensor) - .setyDesc(dxAfterReductionTensor) - .setreductionDesc(reductionAddDesc) - .build(); - - // subtraction (dy - sum(y * dy)) - auto sub_0_desc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_SUB); - auto sub_0_op = - binary_pw_op_create(dyTensor, dxAfterReductionTensor, dxAfterSubtractionTensor, sub_0_desc); - - // mul (y * (dy - sum(y * dy))) - auto mul_2_desc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - auto mul_2_op = - binary_pw_op_create(yTensor, dxAfterSubtractionTensor, dxUnscaleTensor, mul_2_desc); - - // mul (scale * dx) - auto mul_3_desc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - auto mul_3_op = binary_pw_op_create(dxUnscaleTensor, softmaxScaleTensor, dxTensor, mul_3_desc); - - ops.push_back(std::move(mul_1_op)); - ops.push_back(std::move(reductionAdd_op)); - ops.push_back(std::move(sub_0_op)); - ops.push_back(std::move(mul_2_op)); - ops.push_back(std::move(mul_3_op)); - - return dxTensor; -} - -void fused_attn_max_512_fwd_impl( - int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, bool is_training, - float scaling_factor, float dropout_probability, NVTE_QKV_Layout layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, void *devPtrQ, void *devPtrK, void *devPtrV, - void *devPtrS, void *devPtrO, void *devPtrBias, void *devPtrCuSeqlenQ, void *devPtrCuSeqlenKV, - void *devPtrDropoutSeed, void *devPtrDropoutOffset, void *workspace, size_t *workspace_size, - cudnnDataType_t tensorType, cudaStream_t stream, cudnnHandle_t handle) { - try { - FADescriptor descriptor{b, h, - s_q, s_kv, - d, scaling_factor, - is_training, dropout_probability, - layout, bias_type, - mask_type, tensorType, - false}; - - using CacheType = std::map; - static thread_local CacheType fmha_fprop_cache; - - // softmax auxiliary is only used in the training mode - bool enable_dropout = is_training && (dropout_probability != 0.0f); - - // two conditions that make softmax auxiliary in virtual - // 1. inference mode (not is_training) - // 2. dropout enabled: the auxiliary becomes the dropout output - bool softmax_output_virtual = !is_training || enable_dropout; - - // Get plan from cache if cache is available, otherwise create one - auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) { - // if hit, return - auto it = cache.find(descriptor); - if (it != cache.end()) { - auto plan = it->second; - return plan; - } - - // otherwise, build the op_graph and the plan. Then update cache - std::vector all_ops; - std::vector ops; - - createScale(b, h, s_q, s_kv, d, layout, tensorType, ops); - - // if bias, we need to memset the S buffer to correctly computate dbias - // WAR: causal_mask without bias needs memset the S buffer - // inference mode doesn't need the S auxiliary - auto zero_s = (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) || - (mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - (mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)) && - is_training; - std::shared_ptr maskInput; - auto bmm1_output = createBMM1(b, h, s_q, s_kv, d, layout, tensorType, zero_s, ops); - - NVTE_CHECK(bias_type != NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS, - "NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS has not been implemented."); - - if (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS) { - auto bias_output = createBias(b, h, s_q, s_kv, d, layout, tensorType, ops, bmm1_output); - maskInput = std::make_shared(std::move(bias_output)); - } - if (bias_type == NVTE_Bias_Type::NVTE_NO_BIAS) { - maskInput = std::make_shared(std::move(bmm1_output)); - } - - auto mask_output = createMask(b, h, s_q, s_kv, d, layout, mask_type, tensorType, ops, - *maskInput.get(), false); - - NVTE_CHECK(dropout_probability != 1.0f, "Dropout probability cannot be 1.0."); - - auto softmax_output = - createSoftmaxForward(b, h, s_q, s_kv, d, layout, enable_dropout, softmax_output_virtual, - tensorType, ops, mask_output); - - if (enable_dropout) { - auto dropout_output = - createDropout(b, h, s_q, s_kv, d, dropout_probability, tensorType, ops, softmax_output); - createBMM2(b, h, s_q, s_kv, d, layout, tensorType, ops, dropout_output); - } else { - createBMM2(b, h, s_q, s_kv, d, layout, tensorType, ops, softmax_output); - } - - for (unsigned int i = 0; i < ops.size(); i++) { - all_ops.push_back(&ops[i]); - } - - // Create an Operation Graph - auto opGraph = cudnn_frontend::OperationGraphBuilder() - .setHandle(handle) - .setOperationGraph(all_ops.size(), all_ops.data()) - .build(); - - cudnn_frontend::EngineConfigList filtered_configs; - auto statuses = cudnn_frontend::get_heuristics_list<1>( - {"heuristics_instant"}, opGraph, allowAllConfig, filtered_configs, true); - - if (filtered_configs.size() == 0) { - cudnn_frontend::set_error_and_throw_exception( - nullptr, CUDNN_STATUS_NOT_SUPPORTED, - "run_mha_fprop: No config returned by the heuristics"); - } - auto plan = cudnn_frontend::ExecutionPlanBuilder() - .setHandle(handle) - .setEngineConfig(filtered_configs[0], opGraph.getTag()) - .build(); - cache.insert({descriptor, plan}); - return plan; - }; - - auto plan = get_plan(fmha_fprop_cache, descriptor); - - auto plan_workspace_size = plan.getWorkspaceSize(); - - // Exit to request upper level API to allocate memory if needed - if (workspace == nullptr) { - size_t actual_seqlen_workspace_size = 2 * b * sizeof(int32_t); - *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; - return; - } - - // cuDNN stream check needs to be moved here to support dummy kernel calls with - // null streams for sizing the cuDNN workspace. - NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); - - // Prepare actual seqlen - constexpr size_t nthreads_per_block = 128; - const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; - void *devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void *devActualSeqlenK = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); - cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrCuSeqlenQ), - static_cast(devPtrCuSeqlenKV), static_cast(devActualSeqlenQ), - static_cast(devActualSeqlenK)); - NVTE_CHECK_CUDA(cudaGetLastError()); - - // change this if you have access to float_min - float negInfinity = -1.0E+10; - float scale_dropout = 1 / (1 - dropout_probability); - - std::set> data_ptrs; - // add all the data pointers to be used in the variant pack - data_ptrs.insert(std::pair(Q_ID, devPtrQ)); - data_ptrs.insert(std::pair(K_ID, devPtrK)); - data_ptrs.insert(std::pair(V_ID, devPtrV)); - data_ptrs.insert(std::pair(Q_SEQLEN_ID, devActualSeqlenQ)); - data_ptrs.insert(std::pair(K_SEQLEN_ID, devActualSeqlenK)); - data_ptrs.insert(std::pair(MASK_VAL_ID, &negInfinity)); - - __half half_cast_scaling_factor{scaling_factor}; - __nv_bfloat16 bfloat_cast_scaling_factor{scaling_factor}; - - if (tensorType == CUDNN_DATA_FLOAT) { - data_ptrs.insert(std::pair(S_CONST_ID, &scaling_factor)); - } else if (tensorType == CUDNN_DATA_HALF) { - data_ptrs.insert(std::pair(S_CONST_ID, &half_cast_scaling_factor)); - } else if (tensorType == CUDNN_DATA_BFLOAT16) { - data_ptrs.insert(std::pair(S_CONST_ID, &bfloat_cast_scaling_factor)); - } else { - NVTE_ERROR("Unsupported tensor type."); - } - - data_ptrs.insert(std::pair(O_ID, devPtrO)); - - if (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) { - data_ptrs.insert(std::pair(B_ID, devPtrBias)); - } - - // if enable_dropout, S is the result after dropout - // if not enable dropout, S is the result after softmax - if (enable_dropout || !softmax_output_virtual) { - data_ptrs.insert(std::pair(S_ID, devPtrS)); - } - - __half half_cast_scale_dropout{scale_dropout}; - __nv_bfloat16 bfloat16_cast_scale_dropout{scale_dropout}; - - if (enable_dropout) { - // TODO(rewang): make a util func - if (tensorType == CUDNN_DATA_FLOAT) { - data_ptrs.insert(std::pair(DROPOUT_CONST_ID, &scale_dropout)); - } else if (tensorType == CUDNN_DATA_HALF) { - data_ptrs.insert(std::pair(DROPOUT_CONST_ID, &half_cast_scale_dropout)); - } else if (tensorType == CUDNN_DATA_BFLOAT16) { - data_ptrs.insert( - std::pair(DROPOUT_CONST_ID, &bfloat16_cast_scale_dropout)); - } else { - NVTE_ERROR("Unsupported tensor type."); - } - data_ptrs.insert(std::pair(DROPOUT_SEED_ID, devPtrDropoutSeed)); - data_ptrs.insert(std::pair(DROPOUT_OFFSET_ID, devPtrDropoutOffset)); - } - - auto variantPack = cudnn_frontend::VariantPackBuilder() - .setWorkspacePointer(workspace) - .setDataPointers(data_ptrs) - .build(); - - NVTE_CHECK_CUDNN(cudnnBackendExecute(handle, plan.get_raw_desc(), variantPack.get_raw_desc())); - } catch (cudnn_frontend::cudnnException &e) { - NVTE_ERROR(e.what()); - } -} - -void fused_attn_max_512_bwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, - float scaling_factor, float dropout_probability, - NVTE_QKV_Layout layout, NVTE_Mask_Type mask_type, - NVTE_Bias_Type bias_type, void *devPtrQ, void *devPtrK, - void *devPtrV, void *devPtrS, void *devPtrdQ, void *devPtrdK, - void *devPtrdV, void *devPtrdO, void *devPtrdS, void *devPtrdBias, - void *devPtrCuSeqlenQ, void *devPtrCuSeqlenKV, void *workspace, - size_t *workspace_size, cudnnDataType_t tensorType, - cudaStream_t stream, cudnnHandle_t handle) { - try { - FADescriptor descriptor{ - b, h, s_q, s_kv, d, scaling_factor, true, dropout_probability, - layout, bias_type, mask_type, tensorType, false}; - - using CacheType = std::map; - static thread_local CacheType fmha_bprop_cache; - - auto get_plan = [&](CacheType &cache, const FADescriptor &descriptor) { - auto it = cache.find(descriptor); - if (it != cache.end()) { - return it->second; - } - - std::vector all_ops; - std::vector ops; - - // Creates the necessary tensor descriptors - int64_t q_dim[4] = {b, h, s_q, d}; - int64_t q_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, q_stride, layout, NVTE_QKV_Matrix::NVTE_Q_Matrix); - - int64_t k_dim[4] = {b, h, s_kv, d}; - int64_t k_stride[4]; - generateMatrixStrides( - b, h, s_q, s_kv, d, k_stride, layout, - NVTE_QKV_Matrix::NVTE_K_Matrix); // type is correct as K is not transposed - - int64_t v_dim[4] = {b, h, d, s_kv}; - int64_t v_stride[4]; - generateMatrixStrides( - b, h, s_q, s_kv, d, v_stride, layout, - NVTE_QKV_Matrix::NVTE_V_Matrix_Transpose); // type is correct as V is transposed - - int64_t p_dim[4] = {b, h, s_q, s_kv}; - int64_t p_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, p_stride, layout, NVTE_QKV_Matrix::NVTE_S_Matrix); - - int64_t p_transpose_dim[4] = {b, h, s_kv, s_q}; - int64_t p_transpose_stride[4]; - p_transpose_stride[0] = p_stride[0]; - p_transpose_stride[1] = p_stride[1]; - p_transpose_stride[2] = p_stride[3]; - p_transpose_stride[3] = p_stride[2]; - - int64_t o_dim[4] = {b, h, s_q, d}; - int64_t o_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, o_stride, layout, NVTE_QKV_Matrix::NVTE_O_Matrix); - - int64_t seqlen_dim[4] = {b, 1, 1, 1}; - int64_t seqlen_stride[4] = {1, 1, 1, 1}; - - int64_t scale_dim[4] = {1, 1, 1, 1}; - int64_t scale_stride[4] = {1, 1, 1, 1}; - - // inputs to fprop - auto qTensor = tensor_create(tensorType, Q_ID, q_dim, q_stride, false, false); - auto kTensor = tensor_create(tensorType, K_ID, k_dim, k_stride, false, false); - auto vTensor = tensor_create(tensorType, V_ID, v_dim, v_stride, false, false); - auto seqlenQTensor = - tensor_create(CUDNN_DATA_INT32, Q_SEQLEN_ID, seqlen_dim, seqlen_stride, false, false); - auto seqlenKTensor = - tensor_create(CUDNN_DATA_INT32, K_SEQLEN_ID, seqlen_dim, seqlen_stride, false, false); - - // gradient of the output - auto doTensor = tensor_create(tensorType, dO_ID, o_dim, o_stride, false, false); - - auto reorder_type = cudnn_frontend::TensorReordering_t::F16x16; - - // activation from fprop - auto pTensor = cudnn_frontend::TensorBuilder() - .setDim(4, p_dim) - .setStride(4, p_stride) - .setId(S_ID) - .setAlignment(16) // 16B alignment is needed to run a tensor core engine - .setDataType(tensorType) - .setVirtual(false) - .setByValue(false) - .setReorderType(reorder_type) - .build(); - - // outputs from bprop - auto dqTensor = tensor_create(tensorType, dQ_ID, q_dim, q_stride, false, false); - auto dkTensor = tensor_create(tensorType, dK_ID, k_dim, k_stride, false, false); - auto dvTensor = tensor_create(tensorType, dV_ID, k_dim, k_stride, false, - false); // not transposed therefore k_dim and k_stride - - //////////////////////////////////////////////////////// - // start creating the ops and the intermediate tensors - auto pReshapeTensor = tensor_create(tensorType, VIRTUAL_ID + 300, p_transpose_dim, - p_transpose_stride, true, false); - - // reshape to perform transpose and make pReshape - auto reshape_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_RESHAPE_DESCRIPTOR) - .setxDesc(pTensor) - .setyDesc(pReshapeTensor) - .build(); - - ops.push_back(std::move(reshape_op)); - - // scale dropout - auto dropoutScaleTensor = tensor_create(CUDNN_DATA_FLOAT, DROPOUT_CONST_ID, scale_dim, - scale_stride, false, true); // is by value - auto pAfterScaleTensor = tensor_create(tensorType, VIRTUAL_ID + 301, p_transpose_dim, - p_transpose_stride, true, false); - - auto scaleMulDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - auto scaleMul_op = - binary_pw_op_create(pReshapeTensor, dropoutScaleTensor, pAfterScaleTensor, scaleMulDesc); - ops.push_back(std::move(scaleMul_op)); - - // perform absolute operation to remove the mask bit - auto pTransposeAfterAbsTensor = tensor_create(tensorType, VIRTUAL_ID + 302, p_transpose_dim, - p_transpose_stride, true, false); - - auto absDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_ABS); - auto abs_op = unary_pw_op_create(pAfterScaleTensor, pTransposeAfterAbsTensor, absDesc); - ops.push_back(std::move(abs_op)); - - // matmul to calculate dvTensor - // set padding value optionally to 0 for writing zeros to dV tensor (if not set, old - // behaviour) - auto matmul_0_Desc = cudnn_frontend::MatMulDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setPaddingValue(0.0f) - .build(); - - auto matmul_op0 = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_MATMUL_DESCRIPTOR) - .setaMatDesc(pTransposeAfterAbsTensor) - .setbMatDesc(doTensor) - .setcMatDesc(dvTensor) - .setmOverrideDesc(seqlenKTensor) - .setkOverrideDesc(seqlenQTensor) - .setmatmulDesc(matmul_0_Desc) - .build(); - - ops.push_back(std::move(matmul_op0)); - - // matmul to calculate dpTensor - auto dpTensor = - tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 303, p_dim, p_stride, true, false); - - auto matmul_1_Desc = - cudnn_frontend::MatMulDescBuilder().setComputeType(CUDNN_DATA_FLOAT).build(); - - auto matmul_op1 = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_MATMUL_DESCRIPTOR) - .setaMatDesc(doTensor) - .setbMatDesc(vTensor) - .setcMatDesc(dpTensor) - .setmOverrideDesc(seqlenQTensor) - .setnOverrideDesc(seqlenKTensor) - .setmatmulDesc(matmul_1_Desc) - .build(); - - ops.push_back(std::move(matmul_op1)); - - // mask the values which were dropped in dropout - auto pAbsTensor = tensor_create(tensorType, VIRTUAL_ID + 304, p_dim, p_stride, true, false); - - auto p_absDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_ABS); - auto p_abs_op = unary_pw_op_create(pTensor, pAbsTensor, p_absDesc); - ops.push_back(std::move(p_abs_op)); - - // create the dropout mask - auto zeroTensor = tensor_create(CUDNN_DATA_FLOAT, MASK_VAL_ID, scale_dim, scale_stride, false, - true); // is by value - auto dropoutMaskTensor = - tensor_create(CUDNN_DATA_BOOLEAN, VIRTUAL_ID + 305, p_dim, p_stride, true, false); - - auto greater_than_0_desc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_CMP_GT); - auto greater_than_0_op = - binary_pw_op_create(pTensor, zeroTensor, dropoutMaskTensor, greater_than_0_desc); - ops.push_back(std::move(greater_than_0_op)); - - // scale for the dropout - auto dpAfterScaleTensor = - tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 306, p_dim, p_stride, true, false); - - auto mul_0_desc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - auto mul_0_op = - binary_pw_op_create(dpTensor, dropoutScaleTensor, dpAfterScaleTensor, mul_0_desc); - ops.push_back(std::move(mul_0_op)); - - // drop the values based on the dropout mask - auto dpAfterDropoutTensor = - tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 307, p_dim, p_stride, true, false); - - auto selection_0_desc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_BINARY_SELECT); - auto selection_0_op = ternary_pw_op_create(dpAfterScaleTensor, zeroTensor, dropoutMaskTensor, - dpAfterDropoutTensor, selection_0_desc); - ops.push_back(std::move(selection_0_op)); - - // softmax backward - auto dsTensor = createSoftmaxBackward(b, h, s_q, s_kv, d, layout, tensorType, ops, pAbsTensor, - dpAfterDropoutTensor); - - // mask - auto dsAfterMaskTensor = - createMask(b, h, s_q, s_kv, d, layout, mask_type, tensorType, ops, dsTensor, true); - - // dbias tensor - int64_t dbias_dim[4] = {1, h, s_q, s_kv}; - int64_t dbias_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1}; - auto dBiasTensor = tensor_create(tensorType, dBias_ID, dbias_dim, dbias_stride, false, false); - - if (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS) { - auto softmaxScaleTensor = - tensor_create(CUDNN_DATA_FLOAT, S_CONST_ID, scale_dim, scale_stride, false, true); - auto softmaxScaleReciprocalTensor = - tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 401, scale_dim, scale_stride, true, false); - auto dbiasBeforeScaleTensor = - tensor_create(CUDNN_DATA_FLOAT, VIRTUAL_ID + 402, dbias_dim, dbias_stride, true, false); - - // Define the reduction descriptor - auto reductionAddDesc = cudnn_frontend::ReductionDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setReductionOp(CUDNN_REDUCE_TENSOR_ADD) - .build(); - - // Create a reduction add node to compute the dbias - auto reductionAdd_op = - cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_REDUCTION_DESCRIPTOR) - .setxDesc(dsAfterMaskTensor) - .setyDesc(dbiasBeforeScaleTensor) - .setreductionDesc(reductionAddDesc) - .build(); - ops.push_back(std::move(reductionAdd_op)); - - // take the reciprocal of the scale - auto reciprocal_scale_desc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_RECIPROCAL); - auto reciprocal_scale_op = unary_pw_op_create( - softmaxScaleTensor, softmaxScaleReciprocalTensor, reciprocal_scale_desc); - ops.push_back(std::move(reciprocal_scale_op)); - - // apply the scale - auto dBias_scale_desc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - auto dBias_scale_op = binary_pw_op_create( - dbiasBeforeScaleTensor, softmaxScaleReciprocalTensor, dBiasTensor, dBias_scale_desc); - ops.push_back(std::move(dBias_scale_op)); - } - - // matmul to calculate dqTensor - // set padding value optionally to 0 for writing zeros to dqTensor (if not set, old - // behaviour) - auto matmul_2_Desc = cudnn_frontend::MatMulDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setPaddingValue(0.0f) - .build(); - - auto matmul_op2 = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_MATMUL_DESCRIPTOR) - .setaMatDesc(dsAfterMaskTensor) - .setbMatDesc(kTensor) - .setcMatDesc(dqTensor) - .setmOverrideDesc(seqlenQTensor) - .setkOverrideDesc(seqlenKTensor) - .setmatmulDesc(matmul_2_Desc) - .build(); - - ops.push_back(std::move(matmul_op2)); - - // reshape for transpose of ds - auto dsAfterMaskReshapeTensor = tensor_create(tensorType, VIRTUAL_ID + 308, p_transpose_dim, - p_transpose_stride, true, false); - - auto reshape_2_op = - cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_RESHAPE_DESCRIPTOR) - .setxDesc(dsAfterMaskTensor) - .setyDesc(dsAfterMaskReshapeTensor) - .build(); - - ops.push_back(std::move(reshape_2_op)); - - // matmul to calculate dkTensor - // set padding value optionally to 0 for writing zeros to dktensor (if not set, old - // behaviour) - auto matmul_3_Desc = cudnn_frontend::MatMulDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setPaddingValue(0.0f) - .build(); - - auto matmul_op3 = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_MATMUL_DESCRIPTOR) - .setaMatDesc(dsAfterMaskReshapeTensor) - .setbMatDesc(qTensor) - .setcMatDesc(dkTensor) - .setmOverrideDesc(seqlenKTensor) - .setkOverrideDesc(seqlenQTensor) - .setmatmulDesc(matmul_3_Desc) - .build(); - - ops.push_back(std::move(matmul_op3)); - - ///////////////////////////////////////////////////////////////// - - for (unsigned int i = 0; i < ops.size(); i++) { - all_ops.push_back(&ops[i]); - } - - // Create an Operation Graph - auto opGraph = cudnn_frontend::OperationGraphBuilder() - .setHandle(handle) - .setOperationGraph(all_ops.size(), all_ops.data()) - .build(); - - cudnn_frontend::EngineConfigList filtered_configs; - auto statuses = cudnn_frontend::get_heuristics_list<1>( - {"heuristics_instant"}, opGraph, allowAllConfig, filtered_configs, true); - - if (filtered_configs.size() == 0) { - cudnn_frontend::set_error_and_throw_exception( - nullptr, CUDNN_STATUS_NOT_SUPPORTED, - "run_mha_bprop: No config returned by the heuristics"); - } - - auto plan = cudnn_frontend::ExecutionPlanBuilder() - .setHandle(handle) - .setEngineConfig(filtered_configs[0], opGraph.getTag()) - .build(); - cache.insert({descriptor, plan}); - return plan; - }; - - auto plan = get_plan(fmha_bprop_cache, descriptor); - - auto plan_workspace_size = plan.getWorkspaceSize(); - - // Exit to request upper level API to allocate memory if needed - if (workspace == nullptr) { - size_t actual_seqlen_workspace_size = 2 * b * sizeof(int32_t); - *workspace_size = plan_workspace_size + actual_seqlen_workspace_size; - return; - } - - // cuDNN stream check needs to be moved here to support dummy kernel calls with - // null streams for sizing the cuDNN workspace. - NVTE_CHECK_CUDNN(cudnnSetStream(handle, stream)); - - constexpr size_t nthreads_per_block = 128; - const size_t grid = (b + nthreads_per_block - 1) / nthreads_per_block; - void *devActualSeqlenQ = static_cast(workspace) + plan_workspace_size; - void *devActualSeqlenK = static_cast(devActualSeqlenQ) + b * sizeof(int32_t); - cu_seqlens_to_actual_seqlens<<>>( - b, b, static_cast(devPtrCuSeqlenQ), - static_cast(devPtrCuSeqlenKV), static_cast(devActualSeqlenQ), - static_cast(devActualSeqlenK)); - NVTE_CHECK_CUDA(cudaGetLastError()); - - std::set> data_ptrs; - // add all the data pointers to be used in the variant pack - data_ptrs.insert(std::pair(dQ_ID, devPtrdQ)); - data_ptrs.insert(std::pair(dK_ID, devPtrdK)); - data_ptrs.insert(std::pair(dV_ID, devPtrdV)); - - data_ptrs.insert(std::pair(Q_ID, devPtrQ)); - data_ptrs.insert(std::pair(K_ID, devPtrK)); - data_ptrs.insert(std::pair(V_ID, devPtrV)); - data_ptrs.insert(std::pair(S_ID, devPtrS)); - data_ptrs.insert(std::pair(dO_ID, devPtrdO)); - data_ptrs.insert(std::pair(dS_ID, devPtrdS)); - data_ptrs.insert(std::pair(Q_SEQLEN_ID, devActualSeqlenQ)); - data_ptrs.insert(std::pair(K_SEQLEN_ID, devActualSeqlenK)); - - if (bias_type != NVTE_Bias_Type::NVTE_NO_BIAS) { - data_ptrs.insert(std::pair(dBias_ID, devPtrdBias)); - } - - float zeroVal = 0.0f; - float dropoutScale = 1.0f / (1.0f - dropout_probability); - - data_ptrs.insert(std::pair(DROPOUT_CONST_ID, &dropoutScale)); - data_ptrs.insert(std::pair(S_CONST_ID, &scaling_factor)); - data_ptrs.insert(std::pair(MASK_VAL_ID, &zeroVal)); - - auto variantPack = cudnn_frontend::VariantPackBuilder() - .setWorkspacePointer(workspace) - .setDataPointers(data_ptrs) - .build(); - - NVTE_CHECK_CUDNN(cudnnBackendExecute(handle, plan.get_raw_desc(), variantPack.get_raw_desc())); - } catch (cudnn_frontend::cudnnException &e) { - NVTE_ERROR(e.what()); - } -} - -} // namespace fused_attn - -using namespace transformer_engine::fused_attn; -void fused_attn_max_512_fwd(size_t batch, size_t num_head, size_t q_max_seqlen, - size_t kv_max_seqlen, size_t head_dim, bool is_training, - float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_Bias, Tensor *output_O, - NVTETensorPack *Aux_CTX_Tensors, const Tensor *q_cu_seqlens, - const Tensor *kv_cu_seqlens, const Tensor *rng_state, Tensor *workspace, - cudaStream_t stream, cudnnHandle_t handle) { - using namespace transformer_engine; - - void *devPtrQ = input_Q->data.dptr; - void *devPtrK = input_K->data.dptr; - void *devPtrV = input_V->data.dptr; - - void *devPtrBias = input_Bias->data.dptr; - - void *devPtrO = output_O->data.dptr; - - void *devPtrS = nullptr; - - const DType q_type = input_Q->data.dtype; - const DType kv_type = input_K->data.dtype; - NVTE_CHECK(q_type == kv_type, "data type of Q must be equal to data type of KV."); - - if (Aux_CTX_Tensors->size == 0) { - Aux_CTX_Tensors->size = 1; - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[0]); - output_S->data.dptr = nullptr; - output_S->data.shape = {batch, num_head, q_max_seqlen, kv_max_seqlen}; - output_S->data.dtype = q_type; - } else if (Aux_CTX_Tensors->size == 1) { - Tensor *output_S = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[0]); - devPtrS = output_S->data.dptr; - } else { - NVTE_ERROR("Unexpected Aux_CTX_Tensors->size."); - } - - void *devQCuSeqlen = q_cu_seqlens->data.dptr; - void *devKVCuSeqlen = kv_cu_seqlens->data.dptr; - - const DType rng_state_type = rng_state->data.dtype; - NVTE_CHECK(rng_state_type == DType::kInt64); - void *devPtrDropoutSeed = rng_state->data.dptr; - void *devPtrDropoutOffset = - static_cast(static_cast(rng_state->data.dptr) + 1); - - size_t workspace_size = 0; - - fused_attn_max_512_fwd_impl( - batch, num_head, q_max_seqlen, kv_max_seqlen, head_dim, is_training, attn_scale, p_dropout, - qkv_layout, bias_type, mask_type, devPtrQ, devPtrK, devPtrV, devPtrS, devPtrO, devPtrBias, - devQCuSeqlen, devKVCuSeqlen, devPtrDropoutSeed, devPtrDropoutOffset, workspace->data.dptr, - &workspace_size, get_cudnn_dtype(q_type), stream, handle); - - if (workspace_size > 0) { - if (workspace->data.dptr == nullptr) { - workspace->data.shape = {workspace_size}; - workspace->data.dtype = DType::kByte; - return; - } - } else if (workspace_size == 0) { - workspace->data.shape = {1}; - workspace->data.dtype = DType::kByte; - return; - } else { - NVTE_ERROR("Unexpected workspace_size."); - } -} - -void fused_attn_max_512_bwd(size_t batch, size_t num_head, size_t q_max_seqlen, - size_t kv_max_seqlen, size_t head_dim, float attn_scale, - float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, const Tensor *input_Q, const Tensor *input_K, - const Tensor *input_V, const Tensor *input_dO, Tensor *output_S, - Tensor *output_dQ, Tensor *output_dK, Tensor *output_dV, - Tensor *output_dBias, const Tensor *q_cu_seqlens, - const Tensor *kv_cu_seqlens, Tensor *workspace, cudaStream_t stream, - cudnnHandle_t handle) { - using namespace transformer_engine; - - void *devPtrQ = input_Q->data.dptr; - void *devPtrK = input_K->data.dptr; - void *devPtrV = input_V->data.dptr; - - void *devPtrdO = input_dO->data.dptr; - - void *devPtrdQ = output_dQ->data.dptr; - void *devPtrdK = output_dK->data.dptr; - void *devPtrdV = output_dV->data.dptr; - - void *devPtrdBias = output_dBias->data.dptr; - - void *devPtrS = output_S->data.dptr; - - // devPtrdS reuses the memory of devPtrS - void *devPtrdS = devPtrS; - - void *devPtrQCuSeqlens = q_cu_seqlens->data.dptr; - void *devPtrKVCuSeqlens = kv_cu_seqlens->data.dptr; - - const auto q_type = input_Q->data.dtype; - const auto kv_type = input_K->data.dtype; - NVTE_CHECK(q_type == kv_type, "data type of Q must be equal to data type of KV."); - size_t workspace_size = 0; - - fused_attn_max_512_bwd_impl( - batch, num_head, q_max_seqlen, kv_max_seqlen, head_dim, attn_scale, p_dropout, qkv_layout, - mask_type, bias_type, devPtrQ, devPtrK, devPtrV, devPtrS, devPtrdQ, devPtrdK, devPtrdV, - devPtrdO, devPtrdS, devPtrdBias, devPtrQCuSeqlens, devPtrKVCuSeqlens, workspace->data.dptr, - &workspace_size, get_cudnn_dtype(q_type), stream, handle); - - if (workspace_size > 0) { - if (workspace->data.dptr == nullptr) { - workspace->data.shape = {workspace_size}; - workspace->data.dtype = DType::kByte; - return; - } - } else if (workspace_size == 0) { - workspace->data.shape = {1}; - workspace->data.dtype = DType::kByte; - return; - } else { - NVTE_ERROR("Unexpected workspace_size."); - } -} -} // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.h b/transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.h deleted file mode 100644 index 1e59d4dc8..000000000 --- a/transformer_engine/common/fused_attn/fused_attn_f16_max512_seqlen.h +++ /dev/null @@ -1,41 +0,0 @@ -/************************************************************************* - * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. - * - * See LICENSE for license information. - ************************************************************************/ - -/*! \file fused_attn_fp16_bf16_max_seqlen_512.h - * \brief Functions for fused attention for half precision with seqlen <= 512 - */ - -#ifndef TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_MAX_512_H_ -#define TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_MAX_512_H_ - -#include - -#include "common/common.h" -#include "transformer_engine/fused_attn.h" - -namespace transformer_engine { -void fused_attn_max_512_fwd(size_t batch, size_t num_head, size_t q_max_seqlen, - size_t kv_max_seqlen, size_t head_dim, bool is_training, - float attn_scale, float p_dropout, NVTE_QKV_Layout qkv_layout, - NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, - const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, - const Tensor *input_Bias, Tensor *output_O, - NVTETensorPack *Aux_CTX_Tensors, const Tensor *q_cu_seqlens, - const Tensor *kv_cu_seqlens, const Tensor *rng_state, Tensor *workspace, - cudaStream_t stream, cudnnHandle_t handle); - -void fused_attn_max_512_bwd(size_t batch, size_t num_head, size_t q_max_seqlen, - size_t kv_max_seqlen, size_t head_dim, float attn_scale, - float p_dropout, NVTE_QKV_Layout qkv_layout, NVTE_Bias_Type bias_type, - NVTE_Mask_Type mask_type, const Tensor *input_Q, const Tensor *input_K, - const Tensor *input_V, const Tensor *input_dO, Tensor *output_S, - Tensor *output_dQ, Tensor *output_dK, Tensor *output_dV, - Tensor *output_dBias, const Tensor *q_cu_seqlens, - const Tensor *kv_cu_seqlens, Tensor *workspace, cudaStream_t stream, - cudnnHandle_t handle); -} // namespace transformer_engine - -#endif // TRANSFORMER_ENGINE_COMMON_FUSED_ATTN_FUSED_ATTN_MAX_512_H_ diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index 912dc32d3..d301be573 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -156,11 +156,9 @@ enum NVTE_Softmax_Type { enum NVTE_Fused_Attn_Backend { /*! No supported backend */ NVTE_No_Backend = -1, - /*! cuDNN-based FP16/BF16 fused attention for <= 512 sequence length */ - NVTE_F16_max512_seqlen = 0, /*! cuDNN-based FP16/BF16 fused attention for any sequence length */ NVTE_F16_arbitrary_seqlen = 1, - /*! cuDNN-based FP8 fused attention for <= 512 sequence length */ + /*! cuDNN-based FP8 fused attention */ NVTE_FP8 = 2, }; @@ -236,8 +234,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * Support Matrix: \verbatim | backend | precision | qkv layout | bias | mask | dropout | sequence length | head_dim | - | 0 | FP16/BF16 | BS3HD,SB3HD,BSHD_BS2HD,SBHD_SB2HD | NO/POST_SCALE_BIAS | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK | Yes | <= 512, % 64 == 0 | 64 | - | 1 | FP16/BF16 | BS3HD,SB3HD,BSH3D,SBH3D | NO/POST_SCALE_BIAS/ALIBI | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK | Yes | > 512, % 64 == 0 | <= 128, % 8 == 0 | + | 1 | FP16/BF16 | BS3HD,SB3HD,BSH3D,SBH3D | NO/POST_SCALE_BIAS/ALIBI | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK | Yes | any, % 64 == 0 | <= 128, % 8 == 0 | | | | BSHD_BS2HD,BSHD_BSH2D,SBHD_SB2HD,SBHD_SBH2D | | | | | | | | | BSHD_BSHD_BSHD,SBHD_SBHD_SBHD | | | | | | | 2 | FP8 | T3HD | NO_BIAS | PADDING_MASK | Yes | <= 512, % 64 == 0 | 64 | @@ -314,8 +311,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso * Support Matrix: \verbatim | backend | precision | qkv layout | bias | mask | dropout | sequence length | head_dim | - | 0 | FP16/BF16 | BS3HD,SB3HD,BSHD_BS2HD,SBHD_SB2HD | NO/POST_SCALE_BIAS | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK | Yes | <= 512, % 64 == 0 | 64 | - | 1 | FP16/BF16 | BS3HD,SB3HD,BSH3D,SBH3D | NO/POST_SCALE_BIAS/ALIBI | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK | Yes | > 512, % 64 == 0 | <= 128, % 8 == 0 | + | 1 | FP16/BF16 | BS3HD,SB3HD,BSH3D,SBH3D | NO/POST_SCALE_BIAS/ALIBI | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK | Yes | any, % 64 == 0 | <= 128, % 8 == 0 | | | | BSHD_BS2HD,BSHD_BSH2D,SBHD_SB2HD,SBHD_SBH2D | | | | | | | | | BSHD_BSHD_BSHD,SBHD_SBHD_SBHD | | | | | | | 2 | FP8 | T3HD | NO_BIAS | PADDING_MASK | Yes | <= 512, % 64 == 0 | 64 | diff --git a/transformer_engine/common/util/pybind_helper.h b/transformer_engine/common/util/pybind_helper.h index fdfa47da8..ef7687e3e 100644 --- a/transformer_engine/common/util/pybind_helper.h +++ b/transformer_engine/common/util/pybind_helper.h @@ -79,7 +79,6 @@ .value("NVTE_Paged_KV_THD_SBHD_SBHD", NVTE_QKV_Layout::NVTE_Paged_KV_THD_SBHD_SBHD) \ .value("NVTE_BHSD_BHSD_BHSD", NVTE_QKV_Layout::NVTE_BHSD_BHSD_BHSD); \ pybind11::enum_(m, "NVTE_Fused_Attn_Backend", pybind11::module_local()) \ - .value("NVTE_F16_max512_seqlen", NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) \ .value("NVTE_F16_arbitrary_seqlen", NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) \ .value("NVTE_FP8", NVTE_Fused_Attn_Backend::NVTE_FP8) \ .value("NVTE_No_Backend", NVTE_Fused_Attn_Backend::NVTE_No_Backend); \ diff --git a/transformer_engine/jax/cpp_extensions/attention.py b/transformer_engine/jax/cpp_extensions/attention.py index 40d02f40e..489bfde99 100644 --- a/transformer_engine/jax/cpp_extensions/attention.py +++ b/transformer_engine/jax/cpp_extensions/attention.py @@ -353,10 +353,7 @@ def abstract( config.window_size, ).get_fused_attn_backend() - if backend == NVTE_Fused_Attn_Backend.NVTE_F16_max512_seqlen: - softmax_shape = (*batch_shape, attn_heads, q_max_seqlen, kv_max_seqlen) - softmax_dtype = q_dtype - elif backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: + if backend == NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen: # cuDNN 9.6 reduces the required softmax shape if get_cudnn_version() >= (9, 6, 0): if config.qkv_layout.is_thd(): diff --git a/transformer_engine/jax/csrc/extensions/attention.cpp b/transformer_engine/jax/csrc/extensions/attention.cpp index 76f2d9289..ed136d7b9 100644 --- a/transformer_engine/jax/csrc/extensions/attention.cpp +++ b/transformer_engine/jax/csrc/extensions/attention.cpp @@ -28,7 +28,6 @@ NVTE_Fused_Attn_Backend GetFusedAttnBackend( /* NOTE: PrepareFusedAttnForwardAuxTensors unifies the auxiliary tensor pack logic from the fused attention forward kernels in: - - common/fused_attn/fused_attn_f16_max512_seqlen.cu lines 594-634 and 773-812 - common/fused_attn/fused_attn_f16_arbitrary_seqlen.cu lines 1270-1281 and 1348-1359 */ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t input_batch, @@ -40,7 +39,6 @@ void PrepareFusedAttnForwardAuxTensors(NVTETensorPack *tensor_pack, const size_t void *bias_buf = nullptr, void *softmax_offset_buf = nullptr) { // all backends need softmax but expect different shapes/dtypes - // start with the max512 sequence length softmax shape/dtype and correct later tensor_pack->size = 1; NVTETensor &softmax_aux = tensor_pack->tensors[0]; NVTEBasicTensor softmax_aux_data; @@ -127,15 +125,6 @@ void PrepareFusedAttnBackwardAuxTensors(NVTETensorPack *tensor_pack, const size_ q_max_seqlen, kv_max_seqlen, dtype, dummy_bias_type, dummy_backend, softmax_buf, rng_state_buf, bias_buf, softmax_offset_buf); - - // correct softmax shape for max512 sequence length kernel - if (backend == NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) { - NVTEBasicTensor softmax_aux_data = - nvte_get_tensor_param(tensor_pack->tensors[0], kNVTERowwiseData); - softmax_aux_data.shape.data[3] = kv_max_seqlen; // {B,H,Qs,1} -> {B,H,Qs,Ks} - softmax_aux_data.dtype = static_cast(dtype); - nvte_set_tensor_param(&(tensor_pack->tensors[0]), kNVTERowwiseData, &softmax_aux_data); - } } pybind11::tuple GetFusedAttnForwardWorkspaceSizes( diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp index b00264394..70d0403b3 100644 --- a/transformer_engine/jax/csrc/extensions/pybind.cpp +++ b/transformer_engine/jax/csrc/extensions/pybind.cpp @@ -189,7 +189,6 @@ PYBIND11_MODULE(transformer_engine_jax, m) { pybind11::enum_(m, "NVTE_Fused_Attn_Backend", pybind11::module_local()) .value("NVTE_No_Backend", NVTE_Fused_Attn_Backend::NVTE_No_Backend) - .value("NVTE_F16_max512_seqlen", NVTE_Fused_Attn_Backend::NVTE_F16_max512_seqlen) .value("NVTE_F16_arbitrary_seqlen", NVTE_Fused_Attn_Backend::NVTE_F16_arbitrary_seqlen) .value("NVTE_FP8", NVTE_Fused_Attn_Backend::NVTE_FP8); diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 4104820a1..79ebbd4af 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -1859,31 +1859,12 @@ def backward(ctx, d_out, *_args): class FusedAttention(torch.nn.Module): - """Dot product attention, with multiple backends: - - 1. FusedAttnBackend["F16_max512_seqlen"] - cuDNN based fused attention for FP16/BF16 and <=512 sequence length. - 2. FusedAttnBackend["F16_arbitrary_seqlen"] - cuDNN based fused attention for FP16/BF16 and any sequence length. - - Support matrix: - - | backend | 1 | 2 | - | flash based | no | yes | - | cuDNN based | yes | yes | - | qkv dtype | fp16/bf16 | fp16/bf16 | - | attn_type | self/cross | self/cross | - | qkv_layout | | | - | - (q,k,v) | sb3hd, bs3hd | sb3hd, bs3hd, sbh3d, bsh3d | - | | sbhd_sb2hd, bshd_bs2hd | sbhd_sb2hd, bshd_bs2hd | - | | bshd_bshd_bshd | sbhd_sbh2d, bshd_bsh2d | - | | | sbhd_sbhd_sbhd, bshd_bshd_bshd | - | mask_type | causal/padding/no_mask | causal/padding/no_mask | - | bias_type | post_scale_bias/no_bias | post_scale_bias/alibi/no_bias | - | dropout | yes | yes | - | max_seqlen | <=512, multiple of 64 | any, multiple of 64 | - | head_dim | 64 | <=128, multiple of 8 | - | output dtype | fp16/bf16 | fp16/bf16 | + """Dot product attention using cuDNN attention: + + FusedAttnBackend["F16_arbitrary_seqlen"] + cuDNN attention for FP16/BF16 with any sequence length. + FusedAttnBackend["FP8"] + cuDNN attention for FP8 with any sequence length. """ def __init__( diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index ed8742353..7df5daabe 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -1217,10 +1217,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt "Disabling FusedAttention as dbias calculation is not supported for 111s" ) use_fused_attention = False - elif not fu_core_attention_bias_requires_grad: - # max512 backend will only support [1, h, s, s] - os.environ["NVTE_FUSED_ATTN_BACKEND"] = "1" - # Filter: cuDNN support fused_attention_backend = None if use_fused_attention: @@ -1254,32 +1250,6 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt logger.debug("Disabling FusedAttention as no backend supports the provided input") use_fused_attention = False fused_attention_backend = None - if ( - use_fused_attention - and window_size is not None - and (window_size[0] != -1 or window_size[1] not in [-1, 0]) - and fused_attention_backend == FusedAttnBackend["F16_max512_seqlen"] - ): - logger.debug( - "Disabling FusedAttention as only sub-backend %s does not support " - "slidng window attention", - int(fused_attention_backend), - ) - use_fused_attention = False - fused_attention_backend = None - if ( - use_fused_attention - and fused_attention_backend == FusedAttnBackend["F16_max512_seqlen"] - and fu_core_attention_bias_type == "post_scale_bias" - and fu_core_attention_bias_shape != "1hss" - ): - logger.debug( - "Disabling FusedAttention as cuDNN sub-backend 0 only supports post_scale_bias in" - " [1, H, S, S] shape" - ) - use_fused_attention = False - fused_attention_backend = None - # Filter: Determinism # backend | deterministic # --------------------------------------------- diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index 01e139da4..d8f301144 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -98,13 +98,12 @@ } FusedAttnBackend = { - "F16_max512_seqlen": NVTE_Fused_Attn_Backend.NVTE_F16_max512_seqlen, "F16_arbitrary_seqlen": NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen, "FP8": NVTE_Fused_Attn_Backend.NVTE_FP8, "No_Backend": NVTE_Fused_Attn_Backend.NVTE_No_Backend, } -BACKEND_F16m512_FP8_THREADS_PER_CTA = 128 +BACKEND_FP8_THREADS_PER_CTA = 128 BACKEND_F16arb_ELTS_PER_THREADS = 16 META_QKV = FP8FwdTensorIdx.GEMM1_OUTPUT @@ -249,22 +248,18 @@ def fused_attn_fwd( if is_training is False, aux_ctx_tensors = None softmax-related tensors: - 1. if fused_attention_backend == FusedAttnBackend["F16_max512_seqlen"] - softmax: torch.Tensor - Softmax(Q*K.T) - shape [batch_size, num_heads, max_seqlen_q, max_seqlen_kv], dtype float32 - 2. if fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"] + 1. if fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"] softmaxStats: torch.Tensor log(sum(e^(x - max(x)))), where x=Q*K.T shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32 - 3. if fused_attention_backend == FusedAttnBackend["FP8"] + 2. if fused_attention_backend == FusedAttnBackend["FP8"] M: torch.Tensor max(Q*K.T) shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32 ZInv: torch.Tensor, only allocated for T3HD path 1/sum(e^(x - max(x))), where x=Q*K.T shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32 - rng_state: torch.Tensor, optional, if backend is not F16_max512_seqlen + rng_state: torch.Tensor state of the random number generator; [seed, offset], dtype uint64 max_logit : if return_max_logit = True, shape [h] and same data type as O; otherwise None @@ -299,19 +294,13 @@ def fused_attn_fwd( f" q.dtype={q.dtype}, backend={fused_attention_backend}." ) - # BF16/FP16 fused attention API from fmha_v1 apex - if fused_attention_backend == FusedAttnBackend["F16_max512_seqlen"]: - rng_elts_per_thread = ( - max_seqlen_q * max_seqlen_kv + BACKEND_F16m512_FP8_THREADS_PER_CTA - 1 - ) // BACKEND_F16m512_FP8_THREADS_PER_CTA - # BF16/FP16 fused attention API from fmha_v2 - elif fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]: + if fused_attention_backend == FusedAttnBackend["F16_arbitrary_seqlen"]: rng_elts_per_thread = BACKEND_F16arb_ELTS_PER_THREADS # FP8 fused attention API from fmha_v2 elif fused_attention_backend == FusedAttnBackend["FP8"]: rng_elts_per_thread = ( - max_seqlen_q * max_seqlen_q + BACKEND_F16m512_FP8_THREADS_PER_CTA - 1 - ) // BACKEND_F16m512_FP8_THREADS_PER_CTA + max_seqlen_q * max_seqlen_q + BACKEND_FP8_THREADS_PER_CTA - 1 + ) // BACKEND_FP8_THREADS_PER_CTA else: raise ValueError(f"Unsupported backend {fused_attention_backend}") @@ -566,13 +555,12 @@ def fused_attn_bwd( f" q.dtype={q.dtype}, backend={fused_attention_backend}." ) - if fused_attention_backend != FusedAttnBackend["F16_max512_seqlen"]: - if len(aux_ctx_tensors) < 1: - raise ValueError( - "aux_ctx_tensors must contain rng_state as its last element," - f" but got len(aux_ctx_tensors)={len(aux_ctx_tensors)}" - f" for backend={fused_attention_backend}." - ) + if len(aux_ctx_tensors) < 1: + raise ValueError( + "aux_ctx_tensors must contain rng_state as its last element," + f" but got len(aux_ctx_tensors)={len(aux_ctx_tensors)}" + f" for backend={fused_attention_backend}." + ) output_tensors = tex.fused_attn_bwd( max_seqlen_q, diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index e6781bd58..8d7a24dce 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -271,7 +271,6 @@ std::vector fused_attn_fwd( nvte_set_tensor_param(&nvte_aux_tensor_pack.tensors[i], kNVTERowwiseData, &temp_data); }; // allocate memory for nvte_aux_tensor_pack.tensors - // f16_max512 : S [b, h, sq, skv] // f16_arbitrary: // return_max_logit=false: S [b, h, sq, 1], rng_state [2], (optional) Bias [1, h, sq, skv], (optional) SoftmaxOffset [1, h, 1, 1] // return_max_logit=true: S [b, h, sq, 1], Max [b, h, sq, 1], rng_state [2], (optional) Bias [1, h, sq, skv], (optional) SoftmaxOffset [1, h, 1, 1] From e8c0dc67b4001a28b1304c22346dd842c4acb283 Mon Sep 17 00:00:00 2001 From: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Date: Thu, 7 May 2026 16:58:53 -0700 Subject: [PATCH 040/180] [PyTorch/Common] Remove legacy FP8DS implementation (#2959) * remove FP8 v0 legacy code Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove FP8 v0 legacy code Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * replace _impl_v1 with _impl Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor tweaks for docstrings Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * minor tweaks for docstring Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * remove T3HD in selection logic Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review: drop dead 8.9 FP8 guard and stale FP8 docs Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> --------- Signed-off-by: Charlene Yang <8636796+cyanguwa@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/attention/test_attention.py | 36 +- .../common/fused_attn/fused_attn.cpp | 75 +- .../common/fused_attn/fused_attn_fp8.cu | 1731 +---------------- .../common/fused_attn/fused_attn_fp8.h | 12 +- transformer_engine/common/fused_attn/utils.cu | 14 - transformer_engine/common/fused_attn/utils.h | 4 - .../include/transformer_engine/fused_attn.h | 22 +- .../dot_product_attention/context_parallel.py | 49 +- .../pytorch/cpp_extensions/fused_attn.py | 11 +- .../pytorch/csrc/extensions/attention.cpp | 14 +- 10 files changed, 102 insertions(+), 1866 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 894137c84..32ea1694e 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -2550,28 +2550,21 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: "fp8_8": ModelConfig(2, 2048, 24, 128, attn_mask_type="causal"), } param_types_fp8 = [torch.float16, torch.bfloat16] -cudnn_frontend_version = int(os.getenv("NVTE_FUSED_ATTN_FE_VER", "1")) -models_v0 = ["fp8_1", "fp8_2", "fp8_5", "fp8_6"] -models_v1 = ["fp8_3", "fp8_4", "fp8_7", "fp8_8"] @pytest.mark.skipif( - ( - get_cudnn_version() < (8, 9, 3) - if cudnn_frontend_version == 0 - else get_cudnn_version() < (9, 2, 1) - ), - reason=f"""cuDNN {"8.9.3" if cudnn_frontend_version == 0 else "9.2.1"}+ is required.""", + get_cudnn_version() < (9, 2, 1), + reason="cuDNN 9.2.1+ is required for FP8 fused attention.", ) @pytest.mark.skipif(not fp8_attn_available, reason=reason_for_no_fp8_attn) @pytest.mark.parametrize("dtype", param_types_fp8) -@pytest.mark.parametrize("model", models_v1 if cudnn_frontend_version == 1 else models_v0) +@pytest.mark.parametrize("model", model_configs_fp8) def test_custom_mha_fp8_vs_f16(dtype, model): """Test FP8 dot product attention implementations based on cuDNN frontend v0.9 and v1.0+. Each test compares results from a custom implementation of an FP8 MHA module, i.e. Custom_MHA_FP8(), to results from an F16 MHA implementation, i.e. transformer_engine.pytorch.attention.MultiHeadAttention. - Both paths take F16 input and output. QKV layout is t3hd or bs3hd""" + Both paths take F16 input and output. QKV layout is bs3hd""" config = model_configs_fp8[model] @@ -2580,7 +2573,7 @@ def test_custom_mha_fp8_vs_f16(dtype, model): available_backends, _, fused_attn_backends = get_available_attention_backends( config, qkv_dtype=torch.float8_e4m3fn, - qkv_layout="t3hd" if cudnn_frontend_version == 0 else "bs3hd", + qkv_layout="bs3hd", is_training=is_training, deterministic=_deterministic, ) @@ -2787,18 +2780,17 @@ def forward( quantization_params=qkv_quantizer, use_split_accumulator=_2X_ACC_FPROP, ) - qkv_layout = "bs3hd" if cudnn_frontend_version == 1 else "t3hd" - o_format = "bshd" if cudnn_frontend_version == 1 else "thd" + qkv_layout = "bs3hd" + o_format = "bshd" qkv = qkv.view(-1, 3, h, d) qkv_fp16 = qkv.dequantize().view(b, max_s, 3, h, d).contiguous() torch.save(qkv_fp16, "qkv.pt") - if cudnn_frontend_version == 1: - qkv = qkv.view(b, max_s, 3, h, d) # bs3hd + qkv = qkv.view(b, max_s, 3, h, d) # bs3hd # FMHA - q_data = qkv._data[:, :, 0, :, :] if cudnn_frontend_version == 1 else qkv._data[:, 0, :, :] - k_data = qkv._data[:, :, 1, :, :] if cudnn_frontend_version == 1 else qkv._data[:, 1, :, :] - v_data = qkv._data[:, :, 2, :, :] if cudnn_frontend_version == 1 else qkv._data[:, 2, :, :] + q_data = qkv._data[:, :, 0, :, :] + k_data = qkv._data[:, :, 1, :, :] + v_data = qkv._data[:, :, 2, :, :] q = qkv.make_like(tensor=qkv, data=q_data, shape=q_data.shape) k = qkv.make_like(tensor=qkv, data=k_data, shape=k_data.shape) v = qkv.make_like(tensor=qkv, data=v_data, shape=v_data.shape) @@ -2820,7 +2812,7 @@ def forward( qkv_layout=qkv_layout, o_format=o_format, attn_bias_type="no_bias", - attn_mask_type=mask_type if cudnn_frontend_version == 1 else "padding", + attn_mask_type=mask_type, rng_gen=None, o_quantizer=o_quantizer, s_quantizer=s_quantizer, @@ -2887,9 +2879,9 @@ def backward(ctx, grad_output: torch.Tensor) -> Tuple[Union[torch.Tensor, None], do_format=ctx.o_format, dqkv_layout=ctx.qkv_layout, attn_bias_type="no_bias", - attn_mask_type=ctx.mask_type if cudnn_frontend_version == 1 else "padding", + attn_mask_type=ctx.mask_type, ) - dim = 2 if cudnn_frontend_version == 1 else 1 + dim = 2 dqkv = torch.Tensor().to(device=dq._data.device, dtype=dq._data.dtype) dqkv_shape = list(dq._data.shape) dqkv_shape.insert(dim, 3) diff --git a/transformer_engine/common/fused_attn/fused_attn.cpp b/transformer_engine/common/fused_attn/fused_attn.cpp index f18a006fc..d2eb1a831 100644 --- a/transformer_engine/common/fused_attn/fused_attn.cpp +++ b/transformer_engine/common/fused_attn/fused_attn.cpp @@ -254,35 +254,31 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( if ((q_dtype == NVTEDType::kNVTEFloat8E4M3 || q_dtype == NVTEDType::kNVTEFloat8E5M2) && sm_arch_ >= 90 && bias_type == NVTE_Bias_Type::NVTE_NO_BIAS && - // 8.9: t3hd, max_s=512, d=64, padding - ((cudnn_runtime_version >= 8900 && sm_arch_ < 100 && - qkv_layout == NVTE_QKV_Layout::NVTE_T3HD && max_seqlen_q == max_seqlen_kv && - max_seqlen_q <= 512 && head_dim_qk == 64 && head_dim_v == 64 && - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK) || - // 9.2.1: {bshd, sbhd}, any seqlen, d=128, {no_mask, causal} - (cudnn_runtime_version >= 90201 && sm_arch_ < 100 && max_seqlen_q % 128 == 0 && - max_seqlen_kv % 128 == 0 && head_dim_qk == 128 && head_dim_v == 128 && - (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK)) || - // 9.7: {bshd, sbhd}, any seqlen, d<=256 for sm90 and d<=128 for sm100, {padding, padding_causal} - (cudnn_runtime_version >= 90700 && - // TODO (cyang): add is_training to nvte_get_fused_attn_backend - // sm90: fwd d<=256, bwd d=128 only - // sm100: fwd d<=128, bwd d<=128 - ((sm_arch_ < 100 && (!is_training) && head_dim_qk <= 256 && head_dim_v <= 256) || - (sm_arch_ < 100 && is_training && head_dim_qk == 128 && head_dim_v == 128) || - (sm_arch_ >= 100 && head_dim_qk <= 128 && head_dim_v <= 128)) && - head_dim_qk % 16 == 0 && head_dim_v % 16 == 0 && - (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)) || - // 9.21: d_qk=192, d_v=128 - (cudnn_runtime_version >= 92100 && sm_arch_ >= 100 && head_dim_qk <= 192 && - head_dim_v <= 128 && head_dim_qk % 16 == 0 && head_dim_v % 16 == 0 && - (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || - attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK))) && + ( + // 9.2.1: {bshd, sbhd}, any seqlen, d=128, {no_mask, causal} + (cudnn_runtime_version >= 90201 && sm_arch_ < 100 && max_seqlen_q % 128 == 0 && + max_seqlen_kv % 128 == 0 && head_dim_qk == 128 && head_dim_v == 128 && + (attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK)) || + // 9.7: {bshd, sbhd}, any seqlen, d<=256 for sm90 and d<=128 for sm100, {padding, padding_causal} + (cudnn_runtime_version >= 90700 && + // TODO (cyang): add is_training to nvte_get_fused_attn_backend + // sm90: fwd d<=256, bwd d=128 only + // sm100: fwd d<=128, bwd d<=128 + ((sm_arch_ < 100 && (!is_training) && head_dim_qk <= 256 && head_dim_v <= 256) || + (sm_arch_ < 100 && is_training && head_dim_qk == 128 && head_dim_v == 128) || + (sm_arch_ >= 100 && head_dim_qk <= 128 && head_dim_v <= 128)) && + head_dim_qk % 16 == 0 && head_dim_v % 16 == 0 && + (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_PADDING_CAUSAL_MASK)) || + // 9.21: d_qk=192, d_v=128 + (cudnn_runtime_version >= 92100 && sm_arch_ >= 100 && head_dim_qk <= 192 && + head_dim_v <= 128 && head_dim_qk % 16 == 0 && head_dim_v % 16 == 0 && + (attn_mask_type == NVTE_Mask_Type::NVTE_NO_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_MASK || + attn_mask_type == NVTE_Mask_Type::NVTE_CAUSAL_BOTTOM_RIGHT_MASK))) && // pre-9.21: {bshd, sbhd}, {vanilla} // 9.21+: {bshd, sbhd, bhsd}, {vanilla, off-by-one, learnable} ((cudnn_runtime_version < 92100 && @@ -294,14 +290,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( !requires_64bit_ragged_offset && // 9.10.0: known bugs with SDPA FP8 (cudnn_runtime_version != 91000) && !return_max_logit) { - if (cudnn_runtime_version >= 8900) { - backend = NVTE_Fused_Attn_Backend::NVTE_FP8; - } else { - backend = NVTE_Fused_Attn_Backend::NVTE_No_Backend; - std::cout << "Warning: FP8 fused attention is supported by cuDNN 8.9.0+." - " Please upgrade your cuDNN version if possible." - << std::endl; - } + backend = NVTE_Fused_Attn_Backend::NVTE_FP8; } else if ((q_dtype == NVTEDType::kNVTEFloat16) || (q_dtype == NVTEDType::kNVTEBFloat16)) { bool flag_arb = false; if ( @@ -727,10 +716,6 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso } else if (fused_attention_backend == NVTE_Fused_Attn_Backend::NVTE_FP8) { size_t i = 0; const Tensor *input_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - const Tensor *input_ZInv = nullptr; - if (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD) { - input_ZInv = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - } const Tensor *input_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); const Tensor *input_SoftmaxOffset = nullptr; if (softmax_type != NVTE_VANILLA_SOFTMAX) { @@ -744,10 +729,10 @@ void nvte_fused_attn_bwd(const NVTETensor Q, const NVTETensor K, const NVTETenso qkv_layout, o_format, do_format, dqkv_layout, qkv_scale_inv_format, do_scale_inv_format, bias_type, attn_mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, deterministic, - input_Q, input_K, input_V, input_O, input_dO, input_dO_f16, input_M, - input_ZInv, input_S, input_SoftmaxOffset, input_output_dP, output_dQ, - output_dK, output_dV, output_dSoftmaxOffset, input_cu_seqlens_q, - input_cu_seqlens_kv, input_rng_state, wkspace, stream, handle); + input_Q, input_K, input_V, input_O, input_dO, input_dO_f16, input_M, input_S, + input_SoftmaxOffset, input_output_dP, output_dQ, output_dK, output_dV, + output_dSoftmaxOffset, input_cu_seqlens_q, input_cu_seqlens_kv, + input_rng_state, wkspace, stream, handle); } else { NVTE_ERROR("Invalid combination of data type and sequence length for fused attention. \n"); } diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.cu b/transformer_engine/common/fused_attn/fused_attn_fp8.cu index d97f38845..eab1ae02e 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.cu +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.cu @@ -15,1648 +15,14 @@ namespace fused_attn { using namespace transformer_engine; -std::unordered_map tensor_name_to_uid = {{"Q", 1}, - {"K", 2}, - {"V", 3}, - {"O", 4}, - {"S", 5}, - {"B", 6}, - {"DROPOUT_SCALE", 7}, - {"S_CONST", 8}, - {"MNK_OVERRIDE", 9}, - {"dQ", 11}, - {"dK", 12}, - {"dV", 13}, - {"dO", 14}, - {"MASK_VAL", 15}, - {"dS", 16}, - {"O_SEQLEN", 17}, - {"M", 18}, - {"Z", 19}, - {"descaleQ", 20}, - {"descaleK", 21}, - {"descaleV", 22}, - {"descaleS", 23}, - {"scaleS", 24}, - {"amaxS", 25}, - {"amaxO", 26}, - {"QKV_RAGGED", 27}, - {"O_RAGGED", 28}, - {"K_TRANSPOSE", 29}, - {"AttnScale", 30}, - {"scaleO", 31}, - {"Z_INV", 32}, - {"descaleO", 33}, - {"descaledO", 34}, - {"descaledS", 35}, - {"descaledQ", 36}, - {"descaledK", 37}, - {"descaledV", 38}, - {"scaledS", 39}, - {"scaledQ", 40}, - {"scaledK", 41}, - {"scaledV", 42}, - {"amaxdS", 43}, - {"amaxdQ", 44}, - {"amaxdK", 45}, - {"amaxdV", 46}, - {"V_TRANSPOSE", 47}, - {"AttnScale_dS_K", 48}, - {"AttnScale_dSTranspose_Q", 49}, - {"DROPOUT_SCALE_dOVt_OdO", 50}, - {"DROPOUT_OFFSET", 51}, - {"DROPOUT_SEED", 52}, - {"VIRTUAL", 80}}; - -static cudnn_frontend::Tensor createAmax(const std::string& amax_tensor_name, - const cudnn_frontend::Tensor& prevBlockOutputTensor, - std::vector* ops) { - int64_t amax_dim[4] = {1, 1, 1, 1}; - int64_t amax_stride[4] = {1, 1, 1, 1}; - auto amaxTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid[amax_tensor_name], amax_dim, - amax_stride, false, false); - - // Define the amax descriptor - auto reductionDesc = cudnn_frontend::ReductionDescBuilder() - .setMathPrecision(CUDNN_DATA_FLOAT) - .setReductionOp(CUDNN_REDUCE_TENSOR_AMAX) - .build(); - - // Create a reduction amax Node - auto reduction_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_REDUCTION_DESCRIPTOR) - .setxDesc(prevBlockOutputTensor) - .setyDesc(amaxTensor) - .setreductionDesc(reductionDesc) - .build(); - ops->push_back(std::move(reduction_op)); - return amaxTensor; -} - -static cudnn_frontend::Tensor createScale(const cudnn_frontend::Tensor& prevBlockOutputTensor, - const std::string& scale_tensor_name, - cudnnDataType_t tensorType, bool isOutputVirtual, - bool isScaleByValue, - std::vector* ops, - const std::string& output_tensor_name = "") { - int64_t scale_dim[4] = {1, 1, 1, 1}; - int64_t scale_stride[4] = {1, 1, 1, 1}; - - int64_t output_dim[4]; - int64_t output_stride[4]; - - for (int i = 0; i < 4; i++) { - output_dim[i] = prevBlockOutputTensor.getDim()[i]; - output_stride[i] = prevBlockOutputTensor.getStride()[i]; - } - - auto scaleTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid[scale_tensor_name], - scale_dim, scale_stride, false, isScaleByValue); // is by value - - int64_t outputUID = - isOutputVirtual ? tensor_name_to_uid["VIRTUAL"] + tensor_name_to_uid[scale_tensor_name] + 5000 - : tensor_name_to_uid[output_tensor_name]; - auto afterScaleKTensor = tensor_create(tensorType, outputUID, output_dim, output_stride, - isOutputVirtual, false); // is virtual - - // Define the scale descriptor - auto scaleDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // Create a Scale Node - auto scale_op = - binary_pw_op_create(prevBlockOutputTensor, scaleTensor, afterScaleKTensor, scaleDesc); - - ops->push_back(std::move(scale_op)); - return afterScaleKTensor; -} - -static cudnn_frontend::Tensor createScale(const cudnn_frontend::Tensor& prevBlockOutputTensor, - const cudnn_frontend::Tensor& scaleTensor, - cudnnDataType_t tensorType, bool isOutputVirtual, - bool isScaleByValue, - std::vector* ops, - int UID_offset, - const std::string& output_tensor_name = "") { - int64_t output_dim[4]; - int64_t output_stride[4]; - for (int i = 0; i < 4; i++) { - output_dim[i] = prevBlockOutputTensor.getDim()[i]; - output_stride[i] = prevBlockOutputTensor.getStride()[i]; - } - - int64_t outputUID = isOutputVirtual ? tensor_name_to_uid["VIRTUAL"] + UID_offset - : tensor_name_to_uid[output_tensor_name]; - auto afterScaleTensor = tensor_create(tensorType, outputUID, output_dim, output_stride, - isOutputVirtual, false); // is virtual - - // Define the scale descriptor - auto scaleDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // Create a Scale Node - auto scale_op = - binary_pw_op_create(prevBlockOutputTensor, scaleTensor, afterScaleTensor, scaleDesc); - - ops->push_back(std::move(scale_op)); - return afterScaleTensor; -} - -static cudnn_frontend::Tensor createScaleWithOffset( - const cudnn_frontend::Tensor& prevBlockOutputTensor, const std::string& scale_tensor_name, - NVTE_QKV_Layout layout, cudnnDataType_t tensorType, bool isOutputVirtual, bool isScaleByValue, - std::vector* ops, - std::shared_ptr offsetTensor, - const std::string& output_tensor_name = "") { - int64_t scale_dim[4] = {1, 1, 1, 1}; - int64_t scale_stride[4] = {1, 1, 1, 1}; - - int64_t output_dim[4]; - int64_t output_stride[4]; - // If output tensor is dQ, dK, or dV, we need to generate QKV interleaved strides - if (output_tensor_name == "dQ" || output_tensor_name == "dK" || output_tensor_name == "dV") { - for (int i = 0; i < 4; i++) { - output_dim[i] = prevBlockOutputTensor.getDim()[i]; - } - generateMatrixStrides(output_dim[0], output_dim[1], output_dim[2], - 0 /*s_kv = 0 for placeholder*/, output_dim[3], output_stride, layout, - NVTE_QKV_Matrix::NVTE_Q_Matrix); - } else { - // Otherwise output dim and stride should be the same as prev block dim and stride - for (int i = 0; i < 4; i++) { - output_dim[i] = prevBlockOutputTensor.getDim()[i]; - output_stride[i] = prevBlockOutputTensor.getStride()[i]; - } - } - - auto scaleTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid[scale_tensor_name], - scale_dim, scale_stride, false, isScaleByValue); // is by value - - cudnnDataType_t outputDataType = isOutputVirtual ? CUDNN_DATA_FLOAT : tensorType; - int64_t outputUID = - isOutputVirtual ? tensor_name_to_uid["VIRTUAL"] + tensor_name_to_uid[scale_tensor_name] + 7000 - : tensor_name_to_uid[output_tensor_name]; - auto afterScaleTensor = - tensor_create_with_offset(outputDataType, outputUID, output_dim, output_stride, - isOutputVirtual, false, offsetTensor); // is virtual - - // Define the scale descriptor - auto scaleDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // Create a Scale Node - auto scale_op = - binary_pw_op_create(prevBlockOutputTensor, scaleTensor, afterScaleTensor, scaleDesc); - - ops->push_back(std::move(scale_op)); - return afterScaleTensor; -} - -static cudnn_frontend::Tensor createSoftmaxForward( - int64_t b, int64_t h, int64_t s_q, int64_t s_kv, std::vector* ops, - const cudnn_frontend::Tensor& prevBlockOutputTensor, bool isTraining) { - int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv}; - int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1}; - - int64_t afterReduction_dim[4] = {b, h, s_q, 1}; - int64_t afterReduction_stride[4] = {h * s_q, s_q, 1, 1}; - - // max (x) (M tensor) - auto afterMaxReductionTensor = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["M"], afterReduction_dim, - afterReduction_stride, !isTraining, false); // not virtual if training is true, - // virtual if training is false - // x - max(x) - auto afterSubtractionTensor = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 151, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - // e^(x - max(x)) - auto afterExponentTensor = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 152, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual; - // sum (e^(x - max(x))) (Z tensor) - auto zTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["Z"], afterReduction_dim, - afterReduction_stride, true, false); // is virtual - // 1 / sum (e^(x - max(x))) (Z_INV tensor) - auto zInvTensor = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["Z_INV"], afterReduction_dim, - afterReduction_stride, !isTraining, false); // not virtual if training is true, - // virtual if training is false - // Final softmax output (After exponent * Z_INV) - auto beforeDropoutTensor = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 153, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - - // Define the reduction descriptor - auto reductionMaxDesc = cudnn_frontend::ReductionDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setReductionOp(CUDNN_REDUCE_TENSOR_MAX) - .build(); - - // Create a reduction max Node - auto reductionMax_op = - cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_REDUCTION_DESCRIPTOR) - .setxDesc(prevBlockOutputTensor) - .setyDesc(afterMaxReductionTensor) - .setreductionDesc(reductionMaxDesc) - .build(); - - // Define the subtract descriptor - auto subtractDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_SUB); - - // Create a subtract Node - auto subtract_op = binary_pw_op_create(prevBlockOutputTensor, afterMaxReductionTensor, - afterSubtractionTensor, subtractDesc); - - // Define the exponent descriptor - auto exponentDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_EXP); - - // Create a exponent Node - auto exponent_op = unary_pw_op_create(afterSubtractionTensor, afterExponentTensor, exponentDesc); - - // Define the reduction descriptor - auto reductionAddDesc = cudnn_frontend::ReductionDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setReductionOp(CUDNN_REDUCE_TENSOR_ADD) - .build(); - - // Create a reduction add Node - auto reductionAdd_op = - cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_REDUCTION_DESCRIPTOR) - .setxDesc(afterExponentTensor) - .setyDesc(zTensor) - .setreductionDesc(reductionAddDesc) - .build(); - - // Define the reciprocal descriptor - auto reciprocalDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_RECIPROCAL); - - // Create a reciprocal Node - auto reciprocal_op = unary_pw_op_create(zTensor, zInvTensor, reciprocalDesc); - - // Define the pw multiply descriptor - auto multiplyDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // Create a multiply Node - auto mutliply_op = - binary_pw_op_create(afterExponentTensor, zInvTensor, beforeDropoutTensor, multiplyDesc); - - ops->push_back(std::move(reductionMax_op)); - ops->push_back(std::move(subtract_op)); - ops->push_back(std::move(exponent_op)); - ops->push_back(std::move(reductionAdd_op)); - ops->push_back(std::move(reciprocal_op)); - ops->push_back(std::move(mutliply_op)); - - return beforeDropoutTensor; -} - -static cudnn_frontend::Tensor createDropoutForward( - int64_t b, int64_t h, int64_t s_q, int64_t s_kv, double probability, - std::vector* ops, - const cudnn_frontend::Tensor& beforeDropoutTensor) { - NVTE_CHECK(ops->size() > 0, "Dropout DAG constructed incorrectly as the first one"); - - int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv}; - int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1}; - - int64_t scale_dim[4] = {1, 1, 1, 1}; - int64_t scale_stride[4] = {1, 1, 1, 1}; - - // Mask for the dropout - auto dropoutMaskTensor = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 250, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - auto dropoutSeedTensor = tensor_create(CUDNN_DATA_INT64, tensor_name_to_uid["DROPOUT_SEED"], - scale_dim, scale_stride, false, false); // is by value - auto dropoutOffsetTensor = tensor_create(CUDNN_DATA_INT64, tensor_name_to_uid["DROPOUT_OFFSET"], - scale_dim, scale_stride, false, false); // is by value - - // After dropout tensor befor scale - auto beforeDropoutScaleTensor = - cudnn_frontend::TensorBuilder() - .setDim(4, afterBMM1_dim) - .setStride(4, afterBMM1_stride) - .setId(tensor_name_to_uid["VIRTUAL"] + 201) - .setAlignment(16) // 16B alignment is needed to run a tensor core engine - .setDataType(CUDNN_DATA_FLOAT) - .setVirtual(true) - .setByValue(false) - .setReorderType(cudnn_frontend::TensorReordering_t::F16x16) - .build(); - // Scale after dropout - auto scaleDropoutTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["DROPOUT_SCALE"], - scale_dim, scale_stride, false, true); // is by value - // After Scale - auto afterDropout_before_quan_S = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 202, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - - // Define the reduction descriptor - auto rngDesc = cudnn_frontend::RngDescBuilder() - .setRngDistribution(CUDNN_RNG_DISTRIBUTION_BERNOULLI) - .setBernoulliDistProbability(1.0 - probability) - .build(); - - // Create a rng Node - auto rng_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_RNG_DESCRIPTOR) - .setyDesc(dropoutMaskTensor) - .setSeedDesc(dropoutSeedTensor) - .setOffsetDesc(dropoutOffsetTensor) - .setRngDesc(rngDesc) - .build(); - - // Define the multiply mask descriptor - auto maskMulDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // Create a multiply mask Node - auto maskMul_op = binary_pw_op_create(beforeDropoutTensor, dropoutMaskTensor, - beforeDropoutScaleTensor, maskMulDesc); - - // Define the multiply scale descriptor - auto scaleMulDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // Create a multiply mask Node - auto scaleMul_op = binary_pw_op_create(beforeDropoutScaleTensor, scaleDropoutTensor, - afterDropout_before_quan_S, scaleMulDesc); - - ops->push_back(std::move(rng_op)); - ops->push_back(std::move(maskMul_op)); - ops->push_back(std::move(scaleMul_op)); - - return afterDropout_before_quan_S; -} - -static cudnn_frontend::Tensor createDropoutBackward( - int64_t b, int64_t h, int64_t s_q, int64_t s_kv, double probability, - std::vector* ops, const cudnn_frontend::Tensor& beforeDropoutTensor, - const cudnn_frontend::Tensor& dropoutMaskTensor) { - NVTE_CHECK(ops->size() > 0, "Dropout DAG constructed incorrectly as the first one"); - - int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv}; - int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1}; - - int64_t scale_dim[4] = {1, 1, 1, 1}; - int64_t scale_stride[4] = {1, 1, 1, 1}; - - auto dropoutSeedTensor = tensor_create(CUDNN_DATA_INT64, tensor_name_to_uid["DROPOUT_SEED"], - scale_dim, scale_stride, false, false); // is by value - auto dropoutOffsetTensor = tensor_create(CUDNN_DATA_INT64, tensor_name_to_uid["DROPOUT_OFFSET"], - scale_dim, scale_stride, false, false); // is by value - - // After dropout tensor befor scale - auto beforeDropoutScaleTensor = - cudnn_frontend::TensorBuilder() - .setDim(4, afterBMM1_dim) - .setStride(4, afterBMM1_stride) - .setId(tensor_name_to_uid["VIRTUAL"] + 201) - .setAlignment(16) // 16B alignment is needed to run a tensor core engine - .setDataType(CUDNN_DATA_FLOAT) - .setVirtual(true) - .setByValue(false) - .setReorderType(cudnn_frontend::TensorReordering_t::F16x16) - .build(); - // Scale after dropout (1 / (1 - p)) - auto scaleDropoutTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["DROPOUT_SCALE"], - scale_dim, scale_stride, false, true); // is by value - // After Scale - auto afterDropout_before_quan_S = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 202, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - - // Define the reduction descriptor - auto rngDesc = cudnn_frontend::RngDescBuilder() - .setRngDistribution(CUDNN_RNG_DISTRIBUTION_BERNOULLI) - .setBernoulliDistProbability(1.0 - probability) - .build(); - - // Create a rng Node - auto rng_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_RNG_DESCRIPTOR) - .setyDesc(dropoutMaskTensor) - .setSeedDesc(dropoutSeedTensor) - .setOffsetDesc(dropoutOffsetTensor) - .setRngDesc(rngDesc) - .build(); - - // Define the multiply mask descriptor - auto maskMulDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // Create a multiply mask Node - auto maskMul_op = binary_pw_op_create(beforeDropoutTensor, dropoutMaskTensor, - beforeDropoutScaleTensor, maskMulDesc); - - // Define the multiply scale descriptor - auto scaleMulDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // Create a multiply mask Node - auto scaleMul_op = binary_pw_op_create(beforeDropoutScaleTensor, scaleDropoutTensor, - afterDropout_before_quan_S, scaleMulDesc); - - ops->push_back(std::move(rng_op)); - ops->push_back(std::move(maskMul_op)); - ops->push_back(std::move(scaleMul_op)); - - return afterDropout_before_quan_S; -} - -static cudnn_frontend::Tensor createSoftmaxBackward(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, - std::vector* ops, - const cudnn_frontend::Tensor& dyTensor) { - NVTE_CHECK(ops->size() > 0, "Softmax backward constructed incorrectly as the first one"); - - int64_t dx_dim[4] = {b, h, s_q, s_kv}; - int64_t dx_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1}; - - int64_t M_Z_dim[4] = {b, h, s_q, 1}; - int64_t M_Z_stride[4] = {h * s_q, s_q, 1, 1}; - - // Creating all tensors - auto MTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["M"], M_Z_dim, M_Z_stride, - false, false); // not virtual - auto ZInvTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["Z_INV"], M_Z_dim, - M_Z_stride, false, false); // not virtual - auto dxAfterSubtractionTensor = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 252, dx_dim, dx_stride, true, - false); // is virtual - auto dxAfterExponentiation = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 253, - dx_dim, dx_stride, true, false); // is virtual - auto dxBeforeDropout_QKt_Tensor = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 254, dx_dim, dx_stride, true, - false); // is virtual - - // Creating all ops - // sub (dy - M) - auto subtractionDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_SUB); - auto subtractionOp = - binary_pw_op_create(dyTensor, MTensor, dxAfterSubtractionTensor, subtractionDesc); - - // Define the exponent descriptor - auto exponentDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_EXP); - - // Create a exponent Node. (exp(dy - M)) - auto exponentOp = - unary_pw_op_create(dxAfterSubtractionTensor, dxAfterExponentiation, exponentDesc); - - // Define the pw multiply descriptor - auto multiplyDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // Create a multiply Node - auto mutliplyOp = binary_pw_op_create(dxAfterExponentiation, ZInvTensor, - dxBeforeDropout_QKt_Tensor, multiplyDesc); - - ops->push_back(std::move(subtractionOp)); - ops->push_back(std::move(exponentOp)); - ops->push_back(std::move(mutliplyOp)); - - return dxBeforeDropout_QKt_Tensor; -} - -static cudnn_frontend::Tensor createQKBMM( - int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, NVTE_QKV_Layout layout, - cudnnDataType_t tensorType, std::vector* ops, - const cudnn_frontend::Tensor& qTensor, const cudnn_frontend::Tensor& kTensor, - const cudnn_frontend::Tensor& mnkOverride, - std::shared_ptr QKVRaggedOffsetTensor) { - // Creates the necessary tensor descriptors - int64_t k_transpose_dim[4] = {b, h, d, s_kv}; - int64_t k_transpose_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, k_transpose_stride, layout, - NVTE_QKV_Matrix::NVTE_K_Matrix_Transpose); - - int64_t s_dim[4] = {b, h, s_q, s_kv}; - int64_t s_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, s_stride, layout, NVTE_QKV_Matrix::NVTE_S_Matrix); - - auto kTransposeTensor = tensor_create_with_offset(tensorType, tensor_name_to_uid["K_TRANSPOSE"], - k_transpose_dim, k_transpose_stride, false, - false, QKVRaggedOffsetTensor); // is virtual - - // First GEMM output - auto afterQKTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 1, s_dim, - s_stride, true, false); // is virtual - - // Define the matmul desc - auto matmulDesc = cudnn_frontend::MatMulDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setPaddingValue(-2000000) - .build(); - - // Create reshape node for K -> K.T - auto reshape_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_RESHAPE_DESCRIPTOR) - .setxDesc(kTensor) - .setyDesc(kTransposeTensor) - .build(); - - // Create a matmul Node - auto matmulOp = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_MATMUL_DESCRIPTOR) - .setaMatDesc(qTensor) - .setbMatDesc(kTransposeTensor) - .setcMatDesc(afterQKTensor) - .setmOverrideDesc(mnkOverride) - .setnOverrideDesc(mnkOverride) - .setmatmulDesc(matmulDesc) - .build(); - - ops->push_back(std::move(reshape_op)); - ops->push_back(std::move(matmulOp)); - - return afterQKTensor; -} - -static cudnn_frontend::Tensor createSVBMM( - int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, NVTE_QKV_Layout layout, - cudnnDataType_t tensorType, std::vector* ops, - const cudnn_frontend::Tensor& softmaxTensor, const cudnn_frontend::Tensor& mnkOverride, - std::shared_ptr QKVRaggedOffsetTensor) { - NVTE_CHECK(ops->size() > 0, "BMM2 op constructed incorrectly as the first one"); - - int64_t v_dim[4] = {b, h, s_kv, d}; - int64_t v_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, v_stride, layout, NVTE_QKV_Matrix::NVTE_V_Matrix); - - int64_t o_dim[4] = {b, h, s_q, d}; - int64_t o_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, o_stride, layout, NVTE_QKV_Matrix::NVTE_O_Matrix); - - auto vTensor = tensor_create_with_offset(tensorType, tensor_name_to_uid["V"], v_dim, v_stride, - false, false, QKVRaggedOffsetTensor); - // Second fprop GEMM output - auto oTensor = tensor_create(tensorType, tensor_name_to_uid["VIRTUAL"] + 300, o_dim, o_stride, - true, false); // is virtual - - // Define the matmul desc - auto matmulDesc = cudnn_frontend::MatMulDescBuilder().setComputeType(CUDNN_DATA_FLOAT).build(); - - // Create a matmul Node - auto matmulOp = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_MATMUL_DESCRIPTOR) - .setaMatDesc(softmaxTensor) - .setbMatDesc(vTensor) - .setcMatDesc(oTensor) - .setmOverrideDesc(mnkOverride) - .setkOverrideDesc(mnkOverride) - .setmatmulDesc(matmulDesc) - .build(); - - ops->push_back(std::move(matmulOp)); - - return oTensor; -} - -static cudnn_frontend::Tensor createSdOBMM(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, - int64_t d, cudnnDataType_t tensorType, - std::vector* ops, - const cudnn_frontend::Tensor& softmaxTensor, - const cudnn_frontend::Tensor& dOTensor, - const cudnn_frontend::Tensor& mnkOverride) { - NVTE_CHECK(ops->size() > 0, "BMM2 op constructed incorrectly as the first one"); - - int64_t s_dim_transpose[4] = {b, h, s_kv, s_q}; - int64_t s_stride_transpose[4] = {h * s_kv * s_q, s_kv * s_q, 1, s_kv}; - - int64_t v_dim[4] = {b, h, s_kv, d}; - int64_t v_stride[4] = {h * s_kv * d, d, h * d, 1}; - - auto sTransposeTensor = - tensor_create(tensorType, tensor_name_to_uid["VIRTUAL"] + 499, s_dim_transpose, - s_stride_transpose, true, false); // is virtual - // S.T * dO - auto dVTensor_before_dequan_S = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 500, v_dim, v_stride, true, - false); // is virtual - - // Create reshape node for softmax -> softmax.T - auto reshape_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_RESHAPE_DESCRIPTOR) - .setxDesc(softmaxTensor) - .setyDesc(sTransposeTensor) - .build(); - - // Define the matmul desc - auto matmulDesc = cudnn_frontend::MatMulDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setPaddingValue(0) - .build(); - - // Create a matmul Node - auto matmulOp = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_MATMUL_DESCRIPTOR) - .setaMatDesc(sTransposeTensor) - .setbMatDesc(dOTensor) - .setcMatDesc(dVTensor_before_dequan_S) - .setmOverrideDesc(mnkOverride) - .setkOverrideDesc(mnkOverride) - .setmatmulDesc(matmulDesc) - .build(); - - ops->push_back(std::move(reshape_op)); - ops->push_back(std::move(matmulOp)); - - return dVTensor_before_dequan_S; -} - -static cudnn_frontend::Tensor createdOVBMM( - int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, NVTE_QKV_Layout layout, - cudnnDataType_t tensorType, std::vector* ops, - const cudnn_frontend::Tensor& dOTensor, const cudnn_frontend::Tensor& mnkOverride, - std::shared_ptr QKVRaggedOffsetTensor) { - // Creates the necessary tensor descriptors - int64_t v_dim[4] = {b, h, s_kv, d}; - int64_t v_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, v_stride, layout, NVTE_QKV_Matrix::NVTE_V_Matrix); - - int64_t v_transpose_dim[4] = {b, h, d, s_kv}; - int64_t v_transpose_stride[4]; - v_transpose_stride[0] = v_stride[0]; - v_transpose_stride[1] = v_stride[1]; - v_transpose_stride[2] = v_stride[3]; - v_transpose_stride[3] = v_stride[2]; - - int64_t s_dim[4] = {b, h, s_q, s_kv}; - int64_t s_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, s_stride, layout, NVTE_QKV_Matrix::NVTE_S_Matrix); - - auto vTensor = tensor_create_with_offset(tensorType, tensor_name_to_uid["V"], v_dim, v_stride, - false, false, QKVRaggedOffsetTensor); - auto vTransposeTensor = tensor_create_with_offset(tensorType, tensor_name_to_uid["V_TRANSPOSE"], - v_transpose_dim, v_transpose_stride, false, - false, QKVRaggedOffsetTensor); // is virtual - - // dO * V.T - auto afterdOVTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 600, s_dim, - s_stride, true, false); // is virtual - - // Define the matmul desc - auto matmulDesc = cudnn_frontend::MatMulDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setPaddingValue(0) - .build(); - - // Create reshape node for V -> V.T - auto reshape_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_RESHAPE_DESCRIPTOR) - .setxDesc(vTensor) - .setyDesc(vTransposeTensor) - .build(); - - // Create a matmul Node - auto matmulOp = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_MATMUL_DESCRIPTOR) - .setaMatDesc(dOTensor) - .setbMatDesc(vTransposeTensor) - .setcMatDesc(afterdOVTensor) - .setmOverrideDesc(mnkOverride) - .setnOverrideDesc(mnkOverride) - .setmatmulDesc(matmulDesc) - .build(); - - ops->push_back(std::move(reshape_op)); - ops->push_back(std::move(matmulOp)); - - return afterdOVTensor; -} - -static cudnn_frontend::Tensor createdOAndORowReductionChain( - int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, NVTE_QKV_Layout layout, - std::vector* ops, const cudnn_frontend::Tensor& O_after_dequan, - const cudnn_frontend::Tensor& dO_after_dequan, - const cudnn_frontend::Tensor& dropoutScale_dOVt_OdO_Tensor) { - int64_t o_dim[4] = {b, h, s_q, d}; - int64_t o_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, o_stride, layout, NVTE_QKV_Matrix::NVTE_O_Matrix); - int64_t o_dim_row_sum[4] = {b, h, s_q, 1}; - int64_t o_dim_row_sum_stride[4] = {s_q * h, s_q, 1, 1}; - - auto O_dO_after_pointwise_multiply = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 700, o_dim, o_stride, true, - false); // is virtual - auto O_dO_after_dropout_scale = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 701, o_dim, o_stride, true, - false); // is virtual - auto O_dO_after_rowsum = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 702, o_dim_row_sum, - o_dim_row_sum_stride, true, false); // is virtual - - // Define the pw multiply descriptor - auto multiplyDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // Create a multiply Node - auto mutliply_op = binary_pw_op_create(O_after_dequan, dO_after_dequan, - O_dO_after_pointwise_multiply, multiplyDesc); - - // Create multiply node with dropout scale - auto dropout_scale_multiply_op = - binary_pw_op_create(O_dO_after_pointwise_multiply, dropoutScale_dOVt_OdO_Tensor, - O_dO_after_dropout_scale, multiplyDesc); - - // Define the reduction descriptor - auto reductionAddDesc = cudnn_frontend::ReductionDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setReductionOp(CUDNN_REDUCE_TENSOR_ADD) - .build(); - - // Create a reduction add Node - auto reductionAdd_op = - cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_REDUCTION_DESCRIPTOR) - .setxDesc(O_dO_after_dropout_scale) - .setyDesc(O_dO_after_rowsum) - .setreductionDesc(reductionAddDesc) - .build(); - - ops->push_back(std::move(mutliply_op)); - ops->push_back(std::move(dropout_scale_multiply_op)); - ops->push_back(std::move(reductionAdd_op)); - - return O_dO_after_rowsum; -} - -static cudnn_frontend::Tensor createBiasSubtractionSoftmaxMulChain( - int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, NVTE_QKV_Layout layout, - std::vector* ops, const cudnn_frontend::Tensor& dS_after_dropout, - const cudnn_frontend::Tensor& AfterDropout_before_quan_S, - const cudnn_frontend::Tensor& O_dO_after_rowsum, const cudnn_frontend::Tensor& attnScale) { - int64_t o_dim[4] = {b, h, s_q, s_kv}; - int64_t o_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, o_stride, layout, NVTE_QKV_Matrix::NVTE_S_Matrix); - auto dS_minus_O_dO = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 800, o_dim, - o_stride, true, false); // is virtual - auto AfterAttnScale_before_dS = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 801, o_dim, o_stride, true, - false); // is virtual - auto S_mul_dS_minus_O_dO = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 802, - o_dim, o_stride, true, false); // is virtual - - // Define the pw subtraction descriptor - auto subDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_SUB); - - // Create a subtraction Node - auto sub_op = binary_pw_op_create(dS_after_dropout, O_dO_after_rowsum, dS_minus_O_dO, subDesc); - - // Define the pw multiplication descriptor - auto multiplyDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // dS_minus_O_dO * attnScale - auto mutliply_attn_scale_op = - binary_pw_op_create(dS_minus_O_dO, attnScale, AfterAttnScale_before_dS, multiplyDesc); - - // AfterDropout_before_quan_S * AfterAttnScale_before_dS - auto mutliply_op = binary_pw_op_create(AfterDropout_before_quan_S, AfterAttnScale_before_dS, - S_mul_dS_minus_O_dO, multiplyDesc); - - ops->push_back(std::move(sub_op)); - ops->push_back(std::move(mutliply_attn_scale_op)); - ops->push_back(std::move(mutliply_op)); - - return S_mul_dS_minus_O_dO; -} - -static cudnn_frontend::Tensor createdSKBMM(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, - int64_t d, std::vector* ops, - const cudnn_frontend::Tensor& dSTensor, - const cudnn_frontend::Tensor& kTensor, - const cudnn_frontend::Tensor& mnkOverride) { - // Creates the necessary tensor descriptors - int64_t after_dSK_dim[4] = {b, h, s_kv, d}; - int64_t after_dSK_stride[4] = {h * s_kv * d, d, h * d, 1}; - // dS * K - auto After_dS_K = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 875, - after_dSK_dim, after_dSK_stride, true, false); // is virtual - - // Define the matmul desc - auto matmulDesc = cudnn_frontend::MatMulDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setPaddingValue(0) - .build(); - - // Create a matmul Node - auto matmulOp = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_MATMUL_DESCRIPTOR) - .setaMatDesc(dSTensor) - .setbMatDesc(kTensor) - .setcMatDesc(After_dS_K) - .setmOverrideDesc(mnkOverride) - .setkOverrideDesc(mnkOverride) - .setmatmulDesc(matmulDesc) - .build(); - - ops->push_back(std::move(matmulOp)); - - return After_dS_K; -} - -static cudnn_frontend::Tensor createdSQBMM(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, - int64_t d, NVTE_QKV_Layout layout, - std::vector* ops, - const cudnn_frontend::Tensor& dSTensor, - const cudnn_frontend::Tensor& qTensor, - const cudnn_frontend::Tensor& mnkOverride) { - // Creates the necessary tensor descriptors - int64_t dS_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, dS_stride, layout, NVTE_QKV_Matrix::NVTE_S_Matrix); - - int64_t dS_transpose_dim[4] = {b, h, s_kv, s_q}; - int64_t dS_transpose_stride[4]; - dS_transpose_stride[0] = dS_stride[0]; - dS_transpose_stride[1] = dS_stride[1]; - dS_transpose_stride[2] = dS_stride[3]; - dS_transpose_stride[3] = dS_stride[2]; - - int64_t after_dSTranspose_Q_dim[4] = {b, h, s_kv, d}; - int64_t after_dSTranspose_Q_stride[4] = {h * s_kv * d, d, h * d, 1}; - - auto dSTransposeTensor = - tensor_create(CUDNN_DATA_FP8_E5M2, tensor_name_to_uid["VIRTUAL"] + 650, dS_transpose_dim, - dS_transpose_stride, true, false); // is virtual - - // dS.T * Q - auto After_dSTranspose_Q = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 651, after_dSTranspose_Q_dim, - after_dSTranspose_Q_stride, true, false); // is virtual - - // Create reshape node for V -> V.T - auto reshape_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_RESHAPE_DESCRIPTOR) - .setxDesc(dSTensor) - .setyDesc(dSTransposeTensor) - .build(); - - // Define the matmul desc - auto matmulDesc = cudnn_frontend::MatMulDescBuilder() - .setComputeType(CUDNN_DATA_FLOAT) - .setPaddingValue(0) - .build(); - - // Create a matmul Node - auto matmulOp = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_MATMUL_DESCRIPTOR) - .setaMatDesc(dSTransposeTensor) - .setbMatDesc(qTensor) - .setcMatDesc(After_dSTranspose_Q) - .setmOverrideDesc(mnkOverride) - .setkOverrideDesc(mnkOverride) - .setmatmulDesc(matmulDesc) - .build(); - - ops->push_back(std::move(reshape_op)); - ops->push_back(std::move(matmulOp)); - - return After_dSTranspose_Q; -} - -// fused attention FWD FP8 with FE 0.9 -void fused_attn_fp8_fwd_impl(int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, - bool isTraining, float attnScale, float dropoutProbability, - NVTE_QKV_Layout layout, void* devPtrQ, void* devPtrK, void* devPtrV, - void* devPtrM, void* devPtrZInv, void* devPtrO, void* devPtrDescaleQ, - void* devPtrDescaleK, void* devPtrDescaleV, void* devPtrDescaleS, - void* devPtrScaleS, void* devPtrScaleO, void* devPtrAmaxO, - void* devPtrAmaxS, void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, - void* devPtrDropoutSeed, void* devPtrDropoutOffset, - cudnnDataType_t tensorType, void* workspace_ptr, - size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle_) { - try { - FADescriptor descriptor{b, - h, - s_q, - s_kv, - d, - attnScale, - isTraining, - dropoutProbability, - layout, - NVTE_Bias_Type::NVTE_NO_BIAS, - NVTE_Mask_Type::NVTE_PADDING_MASK, - tensorType, - false}; - - using CacheType = std::map; - static thread_local CacheType fa_fprop_cache; - - // Get plan from cache if cache is available, otherwise create one - auto get_plan = [&](CacheType& cache, const FADescriptor& descriptor) { - // If hit, return - auto it = cache.find(descriptor); - if (it != cache.end()) { - auto plan = it->second; - return plan; - } - - // Otherwise, build the op_graph and the plan. Then update cache - std::vector all_ops; - std::vector ops; - - NVTE_CHECK(dropoutProbability == 0.0f || isTraining, - "Dropout probability should be 0.0f for inference mode"); - NVTE_CHECK(dropoutProbability != 1.0f, "Dropout probability cannot be 1.0"); - - int64_t raggedDim[4] = {b + 1, 1, 1, 1}; - int64_t raggedStride[4] = {1, 1, 1, 1}; - // Create offset tensors - auto QKVOffsetTensor = tensor_create(CUDNN_DATA_INT32, tensor_name_to_uid["QKV_RAGGED"], - raggedDim, raggedStride, false, false); - auto ORaggedOffsetTensor = tensor_create(CUDNN_DATA_INT32, tensor_name_to_uid["O_RAGGED"], - raggedDim, raggedStride, false, false); - - int64_t seqlen_dim[4] = {b, 1, 1, 1}; - int64_t seqlen_stride[4] = {1, 1, 1, 1}; - // Create override tensors - auto seqlenMNKTensor = tensor_create(CUDNN_DATA_INT32, tensor_name_to_uid["MNK_OVERRIDE"], - seqlen_dim, seqlen_stride, false, false); - - // Create shared ptrs to ragged offset tensors - // for multiple tensors to use ragged offset - std::shared_ptr QKVRaggedOffsetTensorPtr = - std::make_shared(std::move(QKVOffsetTensor)); - std::shared_ptr ORaggedOffsetTensorPtr = - std::make_shared(std::move(ORaggedOffsetTensor)); - - // Create Q and K tensors that are used in different places - int64_t q_dim[4] = {b, h, s_q, d}; - int64_t q_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, q_stride, layout, NVTE_QKV_Matrix::NVTE_Q_Matrix); - - int64_t k_dim[4] = {b, h, s_kv, d}; - int64_t k_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, k_stride, layout, NVTE_QKV_Matrix::NVTE_K_Matrix); - - auto qTensor = tensor_create_with_offset(tensorType, tensor_name_to_uid["Q"], q_dim, q_stride, - false, false, QKVRaggedOffsetTensorPtr); - auto kTensor = tensor_create_with_offset(tensorType, tensor_name_to_uid["K"], k_dim, k_stride, - false, false, QKVRaggedOffsetTensorPtr); - - // Q * K.T - auto afterQKTensor = createQKBMM(b, h, s_q, s_kv, d, layout, tensorType, &ops, qTensor, - kTensor, seqlenMNKTensor, QKVRaggedOffsetTensorPtr); - - // QK.T * attn scale - auto AfterAttnScale_before_dequan_Q_tensor = - createScale(afterQKTensor, // input tensor - "AttnScale", // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - true, // scale is by value - &ops); - - // QK.T * attn scale * dequant_Q - auto AfterAttnScale_before_dequan_K_tensor = - createScale(AfterAttnScale_before_dequan_Q_tensor, // input tensor - "descaleQ", // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops); - - // QK.T * attn scale * dequant_Q * dequant_K - auto AfterAttnScale_tensor = - createScale(AfterAttnScale_before_dequan_K_tensor, // input tensor - "descaleK", // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops); - - auto BeforeDropoutTensor = - createSoftmaxForward(b, h, s_q, s_kv, &ops, AfterAttnScale_tensor, isTraining); - - auto AfterDropout_before_quan_S = - createDropoutForward(b, h, s_q, s_kv, dropoutProbability, &ops, BeforeDropoutTensor); - - // Amax for S - createAmax("amaxS", BeforeDropoutTensor, &ops); - - // After softmax * dropout * scale S -> fp8 input to next bmm with V - auto AfterMultiplyDropout = createScale(AfterDropout_before_quan_S, // input tensor - "scaleS", // scale tensor - tensorType, // output tensor type - true, // output is virtual - false, // scale is by value - &ops); - - // After softmax * Dropout * V - auto OTensor_before_dequan_S_tensor = - createSVBMM(b, h, s_q, s_kv, d, layout, tensorType, &ops, AfterMultiplyDropout, - seqlenMNKTensor, QKVRaggedOffsetTensorPtr); - - // O * dequant_S - auto OTensor_before_dequan_V_tensor = - createScale(OTensor_before_dequan_S_tensor, // input tensor - "descaleS", // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops); - - // O * dequant_S * dequant_V - auto OTensor_before_quan_O_tensor = - createScale(OTensor_before_dequan_V_tensor, // input tensor - "descaleV", // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops); - - // O * dequant_S * dequant_V * scale O - auto OTensor = createScaleWithOffset(OTensor_before_quan_O_tensor, // input tensor - "scaleO", // scale tensor - layout, // qkv layout - tensorType, // output tensor type - false, // output not virtual - false, // scale is by value - &ops, - ORaggedOffsetTensorPtr, // ragged offset - "O"); - - // Amax for O - createAmax("amaxO", OTensor_before_quan_O_tensor, &ops); - - for (unsigned int i = 0; i < ops.size(); i++) { - all_ops.push_back(&ops[i]); - } - - // Create an Operation Graph - auto opGraph = cudnn_frontend::OperationGraphBuilder() - .setHandle(handle_) - .setOperationGraph(all_ops.size(), all_ops.data()) - .build(); - - cudnn_frontend::EngineConfigList filtered_configs; - auto statuses = cudnn_frontend::get_heuristics_list<1>( - {"heuristics_instant"}, opGraph, allowAllConfig, filtered_configs, true); - - if (filtered_configs.size() == 0) { - cudnn_frontend::set_error_and_throw_exception( - nullptr, CUDNN_STATUS_NOT_SUPPORTED, - "run_mha_fprop: No config returned by the heuristics"); - } - - auto plan = cudnn_frontend::ExecutionPlanBuilder() - .setHandle(handle_) - .setEngineConfig(filtered_configs[0], opGraph.getTag()) - .build(); - cache.insert({descriptor, plan}); - return plan; - }; // end of get_plan - - auto plan = get_plan(fa_fprop_cache, descriptor); - size_t wkspace_size = static_cast(plan.getWorkspaceSize()); - - // Exit to request upper level API to allocate memory if needed - if (workspace_ptr == nullptr) { - *workspace_size = wkspace_size + ((b + 1) * 2 + b) * sizeof(int32_t); - return; - } - - // cuDNN stream check needs to be moved here to support dummy kernel calls with - // null streams for sizing the cuDNN workspace. - NVTE_CHECK_CUDNN(cudnnSetStream(handle_, stream)); - - int32_t* qkv_ragged_offset = - reinterpret_cast(reinterpret_cast(workspace_ptr) + wkspace_size); - int32_t* o_ragged_offset = reinterpret_cast(reinterpret_cast(workspace_ptr) + - wkspace_size + (b + 1) * sizeof(int32_t)); - int32_t* actual_seqlens_q = reinterpret_cast( - reinterpret_cast(workspace_ptr) + wkspace_size + (b + 1) * 2 * sizeof(int32_t)); - // FP8 currently only supports self-attention, so doesn't use devPtrcuSeqlensKV - dim3 blockDims(128); - dim3 gridDims((b + blockDims.x) / blockDims.x); - cu_seqlens_to_offsets<<>>( - b, h, d, reinterpret_cast(devPtrcuSeqlensQ), actual_seqlens_q, qkv_ragged_offset, - o_ragged_offset); - NVTE_CHECK_CUDA(cudaGetLastError()); - void* devPtrQKVRaggedOffset = reinterpret_cast(qkv_ragged_offset); - void* devPtrORaggedOffset = reinterpret_cast(o_ragged_offset); - void* devPtrMNKOverride = reinterpret_cast(actual_seqlens_q); - - float dropoutScale = 1.0f / (1.0f - dropoutProbability); - - std::set> data_ptrs; - data_ptrs.emplace(std::pair(tensor_name_to_uid["Q"], devPtrQ)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["K"], devPtrK)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["K_TRANSPOSE"], devPtrK)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["V"], devPtrV)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["AttnScale"], &attnScale)); - data_ptrs.emplace( - std::pair(tensor_name_to_uid["DROPOUT_SCALE"], &dropoutScale)); - data_ptrs.emplace( - std::pair(tensor_name_to_uid["DROPOUT_SEED"], devPtrDropoutSeed)); - data_ptrs.emplace( - std::pair(tensor_name_to_uid["DROPOUT_OFFSET"], devPtrDropoutOffset)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["O"], devPtrO)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["descaleQ"], devPtrDescaleQ)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["descaleK"], devPtrDescaleK)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["descaleV"], devPtrDescaleV)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["descaleS"], devPtrDescaleS)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["scaleS"], devPtrScaleS)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["scaleO"], devPtrScaleO)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["amaxO"], devPtrAmaxO)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["amaxS"], devPtrAmaxS)); - data_ptrs.emplace( - std::pair(tensor_name_to_uid["QKV_RAGGED"], devPtrQKVRaggedOffset)); - data_ptrs.emplace( - std::pair(tensor_name_to_uid["O_RAGGED"], devPtrORaggedOffset)); - data_ptrs.emplace( - std::pair(tensor_name_to_uid["MNK_OVERRIDE"], devPtrMNKOverride)); - - // If training, then we need to write out M and Z_INV - if (isTraining) { - data_ptrs.emplace(std::pair(tensor_name_to_uid["M"], devPtrM)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["Z_INV"], devPtrZInv)); - } - - auto variantPack = cudnn_frontend::VariantPackBuilder() - .setWorkspacePointer(workspace_ptr) - .setDataPointers(data_ptrs) - .build(); - - NVTE_CHECK_CUDNN(cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc())); - } catch (cudnn_frontend::cudnnException& e) { - struct cudaDeviceProp prop; - NVTE_CHECK_CUDA(cudaGetDeviceProperties(&prop, 0)); - - // This example is only for GH100 cards (cudnn Version >= 8900) - if (!((prop.major == 9 && prop.minor == 0 && CUDNN_VERSION >= 8900)) && - (e.getCudnnStatus() == CUDNN_STATUS_ARCH_MISMATCH || - e.getCudnnStatus() == CUDNN_STATUS_NOT_SUPPORTED)) { - std::cout << "Example is only supported for GH100 (cuDNN >= 8900) GPUs" << std::endl; - } else { - std::cout << "[ERROR] Exception " << e.what() << std::endl; - } - } -} - -// fused attention BWD FP8 with FE 0.9 -void fused_attn_fp8_bwd_impl( - int64_t b, int64_t h, int64_t s_q, int64_t s_kv, int64_t d, float attnScale, - float dropoutProbability, NVTE_QKV_Layout layout, void* devPtrQ, void* devPtrK, void* devPtrV, - void* devPtrM, void* devPtrZInv, void* devPtrO, void* devPtrdO, void* devPtrdQ, void* devPtrdK, - void* devPtrdV, void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, - void* devPtrDescaleO, void* devPtrDescaledO, void* devPtrDescaleS, void* devPtrDescaledS, - void* devPtrScaleS, void* devPtrScaledS, void* devPtrScaledQ, void* devPtrScaledK, - void* devPtrScaledV, void* devPtrAmaxdS, void* devPtrAmaxdQ, void* devPtrAmaxdK, - void* devPtrAmaxdV, void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, void* devPtrDropoutSeed, - void* devPtrDropoutOffset, cudnnDataType_t tensorType, void* workspace_ptr, - size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle_) { - try { - FADescriptor descriptor{b, - h, - s_q, - s_kv, - d, - attnScale, - false, - dropoutProbability, - layout, - NVTE_Bias_Type::NVTE_NO_BIAS, - NVTE_Mask_Type::NVTE_PADDING_MASK, - tensorType, - false}; - - using CacheType = std::map; - static thread_local CacheType fa_bprop_cache; - - // Get plan from cache if cache is available, otherwise create one - auto get_plan = [&](CacheType& cache, const FADescriptor& descriptor) { - // If hit, return - auto it = cache.find(descriptor); - if (it != cache.end()) { - auto plan = it->second; - return plan; - } - - // Otherwise, build the op_graph and the plan. Then update cache - std::vector all_ops; - std::vector ops; - - NVTE_CHECK(dropoutProbability != 1.0f, "Dropout probability cannot be 1.0"); - - int64_t raggedDim[4] = {b + 1, 1, 1, 1}; - int64_t raggedStride[4] = {1, 1, 1, 1}; - // Create offset tensors - auto QKVOffsetTensor = tensor_create(CUDNN_DATA_INT32, tensor_name_to_uid["QKV_RAGGED"], - raggedDim, raggedStride, false, false); - auto ORaggedOffsetTensor = tensor_create(CUDNN_DATA_INT32, tensor_name_to_uid["O_RAGGED"], - raggedDim, raggedStride, false, false); - - // Create shared ptrs to ragged offset tensors for multiple tensors - std::shared_ptr QKVRaggedOffsetTensorPtr = - std::make_shared(std::move(QKVOffsetTensor)); - std::shared_ptr ORaggedOffsetTensorPtr = - std::make_shared(std::move(ORaggedOffsetTensor)); - - // Create Q and K tensors that are used in different places - int64_t q_dim[4] = {b, h, s_q, d}; - int64_t q_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, q_stride, layout, NVTE_QKV_Matrix::NVTE_Q_Matrix); - - int64_t k_dim[4] = {b, h, s_kv, d}; - int64_t k_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, k_stride, layout, NVTE_QKV_Matrix::NVTE_K_Matrix); - - auto qTensor = tensor_create_with_offset(tensorType, tensor_name_to_uid["Q"], q_dim, q_stride, - false, false, QKVRaggedOffsetTensorPtr); - auto kTensor = tensor_create_with_offset(tensorType, tensor_name_to_uid["K"], k_dim, k_stride, - false, false, QKVRaggedOffsetTensorPtr); - - int64_t scale_dim[4] = {1, 1, 1, 1}; - int64_t scale_stride[4] = {1, 1, 1, 1}; - - // Create attnScale tensor for multiple ops to use - auto attnScaleTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["AttnScale"], - scale_dim, scale_stride, false, true); // is by value - - // Create descale Q K dO dS global tensors since they are used in multiple places - auto descaleQTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["descaleQ"], - scale_dim, scale_stride, false, false); - auto descaleKTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["descaleK"], - scale_dim, scale_stride, false, false); - auto descaledOTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["descaledO"], - scale_dim, scale_stride, false, false); - auto descaledSTensor = tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["descaledS"], - scale_dim, scale_stride, false, false); - - int64_t seqlen_dim[4] = {b, 1, 1, 1}; - int64_t seqlen_stride[4] = {1, 1, 1, 1}; - // Create MNK override tensor - auto seqlenMNKTensor = tensor_create(CUDNN_DATA_INT32, tensor_name_to_uid["MNK_OVERRIDE"], - seqlen_dim, seqlen_stride, false, false); - - int64_t O_dim[4] = {b, h, s_q, d}; - int64_t O_stride[4]; - generateMatrixStrides(b, h, s_q, s_kv, d, O_stride, layout, NVTE_QKV_Matrix::NVTE_O_Matrix); - // Create O and loss tensor - auto OTensor = tensor_create_with_offset(tensorType, tensor_name_to_uid["O"], O_dim, O_stride, - false, false, ORaggedOffsetTensorPtr); - // dO is used in multiple places and E5M2 - auto dOTensor = - tensor_create_with_offset(CUDNN_DATA_FP8_E5M2, tensor_name_to_uid["dO"], O_dim, O_stride, - false, false, ORaggedOffsetTensorPtr); - - // Q * K.T - auto afterQKTensor = createQKBMM(b, h, s_q, s_kv, d, layout, tensorType, &ops, qTensor, - kTensor, seqlenMNKTensor, QKVRaggedOffsetTensorPtr); - - // QK.T * attn scale - auto AfterAttnScale_before_dequan_Q_tensor = - createScale(afterQKTensor, // input tensor - attnScaleTensor, // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - true, // scale is by value - &ops, 1999 /*UID offset*/); - - // QK.T * attn scale * dequant_Q - auto AfterAttnScale_before_dequan_K_tensor = - createScale(AfterAttnScale_before_dequan_Q_tensor, // input tensor - descaleQTensor, // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops, 2000 /*UID offset*/); - - // QK.T * attn scale * dequant_Q * dequant_K - auto AfterAttnScale_tensor = - createScale(AfterAttnScale_before_dequan_K_tensor, // input tensor - descaleKTensor, // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops, 2001 /*UID offset*/); - - auto beforeDropout_QKt_Tensor = - createSoftmaxBackward(b, h, s_q, s_kv, &ops, AfterAttnScale_tensor); - - int64_t afterBMM1_dim[4] = {b, h, s_q, s_kv}; - int64_t afterBMM1_stride[4] = {h * s_q * s_kv, s_q * s_kv, s_kv, 1}; - - // mask for the dropout. Used in different places - auto dropoutMaskTensor = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 200, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - - auto AfterDropout_before_quan_S = createDropoutBackward( - b, h, s_q, s_kv, dropoutProbability, &ops, beforeDropout_QKt_Tensor, dropoutMaskTensor); - - // After softmax * scale S -> fp8 input to next bmm with V - auto AfterMultiply = createScale(AfterDropout_before_quan_S, // input tensor - "scaleS", // scale tensor - tensorType, // output tensor type - true, // output is virtual - false, // scale is by value - &ops); - - // After softmax * dO - auto dVTensor_before_dequan_S = createSdOBMM(b, h, s_q, s_kv, d, tensorType, &ops, - AfterMultiply, dOTensor, seqlenMNKTensor); - - // O * dequant_S - auto dVTensor_before_dequan_dO = createScale(dVTensor_before_dequan_S, // input tensor - "descaleS", // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops); - - // O * dequant_S * dequant_dO - auto dVTensor_before_quan_dV = createScale(dVTensor_before_dequan_dO, // input tensor - descaledOTensor, // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops, 2002 /*UID offset*/); - - // O * dequant_S * dequant_dO * scale dV - auto dVTensor = createScaleWithOffset(dVTensor_before_quan_dV, // input tensor - "scaledV", // scale tensor - layout, // qkv layout - CUDNN_DATA_FP8_E5M2, // output tensor type - false, // output not virtual - false, // scale is by value - &ops, - QKVRaggedOffsetTensorPtr, // ragged offset - "dV" /*Output tensor name*/); - - // Amax for dV - createAmax("amaxdV", dVTensor_before_quan_dV, &ops); - - auto dS_before_dequan_dO_Tensor = - createdOVBMM(b, h, s_q, s_kv, d, layout, tensorType, &ops, dOTensor, seqlenMNKTensor, - QKVRaggedOffsetTensorPtr); - - // dS * dequant_dO - auto dS_before_dequan_V = createScale(dS_before_dequan_dO_Tensor, // input tensor - descaledOTensor, // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops, 2003 /*UID offset*/); - - // O * dequant_S * dequant_dV - auto dS_after_dequan = createScale(dS_before_dequan_V, // input tensor - "descaleV", // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops); - - // RNG Multiply - auto beforeDropoutScale_dOVt_Tensor = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 350, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - // After dropout mask and scale - auto dS_after_dropout = - tensor_create(CUDNN_DATA_FLOAT, tensor_name_to_uid["VIRTUAL"] + 351, afterBMM1_dim, - afterBMM1_stride, true, false); // is virtual - - // Define the multiply mask descriptor - auto mulDesc = pw_desc_create(CUDNN_DATA_FLOAT, CUDNN_POINTWISE_MUL); - - // Create a multiply mask Node - auto maskMul_op = binary_pw_op_create(dS_after_dequan, dropoutMaskTensor, - beforeDropoutScale_dOVt_Tensor, mulDesc); - - ops.push_back(std::move(maskMul_op)); - - // scale after dropout for dO and O chain - auto dropoutScale_dOVt_OdO_Tensor = - tensor_create(tensorType, tensor_name_to_uid["DROPOUT_SCALE_dOVt_OdO"], scale_dim, - scale_stride, false, true); // is by value - - // Create a multiply dropout scale Node - auto mul_dropout_scale_op = binary_pw_op_create( - beforeDropoutScale_dOVt_Tensor, dropoutScale_dOVt_OdO_Tensor, dS_after_dropout, mulDesc); - - ops.push_back(std::move(mul_dropout_scale_op)); - - // O * dequant_O - auto O_after_dequan_Tensor = createScale(OTensor, // input tensor - "descaleO", // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops); - - // dO * dequant_dO - auto dO_after_dequan_Tensor = createScale(dOTensor, // input tensor - descaledOTensor, // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops, 2004 /*UID offset*/); - - // row reduction sum[(dO * dequant_dO) * (O * dequant_O) * (1 - p)] - auto O_dO_after_rowsum = - createdOAndORowReductionChain(b, h, s_q, s_kv, d, layout, &ops, O_after_dequan_Tensor, - dO_after_dequan_Tensor, dropoutScale_dOVt_OdO_Tensor); - - // (dS_after_dropout - O_dO_after_rowsum) * AfterDropout_before_quan_S * attnScale - auto S_mul_dS_minus_O_dO = createBiasSubtractionSoftmaxMulChain( - b, h, s_q, s_kv, d, layout, &ops, dS_after_dropout, AfterDropout_before_quan_S, - O_dO_after_rowsum, attnScaleTensor); - - // S_mul_dS_minus_O_dO * scaledS - auto S_mul_dS_minus_O_dO_after_quan_dS = - createScale(S_mul_dS_minus_O_dO, // input tensor - "scaledS", // scale tensor - CUDNN_DATA_FP8_E5M2, // output tensor type - true, // output is virtual - false, // scale is by value - &ops); - - // Amax for dS - createAmax("amaxdS", S_mul_dS_minus_O_dO, &ops); - - // dS @ K - auto After_dS_K = createdSKBMM(b, h, s_q, s_kv, d, &ops, S_mul_dS_minus_O_dO_after_quan_dS, - kTensor, seqlenMNKTensor); - - // (dS * K) * descale dS - auto After_dS_K_before_dequan_K = createScale(After_dS_K, // input tensor - descaledSTensor, // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops, 2006 /*UID offset*/); - - // (dS * K) * descale dS * descale K - auto After_dS_K_before_quan_dQ = createScale(After_dS_K_before_dequan_K, // input tensor - descaleKTensor, // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops, 2007 /*UID offset*/); - - // (dS * K) * descale dS * descale K * scale dQ - auto dQ = createScaleWithOffset(After_dS_K_before_quan_dQ, // input tensor - "scaledQ", // scale tensor - layout, // qkv layout - CUDNN_DATA_FP8_E5M2, // output tensor type - false, // output not virtual - false, // scale is by value - &ops, - QKVRaggedOffsetTensorPtr, // ragged offset - "dQ"); - - // Amax for dQ - createAmax("amaxdQ", After_dS_K_before_quan_dQ, &ops); - - // dS.T @ Q - auto After_dSTranspose_Q = - createdSQBMM(b, h, s_q, s_kv, d, layout, &ops, S_mul_dS_minus_O_dO_after_quan_dS, qTensor, - seqlenMNKTensor); - - // (dS.T * Q) * descale dS - auto After_dSTranspose_Q_before_dequan_Q = - createScale(After_dSTranspose_Q, // input tensor - descaledSTensor, // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops, 2009 /*UID offset*/); - - // (dS.T * Q) * descale dS * descale Q - auto After_dSTranspose_Q_before_quan_dK = - createScale(After_dSTranspose_Q_before_dequan_Q, // input tensor - descaleQTensor, // scale tensor - CUDNN_DATA_FLOAT, // output tensor type - true, // output is virtual - false, // scale is by value - &ops, 2010 /*UID offset*/); - - // (dS.T * Q) * descale dS * descale Q * scale dK - auto dK = createScaleWithOffset(After_dSTranspose_Q_before_quan_dK, // input tensor - "scaledK", // scale tensor - layout, // qkv layout - CUDNN_DATA_FP8_E5M2, // output tensor type - false, // output not virtual - false, // scale is by value - &ops, - QKVRaggedOffsetTensorPtr, // ragged offset - "dK"); - - // Amax for dK - createAmax("amaxdK", After_dSTranspose_Q_before_quan_dK, &ops); - - for (unsigned int i = 0; i < ops.size(); i++) { - all_ops.push_back(&ops[i]); - } - - // Create an Operation Graph - auto opGraph = cudnn_frontend::OperationGraphBuilder() - .setHandle(handle_) - .setOperationGraph(all_ops.size(), all_ops.data()) - .build(); - - cudnn_frontend::EngineConfigList filtered_configs; - auto statuses = cudnn_frontend::get_heuristics_list<1>( - {"heuristics_instant"}, opGraph, allowAllConfig, filtered_configs, true); - - if (filtered_configs.size() == 0) { - cudnn_frontend::set_error_and_throw_exception( - nullptr, CUDNN_STATUS_NOT_SUPPORTED, - "run_mha_bprop: No config returned by the heuristics"); - } - - auto plan = cudnn_frontend::ExecutionPlanBuilder() - .setHandle(handle_) - .setEngineConfig(filtered_configs[0], opGraph.getTag()) - .build(); - cache.insert({descriptor, plan}); - return plan; - }; - - auto plan = get_plan(fa_bprop_cache, descriptor); - size_t wkspace_size = static_cast(plan.getWorkspaceSize()); - - // Exit to request upper level API to allocate memory if needed - if (workspace_ptr == nullptr) { - *workspace_size = wkspace_size + ((b + 1) * 2 + b) * sizeof(int32_t); - return; - } - - // cuDNN stream check needs to be moved here to support dummy kernel calls with - // null streams for sizing the cuDNN workspace. - NVTE_CHECK_CUDNN(cudnnSetStream(handle_, stream)); - - int32_t* qkv_ragged_offset = - reinterpret_cast(reinterpret_cast(workspace_ptr) + wkspace_size); - int32_t* o_ragged_offset = reinterpret_cast(reinterpret_cast(workspace_ptr) + - wkspace_size + (b + 1) * sizeof(int32_t)); - int32_t* actual_seqlens_q = reinterpret_cast( - reinterpret_cast(workspace_ptr) + wkspace_size + (b + 1) * 2 * sizeof(int32_t)); - // FP8 currently only supports self-attention, so doesn't use devPtrcuSeqlensKV - dim3 blockDims(128); - dim3 gridDims((b + blockDims.x) / blockDims.x); - cu_seqlens_to_offsets<<>>( - b, h, d, reinterpret_cast(devPtrcuSeqlensQ), actual_seqlens_q, qkv_ragged_offset, - o_ragged_offset); - NVTE_CHECK_CUDA(cudaGetLastError()); - void* devPtrQKVRaggedOffset = reinterpret_cast(qkv_ragged_offset); - void* devPtrORaggedOffset = reinterpret_cast(o_ragged_offset); - void* devPtrMNKOverride = reinterpret_cast(actual_seqlens_q); - - std::set> data_ptrs; - float dropoutScale = 1.0f / (1.0f - dropoutProbability); - float dropoutScale_dOVt_OdO = 1.0f - dropoutProbability; - data_ptrs.emplace(std::pair(tensor_name_to_uid["Q"], devPtrQ)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["K"], devPtrK)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["K_TRANSPOSE"], devPtrK)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["V"], devPtrV)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["V_TRANSPOSE"], devPtrV)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["dQ"], devPtrdQ)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["dK"], devPtrdK)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["dV"], devPtrdV)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["dO"], devPtrdO)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["AttnScale"], &attnScale)); - data_ptrs.emplace( - std::pair(tensor_name_to_uid["DROPOUT_SCALE"], &dropoutScale)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["DROPOUT_SCALE_dOVt_OdO"], - &dropoutScale_dOVt_OdO)); - data_ptrs.emplace( - std::pair(tensor_name_to_uid["DROPOUT_SEED"], devPtrDropoutSeed)); - data_ptrs.emplace( - std::pair(tensor_name_to_uid["DROPOUT_OFFSET"], devPtrDropoutOffset)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["M"], devPtrM)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["Z_INV"], devPtrZInv)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["O"], devPtrO)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["descaleQ"], devPtrDescaleQ)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["descaleK"], devPtrDescaleK)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["descaleV"], devPtrDescaleV)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["descaleS"], devPtrDescaleS)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["descaledS"], devPtrDescaledS)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["descaleO"], devPtrDescaleO)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["descaledO"], devPtrDescaledO)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["scaleS"], devPtrScaleS)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["scaledS"], devPtrScaledS)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["scaledQ"], devPtrScaledQ)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["scaledK"], devPtrScaledK)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["scaledV"], devPtrScaledV)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["amaxdS"], devPtrAmaxdS)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["amaxdQ"], devPtrAmaxdQ)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["amaxdK"], devPtrAmaxdK)); - data_ptrs.emplace(std::pair(tensor_name_to_uid["amaxdV"], devPtrAmaxdV)); - data_ptrs.emplace( - std::pair(tensor_name_to_uid["QKV_RAGGED"], devPtrQKVRaggedOffset)); - data_ptrs.emplace( - std::pair(tensor_name_to_uid["O_RAGGED"], devPtrORaggedOffset)); - data_ptrs.emplace( - std::pair(tensor_name_to_uid["MNK_OVERRIDE"], devPtrMNKOverride)); - - auto variantPack = cudnn_frontend::VariantPackBuilder() - .setWorkspacePointer(workspace_ptr) - .setDataPointers(data_ptrs) - .build(); - NVTE_CHECK_CUDNN(cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc())); - } catch (cudnn_frontend::cudnnException& e) { - struct cudaDeviceProp prop; - NVTE_CHECK_CUDA(cudaGetDeviceProperties(&prop, 0)); - - // This example is only for GH100 cards (cudnn Version >= 8900) - if (!((prop.major == 9 && prop.minor == 0 && CUDNN_VERSION >= 8900)) && - (e.getCudnnStatus() == CUDNN_STATUS_ARCH_MISMATCH || - e.getCudnnStatus() == CUDNN_STATUS_NOT_SUPPORTED)) { - std::cout << "Example is only supported for GH100 (cuDNN >= 8900) GPUs" << std::endl; - } else { - std::cout << "[ERROR] Exception " << e.what() << std::endl; - } - } -} - // fused attention FWD FP8 with FE 1.0+ -void fused_attn_fp8_fwd_impl_v1( +void fused_attn_fp8_fwd_impl( int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, bool is_training, float scaling_factor, float dropout_probability, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, void* devPtrQ, void* devPtrK, void* devPtrV, - void* devPtrSoftmaxOffset, void* devPtrM, void* devPtrZInv, void* devPtrO, void* devPtrDescaleQ, + void* devPtrSoftmaxOffset, void* devPtrM, void* devPtrO, void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, void* devPtrDescaleS, void* devPtrScaleS, void* devPtrScaleO, void* devPtrAmaxO, void* devPtrAmaxS, void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, @@ -2080,26 +446,26 @@ void fused_attn_fp8_fwd_impl_v1( } // fused attention BWD FP8 with FE 1.0+ -void fused_attn_fp8_bwd_impl_v1( +void fused_attn_fp8_bwd_impl( int64_t b, int64_t h, int64_t hg, int64_t s_q, int64_t s_kv, int64_t d_qk, int64_t d_v, float scaling_factor, float dropout_probability, NVTE_QKV_Layout qkv_layout, NVTE_QKV_Format o_format, NVTE_QKV_Format do_format, NVTE_QKV_Layout dqkv_layout, NVTE_Bias_Type bias_type, NVTE_Mask_Type mask_type, NVTE_Softmax_Type softmax_type, int64_t window_size_left, int64_t window_size_right, bool bottom_right_diagonal, - bool deterministic, void* devPtrQ, void* devPtrK, void* devPtrV, void* devPtrM, - void* devPtrZInv, void* devPtrO, void* devPtrdO, void* devPtrSoftmaxOffset, void* devPtrdQ, - void* devPtrdK, void* devPtrdV, void* devPtrdSoftmaxOffset, void* devPtrDescaleQ, - void* devPtrDescaleK, void* devPtrDescaleV, void* devPtrDescaleO, void* devPtrDescaledO, - void* devPtrDescaleS, void* devPtrDescaledP, void* devPtrScaleS, void* devPtrScaledP, - void* devPtrScaledQ, void* devPtrScaledK, void* devPtrScaledV, void* devPtrAmaxdP, - void* devPtrAmaxdQ, void* devPtrAmaxdK, void* devPtrAmaxdV, void* devPtrQ_t, void* devPtrK_t, - void* devPtrdO_f16, void* devPtrdO_t, void* devPtrDescaleQ_t, void* devPtrDescaleK_t, - void* devPtrDescaledO_t, void* devPtrcuSeqlensQ, void* devPtrcuSeqlensKV, - void* devPtrDropoutSeed, void* devPtrDropoutOffset, cudnn_frontend::DataType_t qkv_tensor_type, - cudnn_frontend::DataType_t o_tensor_type, cudnn_frontend::DataType_t do_tensor_type, - cudnn_frontend::DataType_t dqkv_tensor_type, NVTEScalingMode scaling_mode, - NVTE_QKV_Format qkv_scale_inv_format, NVTE_QKV_Format do_scale_inv_format, void* workspace, - size_t* workspace_size, cudaStream_t stream, cudnnHandle_t handle) { + bool deterministic, void* devPtrQ, void* devPtrK, void* devPtrV, void* devPtrM, void* devPtrO, + void* devPtrdO, void* devPtrSoftmaxOffset, void* devPtrdQ, void* devPtrdK, void* devPtrdV, + void* devPtrdSoftmaxOffset, void* devPtrDescaleQ, void* devPtrDescaleK, void* devPtrDescaleV, + void* devPtrDescaleO, void* devPtrDescaledO, void* devPtrDescaleS, void* devPtrDescaledP, + void* devPtrScaleS, void* devPtrScaledP, void* devPtrScaledQ, void* devPtrScaledK, + void* devPtrScaledV, void* devPtrAmaxdP, void* devPtrAmaxdQ, void* devPtrAmaxdK, + void* devPtrAmaxdV, void* devPtrQ_t, void* devPtrK_t, void* devPtrdO_f16, void* devPtrdO_t, + void* devPtrDescaleQ_t, void* devPtrDescaleK_t, void* devPtrDescaledO_t, void* devPtrcuSeqlensQ, + void* devPtrcuSeqlensKV, void* devPtrDropoutSeed, void* devPtrDropoutOffset, + cudnn_frontend::DataType_t qkv_tensor_type, cudnn_frontend::DataType_t o_tensor_type, + cudnn_frontend::DataType_t do_tensor_type, cudnn_frontend::DataType_t dqkv_tensor_type, + NVTEScalingMode scaling_mode, NVTE_QKV_Format qkv_scale_inv_format, + NVTE_QKV_Format do_scale_inv_format, void* workspace, size_t* workspace_size, + cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; const auto cudnn_runtime_version = cudnnGetVersion(); bool is_bias = (bias_type == NVTE_Bias_Type::NVTE_POST_SCALE_BIAS); @@ -2760,19 +1126,12 @@ void fused_attn_fp8_fwd( devPtrSoftmaxOffset = input_SoftmaxOffset->data.dptr; } void* devPtrM = nullptr; - void* devPtrZInv = nullptr; if (Aux_CTX_Tensors->size == 0) { int i = 0; Tensor* output_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_M->data.dptr = nullptr; output_M->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; output_M->data.dtype = DType::kFloat32; - if (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD) { - Tensor* output_ZInv = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - output_ZInv->data.dptr = nullptr; - output_ZInv->data.shape = {batch, num_attn_heads, max_seqlen_q, 1}; - output_ZInv->data.dtype = DType::kFloat32; - } Tensor* output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_rng_state->data.dptr = nullptr; output_rng_state->data.shape = {2}; @@ -2788,11 +1147,6 @@ void fused_attn_fp8_fwd( int i = 0; Tensor* output_M = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); devPtrM = output_M->data.dptr; - devPtrZInv = nullptr; - if (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD) { - Tensor* output_ZInv = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); - devPtrZInv = output_ZInv->data.dptr; - } Tensor* output_rng_state = convertNVTETensorCheck(Aux_CTX_Tensors->tensors[i++]); output_rng_state->data.dptr = rng_state->data.dptr; if (softmax_type != NVTE_VANILLA_SOFTMAX) { @@ -2819,25 +1173,17 @@ void fused_attn_fp8_fwd( NVTE_QKV_Format qkv_format = nvte_get_qkv_format(qkv_layout); if ((qkv_format == NVTE_QKV_Format::NVTE_BSHD) || (qkv_format == NVTE_QKV_Format::NVTE_SBHD) || (qkv_format == NVTE_QKV_Format::NVTE_BHSD)) { - fused_attn::fused_attn_fp8_fwd_impl_v1( + fused_attn::fused_attn_fp8_fwd_impl( batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, is_training, attn_scale, p_dropout, qkv_layout, o_format, bias_type, mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, devPtrQ, devPtrK, - devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrZInv, devPtrO, devPtrDescaleQ, devPtrDescaleK, + devPtrV, devPtrSoftmaxOffset, devPtrM, devPtrO, devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, devPtrAmaxO, devPtrAmaxS, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), get_cudnn_fe_dtype(O_type), input_Q->scaling_mode, qkv_scale_inv_format, workspace->data.dptr, &workspace_size, stream, handle); - } else if (qkv_layout == NVTE_QKV_Layout::NVTE_T3HD) { - fused_attn::fused_attn_fp8_fwd_impl( - batch, num_attn_heads, max_seqlen_q, max_seqlen_kv, head_dim_qk, is_training, attn_scale, - p_dropout, qkv_layout, devPtrQ, devPtrK, devPtrV, devPtrM, devPtrZInv, devPtrO, - devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleS, devPtrScaleS, devPtrScaleO, - devPtrAmaxO, devPtrAmaxS, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, - devPtrDropoutOffset, get_cudnn_dtype(QKV_type), workspace->data.dptr, &workspace_size, - stream, handle); } else { - NVTE_ERROR("FP8 fused attention only supports qkv_layout=t3hd or qkv_format=bshd/sbhd. \n"); + NVTE_ERROR("FP8 fused attention only supports qkv_format=BSHD, SBHD, or BHSD.\n"); } if (workspace_size > 0) { @@ -2862,11 +1208,11 @@ void fused_attn_fp8_bwd( NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, bool bottom_right_diagonal, bool deterministic, const Tensor* input_Q, const Tensor* input_K, const Tensor* input_V, const Tensor* input_O, const Tensor* input_dO, - const Tensor* input_dO_f16, const Tensor* input_M, const Tensor* input_ZInv, - const Tensor* input_S, const Tensor* input_SoftmaxOffset, Tensor* input_output_dP, - const Tensor* output_dQ, const Tensor* output_dK, const Tensor* output_dV, - Tensor* output_dSoftmaxOffset, const Tensor* cu_seqlens_q, const Tensor* cu_seqlens_kv, - const Tensor* rng_state, Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { + const Tensor* input_dO_f16, const Tensor* input_M, const Tensor* input_S, + const Tensor* input_SoftmaxOffset, Tensor* input_output_dP, const Tensor* output_dQ, + const Tensor* output_dK, const Tensor* output_dV, Tensor* output_dSoftmaxOffset, + const Tensor* cu_seqlens_q, const Tensor* cu_seqlens_kv, const Tensor* rng_state, + Tensor* workspace, cudaStream_t stream, cudnnHandle_t handle) { using namespace transformer_engine; void* devPtrQ = input_Q->data.dptr; void* devPtrK = input_K->data.dptr; @@ -2899,7 +1245,6 @@ void fused_attn_fp8_bwd( } void* devPtrM = input_M->data.dptr; - void* devPtrZInv = (input_ZInv != nullptr) ? input_ZInv->data.dptr : nullptr; void *devPtrScaleS = nullptr, *devPtrDescaleS = nullptr, *devPtrAmaxdP = nullptr, *devPtrScaledP = nullptr, *devPtrDescaledP = nullptr; @@ -2949,34 +1294,22 @@ void fused_attn_fp8_bwd( NVTE_QKV_Format dqkv_format = nvte_get_qkv_format(dqkv_layout); if ((dqkv_format == NVTE_QKV_Format::NVTE_BSHD) || (dqkv_format == NVTE_QKV_Format::NVTE_SBHD) || (dqkv_format == NVTE_QKV_Format::NVTE_BHSD)) { - fused_attn::fused_attn_fp8_bwd_impl_v1( + fused_attn::fused_attn_fp8_bwd_impl( batch, num_attn_heads, num_gqa_groups, max_seqlen_q, max_seqlen_kv, head_dim_qk, head_dim_v, attn_scale, p_dropout, qkv_layout, o_format, do_format, dqkv_layout, bias_type, mask_type, softmax_type, window_size_left, window_size_right, bottom_right_diagonal, deterministic, - devPtrQ, devPtrK, devPtrV, devPtrM, devPtrZInv, devPtrO, devPtrdO, devPtrSoftmaxOffset, - devPtrdQ, devPtrdK, devPtrdV, devPtrdSoftmaxOffset, devPtrDescaleQ, devPtrDescaleK, - devPtrDescaleV, devPtrDescaleO, devPtrDescaledO, devPtrDescaleS, devPtrDescaledP, - devPtrScaleS, devPtrScaledP, devPtrScaledQ, devPtrScaledK, devPtrScaledV, devPtrAmaxdP, - devPtrAmaxdQ, devPtrAmaxdK, devPtrAmaxdV, devPtrQ_t, devPtrK_t, devPtrdO_f16, devPtrdO_t, + devPtrQ, devPtrK, devPtrV, devPtrM, devPtrO, devPtrdO, devPtrSoftmaxOffset, devPtrdQ, + devPtrdK, devPtrdV, devPtrdSoftmaxOffset, devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, + devPtrDescaleO, devPtrDescaledO, devPtrDescaleS, devPtrDescaledP, devPtrScaleS, + devPtrScaledP, devPtrScaledQ, devPtrScaledK, devPtrScaledV, devPtrAmaxdP, devPtrAmaxdQ, + devPtrAmaxdK, devPtrAmaxdV, devPtrQ_t, devPtrK_t, devPtrdO_f16, devPtrdO_t, devPtrDescaleQ_t, devPtrDescaleK_t, devPtrDescaledO_t, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, devPtrDropoutOffset, get_cudnn_fe_dtype(QKV_type), get_cudnn_fe_dtype(O_type), get_cudnn_fe_dtype(dO_type), get_cudnn_fe_dtype(dQKV_type), input_dO->scaling_mode, qkv_scale_inv_format, do_scale_inv_format, workspace->data.dptr, &workspace_size, stream, handle); - } else if (dqkv_layout == NVTE_QKV_Layout::NVTE_T3HD) { - // remove this when cuDNN FE supports FP8 + THD - NVTE_CHECK(input_ZInv != nullptr && input_ZInv->data.dptr != nullptr, - "ZInv tensor required for FP8 fused attention backward with T3HD layout."); - fused_attn::fused_attn_fp8_bwd_impl( - batch, num_attn_heads, max_seqlen_q, max_seqlen_kv, head_dim_qk, attn_scale, p_dropout, - qkv_layout, devPtrQ, devPtrK, devPtrV, devPtrM, devPtrZInv, devPtrO, devPtrdO, devPtrdQ, - devPtrdK, devPtrdV, devPtrDescaleQ, devPtrDescaleK, devPtrDescaleV, devPtrDescaleO, - devPtrDescaledO, devPtrDescaleS, devPtrDescaledP, devPtrScaleS, devPtrScaledP, - devPtrScaledQ, devPtrScaledK, devPtrScaledV, devPtrAmaxdP, devPtrAmaxdQ, devPtrAmaxdK, - devPtrAmaxdV, devPtrcuSeqlensQ, devPtrcuSeqlensKV, devPtrDropoutSeed, devPtrDropoutOffset, - get_cudnn_dtype(QKV_type), workspace->data.dptr, &workspace_size, stream, handle); } else { - NVTE_ERROR("FP8 fused attention only supports qkv_layout=t3hd or qkv_format=bshd/sbhd. \n"); + NVTE_ERROR("FP8 fused attention only supports dqkv_format=BSHD, SBHD, or BHSD.\n"); } if (workspace_size > 0) { diff --git a/transformer_engine/common/fused_attn/fused_attn_fp8.h b/transformer_engine/common/fused_attn/fused_attn_fp8.h index aaf5039ee..b9660128c 100644 --- a/transformer_engine/common/fused_attn/fused_attn_fp8.h +++ b/transformer_engine/common/fused_attn/fused_attn_fp8.h @@ -5,7 +5,7 @@ ************************************************************************/ /*! \file fused_attn_fp8.h - * \brief Functions for fused attention for FP8 with seqlen <= 512 + * \brief Functions for fused attention for FP8 */ #include "transformer_engine/fused_attn.h" @@ -34,9 +34,9 @@ void fused_attn_fp8_bwd( NVTE_Softmax_Type softmax_type, size_t window_size_left, size_t window_size_right, bool bottom_right_diagonal, bool deterministic, const Tensor *input_Q, const Tensor *input_K, const Tensor *input_V, const Tensor *input_O, const Tensor *input_dO, - const Tensor *input_dO_f16, const Tensor *input_M, const Tensor *input_ZInv, - const Tensor *input_S, const Tensor *input_SoftmaxOffset, Tensor *input_output_dP, - const Tensor *output_dQ, const Tensor *output_dK, const Tensor *output_dV, - Tensor *output_dSoftmaxOffset, const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, - const Tensor *rng_state, Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); + const Tensor *input_dO_f16, const Tensor *input_M, const Tensor *input_S, + const Tensor *input_SoftmaxOffset, Tensor *input_output_dP, const Tensor *output_dQ, + const Tensor *output_dK, const Tensor *output_dV, Tensor *output_dSoftmaxOffset, + const Tensor *cu_seqlens_q, const Tensor *cu_seqlens_kv, const Tensor *rng_state, + Tensor *workspace, cudaStream_t stream, cudnnHandle_t handle); } // namespace transformer_engine diff --git a/transformer_engine/common/fused_attn/utils.cu b/transformer_engine/common/fused_attn/utils.cu index f37eeb0c6..3e628b658 100644 --- a/transformer_engine/common/fused_attn/utils.cu +++ b/transformer_engine/common/fused_attn/utils.cu @@ -411,20 +411,6 @@ cudnn_frontend::Operation ternary_pw_op_create(cudnn_frontend::Tensor const &xDe return pw_op_created; } -// convert cu_seqlens_q to qkv/o_ragged_offset and actual_seqlens_q -__global__ void cu_seqlens_to_offsets(int64_t b, int64_t h, int64_t d, int32_t *cu_seqlens_q, - int32_t *actual_seqlens_q, int32_t *qkv_ragged_offset, - int32_t *o_ragged_offset) { - size_t tid = blockIdx.x * blockDim.x + threadIdx.x; - if (tid < b) { - actual_seqlens_q[tid] = cu_seqlens_q[tid + 1] - cu_seqlens_q[tid]; - } - if (tid < b + 1) { - qkv_ragged_offset[tid] = cu_seqlens_q[tid] * 3 * h * d; - o_ragged_offset[tid] = cu_seqlens_q[tid] * h * d; - } -} - // convert cu_seqlens to actual_seqlens __global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b, int32_t const *const q_cu_seqlens, diff --git a/transformer_engine/common/fused_attn/utils.h b/transformer_engine/common/fused_attn/utils.h index c3736a6c6..41656062a 100644 --- a/transformer_engine/common/fused_attn/utils.h +++ b/transformer_engine/common/fused_attn/utils.h @@ -333,10 +333,6 @@ struct FADescriptor_v1 { } }; -__global__ void cu_seqlens_to_offsets(int64_t b, int64_t h, int64_t d, int32_t *cu_seqlens_q, - int32_t *actual_seqlens_q, int32_t *qkv_ragged_offset, - int32_t *o_ragged_offset); - __global__ void cu_seqlens_to_actual_seqlens(int64_t actual_b, int64_t max_b, int32_t const *const q_cu_seqlens, int32_t const *const kv_cu_seqlens, int32_t *q_seqlens, diff --git a/transformer_engine/common/include/transformer_engine/fused_attn.h b/transformer_engine/common/include/transformer_engine/fused_attn.h index d301be573..d9d278662 100644 --- a/transformer_engine/common/include/transformer_engine/fused_attn.h +++ b/transformer_engine/common/include/transformer_engine/fused_attn.h @@ -231,15 +231,6 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * - D = Dropout(S) * - O = D * Transpose(V) * - * Support Matrix: - \verbatim - | backend | precision | qkv layout | bias | mask | dropout | sequence length | head_dim | - | 1 | FP16/BF16 | BS3HD,SB3HD,BSH3D,SBH3D | NO/POST_SCALE_BIAS/ALIBI | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK | Yes | any, % 64 == 0 | <= 128, % 8 == 0 | - | | | BSHD_BS2HD,BSHD_BSH2D,SBHD_SB2HD,SBHD_SBH2D | | | | | | - | | | BSHD_BSHD_BSHD,SBHD_SBHD_SBHD | | | | | | - | 2 | FP8 | T3HD | NO_BIAS | PADDING_MASK | Yes | <= 512, % 64 == 0 | 64 | - \endverbatim - * * Notes: * * Tensors `cu_seqlens_q_padded` and `cu_seqlens_kv_padded` @@ -261,7 +252,7 @@ NVTE_Fused_Attn_Backend nvte_get_fused_attn_backend( * \param[in,out] S The S tensor. * \param[out] O The output O tensor. * \param[out] Aux_CTX_Tensors Auxiliary output tensors when training, - * e.g. M, ZInv, rng_state. + * e.g. softmax stats, optional Max, rng_state. * \param[in] cu_seqlens_q Cumulative sequence lengths for Q, [batch_size + 1]. * \param[in] cu_seqlens_kv Cumulative sequence lengths for K and V, [batch_size + 1]. * \param[in] cu_seqlens_q_padded Cumulative sequence offsets for Q, [batch_size + 1]. @@ -308,15 +299,6 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso /*! \brief Compute the backward of the dot product attention with separate Q, K and V. * - * Support Matrix: - \verbatim - | backend | precision | qkv layout | bias | mask | dropout | sequence length | head_dim | - | 1 | FP16/BF16 | BS3HD,SB3HD,BSH3D,SBH3D | NO/POST_SCALE_BIAS/ALIBI | NO/PADDING/CAUSAL/PADDING_CAUSAL_MASK | Yes | any, % 64 == 0 | <= 128, % 8 == 0 | - | | | BSHD_BS2HD,BSHD_BSH2D,SBHD_SB2HD,SBHD_SBH2D | | | | | | - | | | BSHD_BSHD_BSHD,SBHD_SBHD_SBHD | | | | | | - | 2 | FP8 | T3HD | NO_BIAS | PADDING_MASK | Yes | <= 512, % 64 == 0 | 64 | - \endverbatim - * * Notes: * * Tensors `cu_seqlens_q_padded` and `cu_seqlens_kv_padded` @@ -338,7 +320,7 @@ void nvte_fused_attn_fwd(const NVTETensor Q, const NVTETensor K, const NVTETenso * \param[in] S The S tensor. * \param[in,out] dP The gradient of the P tensor. * \param[in] Aux_CTX_Tensors Auxiliary tensors from context when in training mode, - * e.g. M, ZInv, rng_state. + * e.g. softmax stats, optional Max, rng_state. * \param[out] dQ The gradient of the Q tensor. * \param[out] dK The gradient of the K tensor. * \param[out] dV The gradient of the V tensor. diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 7b10593ac..32eb1b597 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -980,10 +980,7 @@ def cp_p2p_fwd_fused_attn( ) if fp8: - if qkv_layout != "t3hd": - softmax_lse_per_step, rng_states = aux_ctx_tensors - else: - softmax_lse_per_step, _, rng_states = aux_ctx_tensors + softmax_lse_per_step, rng_states = aux_ctx_tensors else: softmax_lse_per_step, rng_states, *rest = aux_ctx_tensors attn_bias = rest[0] if len(rest) > 0 else None @@ -1169,17 +1166,7 @@ def cp_p2p_bwd_fused_attn( section, ): """Per-tile backward call of CP P2P with FusedAttention backend""" - if fp8: - if qkv_layout == "t3hd": - aux_tensors = [ - softmax_lse, - softmax_lse, - rng_states[cp_size - step - 1], - ] - else: - aux_tensors = [softmax_lse, rng_states[cp_size - step - 1]] - else: - aux_tensors = [softmax_lse, rng_states[cp_size - step - 1]] + aux_tensors = [softmax_lse, rng_states[cp_size - step - 1]] max_seqlen_q_ = max_seqlen_q max_seqlen_kv_ = max_seqlen_kv @@ -1195,17 +1182,7 @@ def cp_p2p_bwd_fused_attn( attn_mask_type_ = "padding" if "padding" in attn_mask_type else "no_mask" elif section == "upper-triangle": q_part, out_part, dout_part = [x.contiguous() for x in [q_part, out_part, dout_part]] - if fp8: - if qkv_layout == "t3hd": - aux_tensors = [ - softmax_lse_, - softmax_lse_, - rng_states[cp_size - step - 1], - ] - else: - aux_tensors = [softmax_lse_, rng_states[cp_size - step - 1]] - else: - aux_tensors = [softmax_lse_, rng_states[cp_size - step - 1]] + aux_tensors = [softmax_lse_, rng_states[cp_size - step - 1]] max_seqlen_q_ = max_seqlen_q // 2 cu_seqlens_q_padded_ = None if cu_seqlens_q_padded is None else cu_seqlens_q_padded // 2 @@ -3223,10 +3200,7 @@ def forward( **fp8_meta_kwargs, ) if fp8: - if qkv_layout != "t3hd": - softmax_lse_per_step[i], rng_states[i] = aux_ctx_tensors - else: - softmax_lse_per_step[i], _, rng_states[i] = aux_ctx_tensors + softmax_lse_per_step[i], rng_states[i] = aux_ctx_tensors else: softmax_lse_per_step[i], rng_states[i], *_ = aux_ctx_tensors if return_max_logit: @@ -3588,17 +3562,10 @@ def backward(ctx, dout, *_args): out_part = out.select(seq_dim_o, i).contiguous() dout_part = dout.select(seq_dim_o, i).contiguous() if ctx.use_fused_attention: - if ctx.fp8 and ctx.qkv_layout == "t3hd": - aux_ctx_tensors = [ - softmax_lse_per_step[i], - softmax_lse_per_step[i], - rng_states[i], - ] - else: - aux_ctx_tensors = [ - softmax_lse_per_step[i], - rng_states[i], - ] + aux_ctx_tensors = [ + softmax_lse_per_step[i], + rng_states[i], + ] fused_attn_backend = tex.NVTE_Fused_Attn_Backend.NVTE_F16_arbitrary_seqlen fp8_meta_kwargs = {} new_qkv_layout = ctx.qkv_layout diff --git a/transformer_engine/pytorch/cpp_extensions/fused_attn.py b/transformer_engine/pytorch/cpp_extensions/fused_attn.py index d8f301144..2ce939430 100644 --- a/transformer_engine/pytorch/cpp_extensions/fused_attn.py +++ b/transformer_engine/pytorch/cpp_extensions/fused_attn.py @@ -252,12 +252,11 @@ def fused_attn_fwd( softmaxStats: torch.Tensor log(sum(e^(x - max(x)))), where x=Q*K.T shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32 - 2. if fused_attention_backend == FusedAttnBackend["FP8"] - M: torch.Tensor - max(Q*K.T) + Max: torch.Tensor, only when return_max_logit is True shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32 - ZInv: torch.Tensor, only allocated for T3HD path - 1/sum(e^(x - max(x))), where x=Q*K.T + 2. if fused_attention_backend == FusedAttnBackend["FP8"] + softmaxStats: torch.Tensor + log(sum(e^(x - max(x)))), where x=Q*K.T shape [batch_size, num_heads, max_seqlen_q, 1], dtype float32 rng_state: torch.Tensor state of the random number generator; @@ -461,7 +460,7 @@ def fused_attn_bwd( in torch.dtype aux_ctx_tensors : List[torch.Tensor] auxiliary output tensors of the forward pass when its is_training is True, - e.g. aux_ctx_tensors = [M, ZInv, rng_state] + e.g. aux_ctx_tensors = [S, Max, rng_state] fused_attention_backend : tex.NVTE_Fused_Attn_Backend please see FusedAttention module for details on supported backends. cu_seqlens_q_padded : torch.Tensor, default = None diff --git a/transformer_engine/pytorch/csrc/extensions/attention.cpp b/transformer_engine/pytorch/csrc/extensions/attention.cpp index 8d7a24dce..7e8018b3f 100644 --- a/transformer_engine/pytorch/csrc/extensions/attention.cpp +++ b/transformer_engine/pytorch/csrc/extensions/attention.cpp @@ -271,21 +271,17 @@ std::vector fused_attn_fwd( nvte_set_tensor_param(&nvte_aux_tensor_pack.tensors[i], kNVTERowwiseData, &temp_data); }; // allocate memory for nvte_aux_tensor_pack.tensors - // f16_arbitrary: - // return_max_logit=false: S [b, h, sq, 1], rng_state [2], (optional) Bias [1, h, sq, skv], (optional) SoftmaxOffset [1, h, 1, 1] - // return_max_logit=true: S [b, h, sq, 1], Max [b, h, sq, 1], rng_state [2], (optional) Bias [1, h, sq, skv], (optional) SoftmaxOffset [1, h, 1, 1] - // fp8 : M [b, h, sq, 1], optional ZInv [b, h, sq, 1] (T3HD path), rng_state [2] + // f16_arbitrary: S [b, h, sq, 1]/[tq, h, 1], (optional) Max [b, h, sq, 1]/[tq, h, 1], rng_state [2], (optional) Bias [1, h, sq, skv], (optional) SoftmaxOffset [1, h, 1, 1] + // fp8 : S [b, h, sq, 1], rng_state [2] size_t i = 0; at::Tensor output_tensor; - // intermediate softmax tensor, S or M (for fp8) + // intermediate softmax stats tensor S output_tensor = allocateSpace(nvte_shape_to_vector(nvte_tensor_shape(nvte_aux_tensor_pack.tensors[i])), static_cast(nvte_tensor_type(nvte_aux_tensor_pack.tensors[i])), false); set_tensor_param(i++, output_tensor); - // fp8 T3HD has an additional softmax stats tensor, ZInv; return_max_logit=true has an additional Max tensor - if (((qkv_type == DType::kFloat8E4M3 || qkv_type == DType::kFloat8E5M2) && - qkv_layout == NVTE_QKV_Layout::NVTE_T3HD) || - return_max_logit) { + // return_max_logit=true allocates Max after S + if (return_max_logit) { output_tensor = allocateSpace(nvte_shape_to_vector(nvte_tensor_shape(nvte_aux_tensor_pack.tensors[i])), static_cast(nvte_tensor_type(nvte_aux_tensor_pack.tensors[i])), false); From b9df40102752521bafb37c9f4c6bc564d44adf0e Mon Sep 17 00:00:00 2001 From: Alp Dener Date: Thu, 7 May 2026 19:11:47 -0500 Subject: [PATCH 041/180] [Common] Improved fused MoE aux loss kernel for large # of experts (#2758) * added new implementation of fused_moe_aux_loss_forward kernel Signed-off-by: Alp Dener * Fix race condition, type-punning, and namespace bugs in fused_moe_aux_loss_v2 kernel - Accumulate into a float buffer instead of atomicAdd-ing directly into aux_loss (which could be fp16/bf16), fixing a buffer overflow and wrong results for non-float dtypes - Zero the accumulator on the host before launch to eliminate the race between block 0's init and other blocks' atomicAdds - Move kernel into fused_router namespace so symbols resolve correctly - Round block size up to a warp multiple for well-defined shuffles - Allocate Const_buf with 2 elements to hold both C_coeff and the float accumulator Signed-off-by: Alp Dener * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * added shared memory check on number of experts Signed-off-by: Alp Dener * removed duplicate syncwarp Signed-off-by: Alp Dener * updated TE/JAX primitive for fused MoE aux loss to comply with the new V2 API in TE/common Signed-off-by: Alp Dener * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * added missing syncthreads after atomicAdds Signed-off-by: Alp Dener * restored the small 1grid/1block kernel for casting accumulated float result to DataType Signed-off-by: Alp Dener * fixed inter-block race on accumulation coefficient Signed-off-by: Alp Dener * fixed the intermediate coefficient buffer getting passed onto the backward pass correctly Signed-off-by: Alp Dener * removed old kernel, removed _v2 name from new kernel Signed-off-by: Alp Dener * removed unused num_experts from kernel Signed-off-by: Alp Dener * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Alp Dener Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../common/fused_router/fused_moe_aux_loss.cu | 228 ++++++------------ .../jax/cpp_extensions/router.py | 2 +- .../jax/csrc/extensions/router.cpp | 12 +- .../pytorch/csrc/extensions/router.cpp | 2 +- 4 files changed, 79 insertions(+), 165 deletions(-) diff --git a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu index 8aff85450..7e516af97 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -5,7 +5,6 @@ ************************************************************************/ #include -#include #include #include @@ -21,187 +20,102 @@ namespace fused_router { template __global__ void fused_moe_aux_loss_forward_kernel(const DataType* probs, const IndexType* tokens_per_expert, - int total_num_tokens, int num_experts, - int num_rows, int num_cols, int topk, float coeff, - DataType* aux_loss, float* Const_buf) { -#if __CUDA_ARCH__ >= 900 - // Using cooperative_groups to manage the cluster - namespace cg = cooperative_groups; - cg::cluster_group cluster = cg::this_cluster(); - int thread_id = cg::this_grid().thread_rank(); - int lane_id = thread_id % kThreadsPerWarp; - int warp_id = thread_id / kThreadsPerWarp; - int warp_num = blockDim.x * gridDim.x / kThreadsPerWarp; - // Only 1 block in the cluster - int block_id = cluster.block_rank(); - int block_num = cluster.dim_blocks().x; - int cluster_id = blockIdx.x / block_num; - if (cluster_id > 0) return; // Only use the cluster 0 - - extern __shared__ float shmem_aux_loss[]; - CompType* aggregated_probs_per_expert = reinterpret_cast(shmem_aux_loss); - // Clear the shmem - for (int i = threadIdx.x; i < num_cols; i += blockDim.x) { - aggregated_probs_per_expert[i] = CompType(0); - } - __syncthreads(); - - /** - * Section: Reduce the probs to the aggregated_probs_per_expert - * 1. reduce on the block - * 2. reduce on the cluster - */ - // Loop: for all positions in each row - for (int i = lane_id; i < num_cols; i += kThreadsPerWarp) { - CompType tmp = CompType(0); - // Loop: for all rows that this warp is responsible for - for (int j = warp_id; j < num_rows; j += warp_num) { - tmp += CompType(probs[j * num_cols + i]); - } - atomicAdd(&aggregated_probs_per_expert[i], tmp); + int total_num_tokens, int num_rows, int num_cols, + int topk, float coeff, float* Coeff_buf) { + // ----------------------------------------------------------------------- + // 1) Write the CPU-computed coefficient into a device buffer to re-use in BWD + // ----------------------------------------------------------------------- + if (threadIdx.x == 0 && blockIdx.x == 0) { + Coeff_buf[0] = coeff; } - cluster.sync(); - // The block 0 will reduce the results of all blocks - if (block_id == 0) { - for (int i = 1; i < block_num; i++) { - // Map the shared memory of the block i to the current block - CompType* dst_smem = reinterpret_cast(cluster.map_shared_rank(shmem_aux_loss, i)); - for (int j = threadIdx.x; j < num_cols; j += blockDim.x) { - atomicAdd(&aggregated_probs_per_expert[j], dst_smem[j]); - } - } - } - cluster.sync(); - /** - * Section: aggregated_probs_per_expert * tokens_per_expert - * In-place update on shmem - */ - if (block_id == 0) { - for (int i = threadIdx.x; i < num_cols; i += blockDim.x) { - aggregated_probs_per_expert[i] *= CompType(tokens_per_expert[i]); - } - __syncthreads(); + // ----------------------------------------------------------------------- + // 2) Each CTA computes a partial dot-product: + // Sigma_col ( Sigma_row probs[row, col] ) * tokens_per_expert[col] + // ----------------------------------------------------------------------- + CompType thread_sum = CompType(0); - if (warp_id == 0) { - /** - * Section: Reduce to get the sum of aggregated_probs_per_expert - */ - CompType intermediate_result = - warp_reduce_on_shmem(aggregated_probs_per_expert, num_cols, ReduceFuncType::SUM, lane_id); - __syncwarp(); + // Grid-stride over rows so that every row is processed exactly once. + // Each thread processes a subset of columns. + for (int col = threadIdx.x; col < num_cols; col += blockDim.x) { + CompType col_sum = CompType(0); - if (lane_id == 0) { - /** - * Section: Compute the aux_loss - */ - float C_coeff = (num_experts * coeff) / topk / total_num_tokens / total_num_tokens; - aux_loss[0] = static_cast(intermediate_result * C_coeff); - Const_buf[0] = C_coeff; - } + // Accumulate probs over the rows assigned to this CTA (grid-stride). + for (int row = blockIdx.x; row < num_rows; row += gridDim.x) { + col_sum += CompType(probs[row * num_cols + col]); } - } -#else - // Use Only 1 block/1024 threads to avoid the grid sync - if (blockIdx.x > 0) return; - int warp_num = blockDim.x / kThreadsPerWarp; - int warp_id = threadIdx.x / kThreadsPerWarp; - int lane_id = threadIdx.x % kThreadsPerWarp; - extern __shared__ float shmem_aux_loss[]; - CompType* aggregated_probs_per_expert = reinterpret_cast(shmem_aux_loss); - // Clear the shmem - for (int i = threadIdx.x; i < num_cols; i += blockDim.x) { - aggregated_probs_per_expert[i] = CompType(0); - } - __syncthreads(); + // Multiply by the token count for this expert. + col_sum *= CompType(tokens_per_expert[col]); - /** - * Section: Reduce the probs to the aggregated_probs_per_expert - */ - // Loop: for all positions in each row - for (int i = lane_id; i < num_cols; i += kThreadsPerWarp) { - CompType tmp = CompType(0); - // Loop: for all rows that this warp is responsible for - for (int j = warp_id; j < num_rows; j += warp_num) { - tmp += CompType(probs[j * num_cols + i]); - } - atomicAdd(&aggregated_probs_per_expert[i], tmp); + // Accumulate the per-column contribution into the thread-local sum. + thread_sum += col_sum; } - __syncthreads(); - /** - * Section: aggregated_probs_per_expert * tokens_per_expert - * In-place update on shmem - */ - for (int i = threadIdx.x; i < num_cols; i += blockDim.x) { - aggregated_probs_per_expert[i] *= CompType(tokens_per_expert[i]); - } + // ----------------------------------------------------------------------- + // 3) Block-level reduction of thread_sum using warp_reduce_on_shmem + // ----------------------------------------------------------------------- + extern __shared__ float shmem[]; + CompType* shmem_block = reinterpret_cast(shmem); + shmem_block[threadIdx.x] = thread_sum; __syncthreads(); + const int warp_id = threadIdx.x / kThreadsPerWarp; + const int lane_id = threadIdx.x % kThreadsPerWarp; if (warp_id == 0) { - /** - * Section: Reduce to get the sum of aggregated_probs_per_expert - */ - CompType intermediate_result = - warp_reduce_on_shmem(aggregated_probs_per_expert, num_cols, ReduceFuncType::SUM, lane_id); - __syncwarp(); - + CompType block_sum = warp_reduce_on_shmem(shmem_block, static_cast(blockDim.x), + ReduceFuncType::SUM, lane_id); if (lane_id == 0) { - /** - * Section: Compute the aux_loss - */ - float C_coeff = (num_experts * coeff) / topk / total_num_tokens / total_num_tokens; - aux_loss[0] = static_cast(intermediate_result * C_coeff); - Const_buf[0] = C_coeff; + atomicAdd(&Coeff_buf[1], static_cast(block_sum * coeff)); } } -#endif } +// Small kernel to convert the float accumulator to the output DataType. +template +__global__ void convert_accum_to_output(const float* Coeff_buf, DataType* aux_loss) { + aux_loss[0] = static_cast(Coeff_buf[1]); +} + +/* ------------------------------------------------------------------------- + * Kernel launcher -- simplified (no cluster launch). + * ------------------------------------------------------------------------- */ template void fused_moe_aux_loss_forward_kernel_launcher(const DataType* probs, const IndexType* tokens_per_expert, int total_num_tokens, int num_experts, int num_rows, int num_cols, int topk, float coeff, - DataType* aux_loss, float* Const_buf, + DataType* aux_loss, float* Coeff_buf, cudaStream_t stream) { - if (cuda::sm_arch(cuda::current_device()) >= 90) { - cudaLaunchConfig_t config = {0}; - int cluster_size = 8; - config.gridDim = cluster_size; - config.blockDim = 1024; - config.dynamicSmemBytes = sizeof(CompType) * num_cols; - config.stream = stream; - - // Update the max cluster size based on the device - NVTE_CHECK_CUDA(cudaOccupancyMaxPotentialClusterSize( - &cluster_size, - reinterpret_cast(fused_moe_aux_loss_forward_kernel), &config)); - - cudaLaunchAttribute attribute[1]; - attribute[0].id = cudaLaunchAttributeClusterDimension; - attribute[0].val.clusterDim.x = cluster_size; - attribute[0].val.clusterDim.y = 1; - attribute[0].val.clusterDim.z = 1; - config.numAttrs = 1; - config.attrs = attribute; + NVTE_CHECK(num_experts == num_cols, "Number of experts (", num_experts, + ") must be equal to number of input columns (", num_cols, ")."); + + // Round up to a multiple of warp size for correct warp shuffles. + const int block_size = ((std::min(1024, num_cols) + static_cast(kThreadsPerWarp) - 1) / + static_cast(kThreadsPerWarp)) * + static_cast(kThreadsPerWarp); + const int grid_size = cuda::sm_count() * 2; + + // One CompType per thread in shared memory. + const size_t smem_size = block_size * sizeof(CompType); + check_shared_memory_capacity_num_experts(smem_size, num_experts); + + // Compute final coefficient and zero the float accumulator (Coeff_buf[1]) before launch. + const float C_coeff = (num_experts * coeff) / topk / total_num_tokens / total_num_tokens; + NVTE_CHECK_CUDA(cudaMemsetAsync(Coeff_buf + 1, 0, sizeof(float), stream)); + fused_moe_aux_loss_forward_kernel + <<>>(probs, tokens_per_expert, total_num_tokens, + num_rows, num_cols, topk, C_coeff, Coeff_buf); + NVTE_CHECK_CUDA(cudaGetLastError()); - NVTE_CHECK_CUDA(cudaLaunchKernelEx( - &config, fused_moe_aux_loss_forward_kernel, probs, tokens_per_expert, - total_num_tokens, num_experts, num_rows, num_cols, topk, coeff, aux_loss, Const_buf)); - } else { - size_t smem_size = sizeof(CompType) * num_cols; - fused_moe_aux_loss_forward_kernel - <<<1, 1024, smem_size, stream>>>(probs, tokens_per_expert, total_num_tokens, num_experts, - num_rows, num_cols, topk, coeff, aux_loss, Const_buf); - NVTE_CHECK_CUDA(cudaGetLastError()); - } + // Convert the float accumulator to the output DataType. + convert_accum_to_output<<<1, 1, 0, stream>>>(Coeff_buf, aux_loss); + NVTE_CHECK_CUDA(cudaGetLastError()); } void fused_moe_aux_loss_forward(const Tensor& probs, const Tensor& tokens_per_expert, int total_num_tokens, int num_experts, int num_rows, int num_cols, - int topk, float coeff, Tensor& aux_loss, Tensor& Const_buf, + int topk, float coeff, Tensor& aux_loss, Tensor& Coeff_buf, cudaStream_t stream) { TE_ROUTER_PROBS_TYPE_SWITCH_ALL( probs.data.dtype, DataType, @@ -212,7 +126,7 @@ void fused_moe_aux_loss_forward(const Tensor& probs, const Tensor& tokens_per_ex reinterpret_cast(tokens_per_expert.data.dptr), total_num_tokens, num_experts, num_rows, num_cols, topk, coeff, reinterpret_cast(aux_loss.data.dptr), - reinterpret_cast(Const_buf.data.dptr), stream););); + reinterpret_cast(Coeff_buf.data.dptr), stream););); } template @@ -269,13 +183,13 @@ void fused_moe_aux_loss_backward(const Tensor& Const_buf, const Tensor& tokens_p void nvte_fused_moe_aux_loss_forward(const NVTETensor probs, const NVTETensor tokens_per_expert, int total_num_tokens, int num_experts, int num_rows, int num_cols, int topk, float coeff, NVTETensor aux_loss, - NVTETensor Const_buf, cudaStream_t stream) { + NVTETensor Coeff_buf, cudaStream_t stream) { NVTE_API_CALL(nvte_fused_moe_aux_loss_forward); using namespace transformer_engine; fused_router::fused_moe_aux_loss_forward( *convertNVTETensorCheck(probs), *convertNVTETensorCheck(tokens_per_expert), total_num_tokens, num_experts, num_rows, num_cols, topk, coeff, *convertNVTETensorCheck(aux_loss), - *convertNVTETensorCheck(Const_buf), stream); + *convertNVTETensorCheck(Coeff_buf), stream); } void nvte_fused_moe_aux_loss_backward(const NVTETensor Const_buf, diff --git a/transformer_engine/jax/cpp_extensions/router.py b/transformer_engine/jax/cpp_extensions/router.py index f2affacda..0ae267cbf 100644 --- a/transformer_engine/jax/cpp_extensions/router.py +++ b/transformer_engine/jax/cpp_extensions/router.py @@ -401,7 +401,7 @@ def abstract(probs_aval, tokens_per_expert_aval, topk, coeff): del topk, coeff, tokens_per_expert_aval i_dtype = dtypes.canonicalize_dtype(probs_aval.dtype) aux_loss_aval = probs_aval.update(shape=(), dtype=i_dtype) - const_buf_aval = probs_aval.update(shape=(1,), dtype=jnp.float32) + const_buf_aval = probs_aval.update(shape=(2,), dtype=jnp.float32) return aux_loss_aval, const_buf_aval @staticmethod diff --git a/transformer_engine/jax/csrc/extensions/router.cpp b/transformer_engine/jax/csrc/extensions/router.cpp index c81671f10..79daec3f0 100644 --- a/transformer_engine/jax/csrc/extensions/router.cpp +++ b/transformer_engine/jax/csrc/extensions/router.cpp @@ -177,12 +177,12 @@ Error_Type FusedMoEAuxLossForwardFFI(cudaStream_t stream, std::vector{static_cast(num_tokens), static_cast(num_experts)}; auto tpe_dtype = convert_ffi_datatype_to_te_dtype(tokens_per_expert_buf.element_type()); auto tpe_shape = std::vector{static_cast(num_experts)}; - auto scalar_shape = std::vector{1}; auto probs_tensor = TensorWrapper(probs_buf.untyped_data(), probs_shape, dtype); auto tpe_tensor = TensorWrapper(tokens_per_expert_buf.untyped_data(), tpe_shape, tpe_dtype); - auto aux_loss_tensor = TensorWrapper(aux_loss_buf->untyped_data(), scalar_shape, dtype); - auto const_buf_tensor = TensorWrapper(const_buf->untyped_data(), scalar_shape, DType::kFloat32); + auto aux_loss_tensor = TensorWrapper(aux_loss_buf->untyped_data(), std::vector{1}, dtype); + auto const_buf_tensor = + TensorWrapper(const_buf->untyped_data(), std::vector{2}, DType::kFloat32); nvte_fused_moe_aux_loss_forward(probs_tensor.data(), tpe_tensor.data(), num_tokens, num_experts, num_tokens, num_experts, static_cast(topk), @@ -219,16 +219,16 @@ Error_Type FusedMoEAuxLossBackwardFFI(cudaStream_t stream, auto num_tokens = static_cast(grad_probs_dims[0]); auto num_experts = static_cast(grad_probs_dims[1]); - auto scalar_shape = std::vector{1}; auto tpe_dims = tokens_per_expert_buf.dimensions(); auto tpe_shape = std::vector{static_cast(tpe_dims[0])}; auto grad_probs_shape = std::vector{static_cast(num_tokens), static_cast(num_experts)}; - auto const_buf_tensor = TensorWrapper(const_buf_in.untyped_data(), scalar_shape, DType::kFloat32); + auto const_buf_tensor = + TensorWrapper(const_buf_in.untyped_data(), std::vector{2}, DType::kFloat32); auto tpe_tensor = TensorWrapper(tokens_per_expert_buf.untyped_data(), tpe_shape, tpe_dtype); auto grad_aux_loss_tensor = - TensorWrapper(grad_aux_loss_buf.untyped_data(), scalar_shape, grad_dtype); + TensorWrapper(grad_aux_loss_buf.untyped_data(), std::vector{1}, grad_dtype); auto grad_probs_tensor = TensorWrapper(grad_probs_buf->untyped_data(), grad_probs_shape, grad_dtype); diff --git a/transformer_engine/pytorch/csrc/extensions/router.cpp b/transformer_engine/pytorch/csrc/extensions/router.cpp index 94625c0f1..4df64d8e2 100644 --- a/transformer_engine/pytorch/csrc/extensions/router.cpp +++ b/transformer_engine/pytorch/csrc/extensions/router.cpp @@ -148,7 +148,7 @@ std::tuple fused_moe_aux_loss_fwd(at::Tensor probs, // Create the output tensor at::Tensor aux_loss = at::empty({}, at::dtype(probs.scalar_type()).device(at::kCUDA)); - at::Tensor Const_buf = at::empty({}, at::dtype(at::kFloat).device(at::kCUDA)); + at::Tensor Const_buf = at::empty({2}, at::dtype(at::kFloat).device(at::kCUDA)); auto probs_cu = makeTransformerEngineTensor(probs); auto tokens_per_expert_cu = makeTransformerEngineTensor(tokens_per_expert); From c74e5aa37a65eda5c1680562119d466c123ca6ae Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Thu, 7 May 2026 19:54:09 -0700 Subject: [PATCH 042/180] Implement row-scaled NVFP4 fprop recipe (#2931) * Adapt initial implementation and make quantization bitwise exact Signed-off-by: Ziang Li Co-authored-by: Yigong Qin * Add col Signed-off-by: Ziang Li * Add fp32 Signed-off-by: Ziang Li * Clean up tests Signed-off-by: Ziang Li * Clean up ref Signed-off-by: Ziang Li * Clean up gemm wrapper Signed-off-by: Ziang Li * Clean up test Signed-off-by: Ziang Li * Clean up Signed-off-by: Ziang Li * Rename and reformat Signed-off-by: Ziang Li * Avoid partial amax folding in gemm Signed-off-by: Ziang Li * Expand test coverage Signed-off-by: Ziang Li * Expand more tests Signed-off-by: Ziang Li * Turn on test for grouped linear sanity Signed-off-by: Ziang Li * Rename pertoken to per_token Signed-off-by: Ziang Li * Expand .cu test Signed-off-by: Ziang Li * Format after rebase Signed-off-by: Ziang Li * Fix test after rebase Signed-off-by: Ziang Li * Clean up cpp test Signed-off-by: Ziang Li * Extend cpp dequantize test Signed-off-by: Ziang Li * Only pass `per_token_activation` to forward activation quantizer and clean up Signed-off-by: Ziang Li * Minor fix test Signed-off-by: Ziang Li * Improve accuracy by unfolding weight per-tensor fp32 Signed-off-by: Ziang Li * Fold row-wise quantization Signed-off-by: Ziang Li * Drop column wise Signed-off-by: Ziang Li * Clean up Signed-off-by: Ziang Li * Clean up Signed-off-by: Ziang Li * Clean up column wise Signed-off-by: Ziang Li * Move shared test helpers Signed-off-by: Ziang Li * Minor clean up test Signed-off-by: Ziang Li * Readability Signed-off-by: Ziang Li * Rename Signed-off-by: Ziang Li * Further refactor Signed-off-by: Ziang Li * Clean up bias Signed-off-by: Ziang Li * Clean up cast Signed-off-by: Ziang Li * Avoid silently disable column wise Signed-off-by: Ziang Li * Clean up Signed-off-by: Ziang Li * `is_quantizable` returns false Signed-off-by: Ziang Li * Error out grouped gemm Signed-off-by: Ziang Li * Tighten test Signed-off-by: Ziang Li * Rename verbose rowwise_amax_is_row_scaled Signed-off-by: Ziang Li * Clean up Signed-off-by: Ziang Li * Explicitly handle both gemm input and error out Signed-off-by: Ziang Li * Minor Signed-off-by: Ziang Li * Nits and lint Signed-off-by: Ziang Li * Minor fix A100 ci Signed-off-by: Ziang Li * Update tests/pytorch/utils.py Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Ziang Li --------- Signed-off-by: Ziang Li Co-authored-by: Yigong Qin Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- docs/envvars.rst | 6 + .../cpp/operator/test_cast_nvfp4_transpose.cu | 171 ++++++--- tests/cpp/operator/test_dequantize_nvfp4.cu | 60 +++- tests/cpp/test_common.cu | 53 +++ tests/cpp/test_common.h | 8 + tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 334 ++++++++++++++++-- .../nvfp4/test_nvfp4_quantize_exact.py | 53 +++ tests/pytorch/test_backward_override.py | 45 ++- tests/pytorch/test_cuda_graphs.py | 20 +- tests/pytorch/test_recipe.py | 41 ++- tests/pytorch/test_sanity.py | 30 +- tests/pytorch/test_torch_compile.py | 4 +- tests/pytorch/utils.py | 26 +- .../common/cast/dispatch/quantize.cuh | 17 +- .../common/cast/nvfp4/dequantize_nvfp4.cuh | 13 +- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 231 ++++++++++-- .../quantize_transpose_nvfp4_tuned_1D.cuh | 68 ++-- .../comm_gemm_overlap/comm_gemm_overlap.cpp | 8 + transformer_engine/common/common.h | 9 +- .../common/gemm/cublaslt_gemm.cu | 3 + .../common/include/transformer_engine/gemm.h | 2 +- .../transformer_engine/transformer_engine.h | 12 + transformer_engine/common/recipe/__init__.py | 6 + .../common/transformer_engine.cpp | 6 + .../common/transpose/cast_transpose.h | 2 +- ...quantize_transpose_vector_blockwise_fp4.cu | 87 +++-- .../pytorch/cpp_extensions/gemm.py | 139 +++++++- transformer_engine/pytorch/csrc/common.h | 2 + .../pytorch/csrc/extensions/activation.cpp | 10 +- .../pytorch/csrc/extensions/bias.cpp | 5 +- .../pytorch/csrc/extensions/cast.cpp | 52 ++- .../pytorch/csrc/extensions/normalization.cpp | 10 +- transformer_engine/pytorch/csrc/quantizer.cpp | 72 +++- .../pytorch/csrc/type_converters.cpp | 2 + .../custom_recipes/quantization_nvfp4.py | 56 ++- transformer_engine/pytorch/quantization.py | 2 + .../pytorch/tensor/grouped_tensor.py | 3 + .../pytorch/tensor/nvfp4_tensor.py | 23 +- .../tensor/storage/grouped_tensor_storage.py | 42 ++- .../tensor/storage/nvfp4_tensor_storage.py | 8 + 40 files changed, 1482 insertions(+), 259 deletions(-) diff --git a/docs/envvars.rst b/docs/envvars.rst index 29ca49814..ffbad409d 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -281,6 +281,12 @@ Kernel Configuration :Default: ``0`` :Description: Emit a warning when falling back from CUTLASS to cuBLAS for grouped GEMM operations. +.. envvar:: NVTE_NVFP4_ROW_SCALED_ACTIVATION + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Enable row-scaled NVFP4 tensors for forward activation quantizers in the ``NVFP4BlockScaling`` recipe. When set to ``1`` (or when ``NVFP4BlockScaling(row_scaled_activation=True)`` is used), rowwise ``amax`` metadata is stored as one FP32 value per tensor row instead of a single scalar. + Torch Compilation and Fusion ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index 15d7c695c..1f37520bc 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -114,16 +114,14 @@ void quantize_nvfp4_1d(float (*OP)(const float), block_amax = std::max(block_amax, std::abs(elt)); } - // 2. Compute E4M3 scaling factor - // Compute per-block encoding/decoding scaling factor - const float S_dec_b = block_amax / 6.0f; - - // Scale & Store per-block decoding scaling factor - const fp8e4m3 S_dec_b_fp8 = static_cast(S_dec_b * S_enc); + // Compute and store the per-block FP8 decode scale + const float S_dec_b = block_amax * (S_enc * (1.0f / 6.0f)); + const fp8e4m3 S_dec_b_fp8 = static_cast(fminf(S_dec_b, Numeric_Traits::maxNorm)); const float S_dec_b_fp32 = static_cast(S_dec_b_fp8); // Compute "correct" per-block encoding scaling factor - const float S_enc_b_fp8 = S_dec_b_fp32 == 0.f ? 0.f : S_enc / S_dec_b_fp32; + const float S_enc_b_fp8 = S_dec_b_fp32 == 0.f ? 0.f : + fminf(1.0f / (S_dec_b_fp32 * (1.0f / S_enc)), Numeric_Traits::maxNorm); const size_t scale_idx = i * scales_stride + block_X; scales[scale_idx] = S_dec_b_fp8; @@ -317,11 +315,31 @@ void compute_ref(float (*OP)(const float), const size_t scales_stride, const size_t scales_stride_t, const bool use_fast_math, - const bool use_2d_quantization = false) + const bool use_2d_quantization = false, + std::vector *rowwise_amax = nullptr) { std::vector input_t = create_transpose(input, rows, cols); - if (use_2d_quantization) { + if (rowwise_amax != nullptr) { + rowwise_amax->resize(rows, 0.0f); + for (size_t row = 0; row < rows; ++row) { + float row_amax = 0.0f; + for (size_t col = 0; col < cols; ++col) { + row_amax = fmaxf(row_amax, fabsf(static_cast(input[row * cols + col]))); + } + (*rowwise_amax)[row] = row_amax; + quantize_nvfp4(OP, + input + row * cols, + output + row * (cols / 2), + scales + row * scales_stride, + 1, + cols, + scales_stride, + row_amax, + use_fast_math, + use_2d_quantization); + } + } else if (use_2d_quantization) { // Step 1: Compute mathematical 8×8 scaling factors std::vector> math_scales; compute_2d_mathematical_scales(OP, input, rows, cols, global_amax, math_scales, use_fast_math); @@ -504,13 +522,12 @@ void print_detailed_tensor_comparison(const std::string& name, void compareResults_nvfp4(const Tensor &test, const void *ref, const void *ref_t, const int rows, const int cols, - double atol = 1e-5, double rtol = 1e-8, bool if_on_gpus = true, bool dump_data = false) { + double atol = 1e-5, double rtol = 1e-8, bool if_on_gpus = true, + bool dump_data = false, bool compare_columnwise = true) { if (if_on_gpus) test.to_cpu(); const fp4e2m1 *test_data = test.rowwise_cpu_dptr(); - const fp4e2m1 *test_data_t = test.columnwise_cpu_dptr(); const fp4e2m1 *ref_data = reinterpret_cast(ref); - const fp4e2m1 *ref_data_t = reinterpret_cast(ref_t); // Print detailed element-by-element comparison // print_detailed_tensor_comparison("output", test_data, ref_data, rows, cols); @@ -519,17 +536,33 @@ void compareResults_nvfp4(const Tensor &test, // Optionally dump tensor data to files for detailed analysis if (dump_data) { dump_nvfp4_tensor_data("output", test_data, ref_data, rows, cols); - dump_nvfp4_tensor_data("output_t", test_data_t, ref_data_t, cols, rows); } compare_nvfp4_tensors("output", test_data, ref_data, rows, cols, atol, rtol); - compare_nvfp4_tensors("output_t", test_data_t, ref_data_t, cols, rows, atol, rtol); + if (compare_columnwise) { + const fp4e2m1 *test_data_t = test.columnwise_cpu_dptr(); + const fp4e2m1 *ref_data_t = reinterpret_cast(ref_t); + if (dump_data) { + dump_nvfp4_tensor_data("output_t", test_data_t, ref_data_t, cols, rows); + } + compare_nvfp4_tensors("output_t", test_data_t, ref_data_t, cols, rows, atol, rtol); + } +} + +void compare_rowwise_amax(const Tensor &output, const std::vector &ref_amax) { + const std::vector test_amax_data = output.tensor_amax_values(); + ASSERT_EQ(test_amax_data.size(), ref_amax.size()); + for (size_t row = 0; row < ref_amax.size(); ++row) { + ASSERT_EQ(test_amax_data[row], ref_amax[row]) + << "Row-scaled amax mismatch at row " << row; + } } template void performTest(float (*OP)(const float), const std::vector& shape, - const bool use_fast_math) { + const bool use_fast_math, + const bool row_scaled_nvfp4 = false) { using namespace test; DType itype = TypeInfo::dtype; @@ -556,7 +589,7 @@ void performTest(float (*OP)(const float), const size_t scales_stride_t = blocks_X_t; Tensor input("input", shape, itype); - Tensor output("output", shape, otype, true, true, NVTE_NVFP4_1D_SCALING); + Tensor output("output", shape, otype, true, !row_scaled_nvfp4, NVTE_NVFP4_1D_SCALING); std::unique_ptr ref_output = std::make_unique(rows * (cols / 2)); std::unique_ptr ref_output_t = std::make_unique(cols * (rows / 2)); @@ -567,26 +600,44 @@ void performTest(float (*OP)(const float), // Golden value of amax chosen to make the 2nd-stage scaling mantissa zero and avoid rounding issues const float amax = 448.0f * 6.0f * 8.0f; - - // Set 2nd stage NVFP4 scaling factor - output.set_tensor_amax(amax); - output.set_tensor_amax_columnwise(amax); - + std::vector ref_rowwise_amax; bool use_2d_quantization = false; + if (row_scaled_nvfp4) { + output.set_tensor_amax_shape({rows}); + output.set_row_scaled_nvfp4(true); + compute_ref(OP, + input.rowwise_cpu_dptr(), + ref_output.get(), + ref_output_t.get(), + ref_scales.get(), + ref_scales_t.get(), + 0.0f, + rows, + cols, + scales_stride, + scales_stride_t, + use_fast_math, + use_2d_quantization, + &ref_rowwise_amax); + } else { + // Set 2nd stage NVFP4 scaling factor + output.set_tensor_amax(amax); + output.set_tensor_amax_columnwise(amax); + compute_ref(OP, + input.rowwise_cpu_dptr(), + ref_output.get(), + ref_output_t.get(), + ref_scales.get(), + ref_scales_t.get(), + amax, + rows, + cols, + scales_stride, + scales_stride_t, + use_fast_math, + use_2d_quantization); + } - compute_ref(OP, - input.rowwise_cpu_dptr(), - ref_output.get(), - ref_output_t.get(), - ref_scales.get(), - ref_scales_t.get(), - amax, - rows, - cols, - scales_stride, - scales_stride_t, - use_fast_math, - use_2d_quantization); // Initialize stochastic rounding Tensor rng_state("rng_state", std::vector{2}, DType::kInt64); rng_state.rowwise_cpu_dptr()[0] = 123; // rng_seed @@ -629,12 +680,8 @@ void performTest(float (*OP)(const float), const double rtol = 1.0E-6; // Set dump_data=true to enable dumping tensor data to files for analysis - compareResults_nvfp4(output, ref_output.get(), ref_output_t.get(), rows, cols, atol, rtol, true, false); - - const fp8e4m3* kernel_scales = output.rowwise_cpu_scale_inv_ptr(); - const fp8e4m3* ref_scales_ptr = ref_scales.get(); - const fp8e4m3* kernel_scales_t = output.columnwise_cpu_scale_inv_ptr(); - const fp8e4m3* ref_scales_t_ptr = ref_scales_t.get(); + compareResults_nvfp4(output, ref_output.get(), ref_output_t.get(), rows, cols, atol, rtol, true, + false, !row_scaled_nvfp4); size_t scale_mismatches_num = 0; compare_scaling_factors("scales", output.rowwise_cpu_scale_inv_ptr(), @@ -642,10 +689,16 @@ void performTest(float (*OP)(const float), unpadded_blocks_Y, unpadded_blocks_X, scales_stride, scale_mismatches_num); - compare_scaling_factors("scales_t", output.columnwise_cpu_scale_inv_ptr(), - ref_scales_t.get(), - unpadded_blocks_Y_t, unpadded_blocks_X_t, scales_stride_t, - scale_mismatches_num); + if (!row_scaled_nvfp4) { + compare_scaling_factors("scales_t", output.columnwise_cpu_scale_inv_ptr(), + ref_scales_t.get(), + unpadded_blocks_Y_t, unpadded_blocks_X_t, scales_stride_t, + scale_mismatches_num); + } + + if (row_scaled_nvfp4) { + compare_rowwise_amax(output, ref_rowwise_amax); + } } std::vector> tensor_dims = { @@ -678,6 +731,7 @@ class FusedCastTransposeNVFP4TestSuite : public ::testing::TestWithParam , transformer_engine::DType, + bool, bool>> {}; TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { @@ -693,6 +747,7 @@ TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { const auto tensor_dims = std::get<1>(GetParam()); const DType input_type = std::get<2>(GetParam()); const bool use_fast_math = std::get<3>(GetParam()); + const bool row_scaled_nvfp4 = std::get<4>(GetParam()); // Skip tests if the input tensor is 1D if (tensor_dims.size() < 2) { @@ -710,7 +765,7 @@ TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { } TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(input_type, InputType, - performTest(OP, tensor_dims, use_fast_math); + performTest(OP, tensor_dims, use_fast_math, row_scaled_nvfp4); ); } @@ -733,6 +788,7 @@ INSTANTIATE_TEST_SUITE_P( ::testing::ValuesIn(Activation_types), ::testing::ValuesIn(tensor_dims), ::testing::Values(DType::kBFloat16), + ::testing::Values(false), ::testing::Values(false)), [](const testing::TestParamInfo& info) { std::string name = to_string(std::get<0>(info.param)); @@ -746,3 +802,28 @@ INSTANTIATE_TEST_SUITE_P( } return name; }); + +INSTANTIATE_TEST_SUITE_P( + OperatorTestRowScaled, + FusedCastTransposeNVFP4TestSuite, + ::testing::Combine( + ::testing::Values(ActivationType::Identity), + ::testing::Values(tensor_dims[4], tensor_dims[9], tensor_dims[12]), + ::testing::Values(DType::kBFloat16, DType::kFloat32), + ::testing::Values(false), + ::testing::Values(true)), + [](const testing::TestParamInfo& info) { + std::string name = to_string(std::get<0>(info.param)); + const auto& shape = std::get<1>(info.param); + for (const auto& s: shape) { + name += "X" + std::to_string(s); + } + name += "X" + test::typeName(std::get<2>(info.param)); + if (std::get<3>(info.param)) { + name += "X_FAST_SCALING"; + } + if (std::get<4>(info.param)) { + name += "XROW_SCALED"; + } + return name; + }); diff --git a/tests/cpp/operator/test_dequantize_nvfp4.cu b/tests/cpp/operator/test_dequantize_nvfp4.cu index 96e85cb5e..ec405b1d9 100644 --- a/tests/cpp/operator/test_dequantize_nvfp4.cu +++ b/tests/cpp/operator/test_dequantize_nvfp4.cu @@ -42,7 +42,7 @@ float2 cvt_fp4x2_to_float2(fp4e2m1x2 fp4_pair) { template void compute_ref_dequantize_nvfp4(const uint8_t *packed_data, const fp8e4m3 *scales, - float amax, + const std::vector &amax, OType *output, size_t rows, size_t cols, @@ -55,7 +55,8 @@ void compute_ref_dequantize_nvfp4(const uint8_t *packed_data, for (size_t row = 0; row < rows; ++row) { for (size_t block = 0; block < Mread; ++block) { const fp8e4m3 scale = scales[row * scale_stride + block]; - const float final_scale = static_cast(scale) * amax * factor_inv; + const float final_scale = + static_cast(scale) * (amax.size() == 1 ? amax[0] : amax[row]) * factor_inv; for (size_t pair_idx = 0; pair_idx < bytes_per_block; ++pair_idx) { const size_t byte_idx = @@ -88,7 +89,8 @@ float compute_amax(const test::Tensor &t, size_t rows, size_t cols) { // Quantize a high-precision input to NVFP4, then dequantize and compare // against a CPU reference computed from the quantized data. template -void performTest_dequantize_nvfp4(const size_t rows, const size_t cols) { +void performTest_dequantize_nvfp4(const size_t rows, const size_t cols, + const bool row_scaled_nvfp4) { using namespace test; DType otype = TypeInfo::dtype; @@ -97,7 +99,10 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols) { Tensor quantized("quantized", std::vector{rows, cols}, DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING); - if (rows > 0 && cols > 0) { + if (row_scaled_nvfp4) { + quantized.set_tensor_amax_shape({rows}); + quantized.set_row_scaled_nvfp4(true); + } else if (rows > 0 && cols > 0) { quantized.set_tensor_amax(compute_amax(input, rows, cols)); } else { quantized.set_tensor_amax(0.0f); @@ -120,7 +125,7 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols) { const uint8_t *fp4_data = reinterpret_cast(quantized.rowwise_cpu_dptr()); const fp8e4m3 *scales = quantized.rowwise_cpu_scale_inv_ptr(); - const float amax_val = quantized.amax(); + const std::vector amax_val = quantized.tensor_amax_values(); const NVTEShape scale_shape = quantized.rowwise_scale_inv_shape(); const size_t scale_stride = scale_shape.data[scale_shape.ndim - 1]; @@ -137,7 +142,8 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols) { // Dequantize NVFP4 with GEMM-swizzled scales and compare against compact path. template -void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols) { +void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols, + const bool row_scaled_nvfp4) { using namespace test; DType otype = TypeInfo::dtype; @@ -146,7 +152,10 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols) Tensor quantized_compact("quantized_compact", std::vector{rows, cols}, DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING); - if (rows > 0 && cols > 0) { + if (row_scaled_nvfp4) { + quantized_compact.set_tensor_amax_shape({rows}); + quantized_compact.set_row_scaled_nvfp4(true); + } else if (rows > 0 && cols > 0) { quantized_compact.set_tensor_amax(compute_amax(input, rows, cols)); } else { quantized_compact.set_tensor_amax(0.0f); @@ -157,7 +166,7 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols) cudaDeviceSynchronize(); } - // Dequantize with compact scales → reference output + // Dequantize with compact scales to get the reference output. Tensor output_compact("output_compact", std::vector{rows, cols}, otype, true, false); nvte_dequantize(quantized_compact.data(), output_compact.data(), 0); cudaDeviceSynchronize(); @@ -165,13 +174,22 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols) // Create tensor with same FP4 data but swizzled scales Tensor quantized_swizzled("quantized_swizzled", std::vector{rows, cols}, DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING); - quantized_swizzled.set_tensor_amax(0.0f); + if (row_scaled_nvfp4) { + quantized_swizzled.set_tensor_amax_shape({rows}); + quantized_swizzled.set_row_scaled_nvfp4(true); + } else { + quantized_swizzled.set_tensor_amax(0.0f); + } quantized_swizzled.set_with_gemm_swizzled_scales(true); // Copy amax and scale from compact to swizzled before FP4 data, // since from_cpu() uploads all CPU buffers (including zero-init data). quantized_compact.to_cpu(); - quantized_swizzled.set_tensor_amax(quantized_compact.amax()); + if (row_scaled_nvfp4) { + quantized_swizzled.copy_tensor_amax_from(quantized_compact); + } else { + quantized_swizzled.set_tensor_amax(quantized_compact.amax()); + } // Copy FP4 data after from_cpu() to avoid being overwritten const size_t data_bytes = rows * cols / 2; @@ -227,7 +245,8 @@ std::vector> nvfp4_tensor_dims = { class DequantizeNVFP4TestSuite : public ::testing::TestWithParam , - transformer_engine::DType>> {}; + transformer_engine::DType, + bool>> {}; TEST_P(DequantizeNVFP4TestSuite, TestDequantizeNVFP4) { @@ -237,10 +256,11 @@ TEST_P(DequantizeNVFP4TestSuite, TestDequantizeNVFP4) const auto tensor_size = std::get<0>(GetParam()); const DType output_type = std::get<1>(GetParam()); + const bool row_scaled_nvfp4 = std::get<2>(GetParam()); TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(output_type, OutputType, performTest_dequantize_nvfp4( - tensor_size.first, tensor_size.second); + tensor_size.first, tensor_size.second, row_scaled_nvfp4); ); } @@ -249,19 +269,22 @@ INSTANTIATE_TEST_SUITE_P( DequantizeNVFP4TestSuite, ::testing::Combine( ::testing::ValuesIn(nvfp4_tensor_dims), - ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16)), + ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16), + ::testing::Bool()), [](const testing::TestParamInfo& info) { std::string name = std::to_string(std::get<0>(info.param).first) + "X" + std::to_string(std::get<0>(info.param).second) + "X" + - test::typeName(std::get<1>(info.param)); + test::typeName(std::get<1>(info.param)) + "X" + + (std::get<2>(info.param) ? "RowScaled" : "PerTensor"); return name; } ); class DequantizeNVFP4SwizzledTestSuite : public ::testing::TestWithParam , - transformer_engine::DType>> {}; + transformer_engine::DType, + bool>> {}; TEST_P(DequantizeNVFP4SwizzledTestSuite, TestDequantizeNVFP4Swizzled) { @@ -271,10 +294,11 @@ TEST_P(DequantizeNVFP4SwizzledTestSuite, TestDequantizeNVFP4Swizzled) const auto tensor_size = std::get<0>(GetParam()); const DType output_type = std::get<1>(GetParam()); + const bool row_scaled_nvfp4 = std::get<2>(GetParam()); TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(output_type, OutputType, performTest_dequantize_nvfp4_swizzled( - tensor_size.first, tensor_size.second); + tensor_size.first, tensor_size.second, row_scaled_nvfp4); ); } @@ -283,12 +307,14 @@ INSTANTIATE_TEST_SUITE_P( DequantizeNVFP4SwizzledTestSuite, ::testing::Combine( ::testing::ValuesIn(nvfp4_tensor_dims), - ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16)), + ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16), + ::testing::Bool()), [](const testing::TestParamInfo& info) { std::string name = std::to_string(std::get<0>(info.param).first) + "X" + std::to_string(std::get<0>(info.param).second) + "X" + test::typeName(std::get<1>(info.param)) + "X" + + (std::get<2>(info.param) ? "RowScaled" : "PerTensor") + "X" + "Swizzled"; return name; } diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index c756b8381..96e71f951 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -543,6 +543,59 @@ void Tensor::set_scale(float scale) { } } +void Tensor::set_tensor_amax_shape(const std::vector &shape) { + const size_t numel = product(shape); + NVTE_CHECK(tensor_.scaling_mode() == NVTE_NVFP4_1D_SCALING, + "Amax shape override is only supported for NVFP4 test tensors."); + + auto old_amax = tensor_.get_amax(); + if (old_amax.data_ptr != nullptr) { + NVTE_CHECK_CUDA(cudaFree(old_amax.data_ptr)); + } + + float *amax = nullptr; + NVTE_CHECK_CUDA(cudaMalloc(&amax, numel * sizeof(float))); + NVTE_CHECK_CUDA(cudaMemset(amax, 0, numel * sizeof(float))); + tensor_.set_amax(amax, DType::kFloat32, shape); +} + +std::vector Tensor::tensor_amax_values() const { + const auto amax = tensor_.get_amax(); + NVTE_CHECK(static_cast(amax.dtype) == DType::kFloat32, "Tensor amax must be FP32."); + + const size_t numel = product(amax.shape); + if (numel == 0) { + return {}; + } + NVTE_CHECK(amax.data_ptr != nullptr, "Tensor amax is not allocated."); + + std::vector values(numel); + NVTE_CHECK_CUDA( + cudaMemcpy(values.data(), amax.data_ptr, numel * sizeof(float), cudaMemcpyDeviceToHost)); + return values; +} + +void Tensor::copy_tensor_amax_from(const Tensor &other) { + const auto other_amax = other.tensor_.get_amax(); + NVTE_CHECK(static_cast(other_amax.dtype) == DType::kFloat32, + "Source tensor amax must be FP32."); + + auto my_amax = tensor_.get_amax(); + NVTE_CHECK(static_cast(my_amax.dtype) == DType::kFloat32, + "Destination tensor amax must be FP32."); + NVTE_CHECK(areShapesEqual(my_amax.shape, other_amax.shape), "Amax shape mismatch."); + + const size_t numel = product(other_amax.shape); + if (numel == 0) { + return; + } + + NVTE_CHECK(other_amax.data_ptr != nullptr, "Source tensor amax is not allocated."); + NVTE_CHECK(my_amax.data_ptr != nullptr, "Destination tensor amax is not allocated."); + NVTE_CHECK_CUDA(cudaMemcpy(my_amax.data_ptr, other_amax.data_ptr, numel * sizeof(float), + cudaMemcpyDeviceToDevice)); +} + void Tensor::set_scale_inv(float scale_inv) { if (isFp8Type(dtype()) || isFp4Type(dtype())) { if (rowwise_) { diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index b8389d583..b2a7da89c 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -319,10 +319,18 @@ class Tensor { tensor_.set_amax(nullptr, DType::kFloat32, tensor_.defaultShape); } + void set_tensor_amax_shape(const std::vector &shape); + std::vector tensor_amax_values() const; + void copy_tensor_amax_from(const Tensor &other); + void set_with_gemm_swizzled_scales(bool with_gemm_swizzled_scales){ tensor_.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); } + void set_row_scaled_nvfp4(bool row_scaled_nvfp4) { + tensor_.set_row_scaled_nvfp4(row_scaled_nvfp4); + } + void to_cpu() const; void from_cpu() const; void set_scale(float scale); diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index 911b7660d..b93933627 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -8,6 +8,7 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.constants import TE_DType from transformer_engine.pytorch import NVFP4Quantizer +from transformer_engine.pytorch.cpp_extensions import general_gemm, general_grouped_gemm from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef from transformer_engine.pytorch.custom_recipes import utils @@ -26,6 +27,7 @@ def check_nvfp4_gemm_versus_reference( *, x_columnwise: bool = False, w_columnwise: bool = False, + row_scaled_nvfp4: bool = False, ): te_dtype = tex.DType.kFloat4E2M1 @@ -51,11 +53,12 @@ def check_nvfp4_gemm_versus_reference( x_quantizer = NVFP4Quantizer( fp4_dtype=te_dtype, rowwise=True, - columnwise=True, + columnwise=not row_scaled_nvfp4, with_amax_reduction=False, amax_reduction_group=None, with_rht=False, with_post_rht_amax=False, + row_scaled_nvfp4=row_scaled_nvfp4, ) w_quantizer = NVFP4Quantizer( fp4_dtype=te_dtype, @@ -112,7 +115,16 @@ def check_nvfp4_gemm_versus_reference( sw_trimmed = sw_trimmed.view(torch.float8_e4m3fn) # Create reference quantizer for reference GEMM - ref_quantizer = NVFP4QuantizerRef( + x_ref_quantizer = NVFP4QuantizerRef( + dtype=utils.Fp4Formats.E2M1, + rowwise=True, + columnwise=not row_scaled_nvfp4, + pow_2_scales=False, + eps=0.0, + quant_tile_shape=(1, 16), + row_scaled_nvfp4=row_scaled_nvfp4, + ) + w_ref_quantizer = NVFP4QuantizerRef( dtype=utils.Fp4Formats.E2M1, rowwise=True, columnwise=True, @@ -124,16 +136,16 @@ def check_nvfp4_gemm_versus_reference( # Create reference quantized tensors needed by reference GEMM # Reference GEMM is only rowwise. if x_columnwise: - x_nvfp4_ref = ref_quantizer.quantize(x.t().contiguous()) + x_nvfp4_ref = x_ref_quantizer.quantize(x.t().contiguous()) else: - x_nvfp4_ref = ref_quantizer.quantize(x) + x_nvfp4_ref = x_ref_quantizer.quantize(x) if w_columnwise: - w_nvfp4_ref = ref_quantizer.quantize(w.t().contiguous()) + w_nvfp4_ref = w_ref_quantizer.quantize(w.t().contiguous()) else: - w_nvfp4_ref = ref_quantizer.quantize(w) + w_nvfp4_ref = w_ref_quantizer.quantize(w) # Reference GEMM using quantizer's qgemm method - y_ref = ref_quantizer.qgemm( + y_ref = x_ref_quantizer.qgemm( qx=qx_data, qw=qw_data, m_params=None, # MMParams not used in reference @@ -166,27 +178,38 @@ def check_nvfp4_gemm_versus_reference( x_nvfp4_native.update_usage(rowwise_usage=False) if w_columnwise: w_nvfp4_native.update_usage(rowwise_usage=False) - # Native cuBLAS GEMM - # return type is out, bias_grad, gelu_input, extra_output - # We are just capturing out. - y_native = tex.generic_gemm( - w_nvfp4_native, - transa, - x_nvfp4_native, - transb, - out.clone() if accumulate else None, - out_quantizer, - TE_DType[out_dtype], - bias, - bias_dtype, - use_gelu, - gelu_input, - use_grad, - workspace, - workspace.shape[0], - accumulate, - use_split_accumulator, - )[0] + if row_scaled_nvfp4: + layout = ("T" if transa else "N") + ("T" if transb else "N") + y_native = general_gemm( + w_nvfp4_native, + x_nvfp4_native, + out_dtype=out_dtype, + accumulate=accumulate, + layout=layout, + out=out.clone() if accumulate else None, + )[0] + else: + # Native cuBLAS GEMM + # return type is out, bias_grad, gelu_input, extra_output + # We are just capturing out. + y_native = tex.generic_gemm( + w_nvfp4_native, + transa, + x_nvfp4_native, + transb, + out.clone() if accumulate else None, + out_quantizer, + TE_DType[out_dtype], + bias, + bias_dtype, + use_gelu, + gelu_input, + use_grad, + workspace, + workspace.shape[0], + accumulate, + use_split_accumulator, + )[0] # just in case of accumulation, make sure y_ref and y_native are not the same tensor assert y_ref is not y_native, "y_ref and y_native should not be the same tensor" @@ -199,6 +222,170 @@ def check_nvfp4_gemm_versus_reference( torch.testing.assert_close(y_native, y_ref, atol=8e-3, rtol=8e-3) +def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( + x_dtype: torch.dtype, + w_dtype: torch.dtype, + out_dtype: torch.dtype, + m_splits: list[int], + k: int, + n: int, + *, + use_bias: bool, + single_output: bool, +): + te_dtype = tex.DType.kFloat4E2M1 + device = "cuda" + torch.manual_seed(23) + torch.cuda.manual_seed(23) + + num_gemms = len(m_splits) + + x_quantizer = NVFP4Quantizer( + fp4_dtype=te_dtype, + rowwise=True, + columnwise=False, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=True, + ) + w_quantizer = NVFP4Quantizer( + fp4_dtype=te_dtype, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + ) + + x_nvfp4 = [] + w_nvfp4 = [] + bias = [] + expected = [] + for m in m_splits: + x = torch.randn((m, k), dtype=x_dtype, device=device) + w = torch.randn((n, k), dtype=w_dtype, device=device) + x_nvfp4.append( + x_quantizer.update_quantized( + x, x_quantizer.make_empty(x.shape, dtype=x_dtype, device=device) + ) + ) + w_nvfp4.append( + w_quantizer.update_quantized( + w, w_quantizer.make_empty(w.shape, dtype=w_dtype, device=device) + ) + ) + bias.append(torch.randn(n, dtype=torch.bfloat16, device=device) if use_bias else None) + expected.append( + general_gemm( + w_nvfp4[-1], + x_nvfp4[-1], + out_dtype=out_dtype, + layout="TN", + bias=bias[-1], + )[0] + ) + + if single_output: + out = [torch.empty((sum(m_splits), n), dtype=out_dtype, device=device)] + else: + out = [torch.empty((m, n), dtype=out_dtype, device=device) for m in m_splits] + + grouped_out, _, _ = general_grouped_gemm( + w_nvfp4, + x_nvfp4, + out, + quantization_params=[None] * num_gemms, + out_dtype=out_dtype, + layout="TN", + m_splits=m_splits, + bias=bias, + use_bias=use_bias, + single_output=single_output, + ) + + if single_output: + grouped_slices = torch.split(grouped_out, m_splits, dim=0) + else: + grouped_slices = grouped_out + for grouped, ref in zip(grouped_slices, expected): + torch.testing.assert_close(grouped, ref, atol=0.0, rtol=0.0) + + +def check_nvfp4_row_scaled_gemm_matches_emulated( + x_dtype: torch.dtype, + w_dtype: torch.dtype, + out_dtype: torch.dtype, + M: int, + K: int, + N: int, +): + te_dtype = tex.DType.kFloat4E2M1 + device = "cuda" + torch.manual_seed(37) + torch.cuda.manual_seed(37) + + x = torch.randn((M, K), dtype=x_dtype, device=device) + w = torch.randn((N, K), dtype=w_dtype, device=device) + + x_row_scaled_quantizer = NVFP4Quantizer( + fp4_dtype=te_dtype, + rowwise=True, + columnwise=False, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + row_scaled_nvfp4=True, + ) + x_tensorwise_quantizer = NVFP4Quantizer( + fp4_dtype=te_dtype, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + ) + w_quantizer = NVFP4Quantizer( + fp4_dtype=te_dtype, + rowwise=True, + columnwise=True, + with_amax_reduction=False, + amax_reduction_group=None, + with_rht=False, + with_post_rht_amax=False, + ) + + x_row_scaled = x_row_scaled_quantizer.update_quantized( + x, x_row_scaled_quantizer.make_empty(x.shape, dtype=x_dtype, device=device) + ) + w_nvfp4 = w_quantizer.update_quantized( + w, w_quantizer.make_empty(w.shape, dtype=w_dtype, device=device) + ) + y_row_scaled = general_gemm(w_nvfp4, x_row_scaled, out_dtype=out_dtype, layout="TN")[0] + + emulated_rows = [] + for i in range(M): + x_padded = torch.zeros((16, K), dtype=x_dtype, device=device) + x_padded[0].copy_(x[i]) + x_tensorwise = x_tensorwise_quantizer.update_quantized( + x_padded, + x_tensorwise_quantizer.make_empty(x_padded.shape, dtype=x_dtype, device=device), + ) + emulated_rows.append( + general_gemm(w_nvfp4, x_tensorwise, out_dtype=out_dtype, layout="TN")[0][:1] + ) + + y_emulated = torch.cat(emulated_rows, dim=0) + if out_dtype == torch.bfloat16: + torch.testing.assert_close(y_row_scaled, y_emulated, atol=0.0, rtol=7.8e-3) + else: + torch.testing.assert_close(y_row_scaled, y_emulated, atol=3.0517578125e-5, rtol=0.0) + + @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) @pytest.mark.parametrize( "M, K, N", @@ -229,6 +416,7 @@ def check_nvfp4_gemm_versus_reference( ], ids=["rowxrow", "colxrow", "colxcol"], ) +@pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"]) def test_nvfp4_gemm_versus_reference( M: int, K: int, @@ -239,7 +427,14 @@ def test_nvfp4_gemm_versus_reference( accumulate: bool, is_x_columnwise: bool, is_w_columnwise: bool, + row_scaled_nvfp4: bool, ): + if row_scaled_nvfp4: + if accumulate: + pytest.skip("Row-scaled NVFP4 GEMM output rescale does not support accumulation") + if is_x_columnwise: + pytest.skip("Row-scaled NVFP4 GEMM output rescale requires rowwise RHS usage") + check_nvfp4_gemm_versus_reference( x_dtype=x_dtype, w_dtype=w_dtype, @@ -250,4 +445,87 @@ def test_nvfp4_gemm_versus_reference( accumulate=accumulate, x_columnwise=is_x_columnwise, w_columnwise=is_w_columnwise, + row_scaled_nvfp4=row_scaled_nvfp4, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "m_splits, k, n", + [ + ([32, 48, 48], 128, 128), + ([64, 80, 112], 128, 256), + ([64, 80, 112], 256, 256), + ([64, 80, 112], 1024, 256), + ([256, 256, 512], 1024, 1024), + ([1024, 1536, 1536], 512, 3072), + ([16, 32, 64], 128, 96), + ([80, 96, 128], 640, 304), + ([320, 336, 352], 3072, 992), + ([64, 80, 112], 64, 256), + ([32, 48, 48], 128, 112), + ], +) +@pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) +@pytest.mark.parametrize("w_dtype", [torch.float32, torch.bfloat16], ids=str) +@pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16], ids=str) +@pytest.mark.parametrize("use_bias", [False, True], ids=["no_bias", "bias"]) +@pytest.mark.parametrize("single_output", [False, True], ids=["list_output", "single_output"]) +def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( + m_splits: list[int], + k: int, + n: int, + x_dtype: torch.dtype, + w_dtype: torch.dtype, + out_dtype: torch.dtype, + use_bias: bool, + single_output: bool, +): + check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( + x_dtype=x_dtype, + w_dtype=w_dtype, + out_dtype=out_dtype, + m_splits=m_splits, + k=k, + n=n, + use_bias=use_bias, + single_output=single_output, + ) + + +@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe) +@pytest.mark.parametrize( + "M, K, N", + [ + (128, 128, 128), + (256, 128, 256), + (256, 256, 256), + (256, 1024, 256), + (1024, 1024, 1024), + (4096, 512, 3072), + (112, 128, 96), + (304, 640, 304), + (1008, 3072, 992), + (256, 64, 256), + (128, 128, 112), + ], +) +@pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) +@pytest.mark.parametrize("w_dtype", [torch.float32, torch.bfloat16], ids=str) +@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str) +def test_nvfp4_row_scaled_gemm_matches_emulated( + M: int, + K: int, + N: int, + x_dtype: torch.dtype, + w_dtype: torch.dtype, + out_dtype: torch.dtype, +): + check_nvfp4_row_scaled_gemm_matches_emulated( + x_dtype=x_dtype, + w_dtype=w_dtype, + out_dtype=out_dtype, + M=M, + K=K, + N=N, ) diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index bf3f545b8..0824a5e7b 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -16,6 +16,19 @@ recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True) +def maybe_skip_row_scaled_unsupported_quantization( + row_scaled_nvfp4: bool, + return_transpose: bool, + with_2d_quantization: bool = False, +) -> None: + if not row_scaled_nvfp4: + return + if return_transpose: + pytest.skip("Row-scaled NVFP4 does not support columnwise usage") + if with_2d_quantization: + pytest.skip("Row-scaled NVFP4 does not support 2D quantization") + + def unpack_fp4(x: torch.Tensor) -> torch.Tensor: repeated = x.repeat_interleave(2, dim=1) repeated[:, 0::2] &= 0x0F @@ -31,7 +44,12 @@ def check_quantization_nvfp4_versus_reference( swizzled_scale: bool, use_cpp_allocator: bool, with_2d_quantization: bool, + row_scaled_nvfp4: bool = False, ) -> None: + maybe_skip_row_scaled_unsupported_quantization( + row_scaled_nvfp4, return_transpose, with_2d_quantization + ) + te_dtype = tex.DType.kFloat4E2M1 # Setup device and random seed @@ -52,6 +70,7 @@ def check_quantization_nvfp4_versus_reference( with_rht=False, with_post_rht_amax=False, with_2d_quantization=with_2d_quantization, + row_scaled_nvfp4=row_scaled_nvfp4, ) if use_cpp_allocator: x_nvfp4_sut = nvfp4_quantizer(x) @@ -73,6 +92,7 @@ def check_quantization_nvfp4_versus_reference( ) sx_t = x_nvfp4_sut._columnwise_scale_inv qx_amax = x_nvfp4_sut._amax_rowwise + qx_amax_t = x_nvfp4_sut._amax_columnwise # Reference quantization quant_tile_shape = (1, 16) if not with_2d_quantization else (16, 16) @@ -83,6 +103,7 @@ def check_quantization_nvfp4_versus_reference( pow_2_scales=False, eps=0.0, quant_tile_shape=quant_tile_shape, + row_scaled_nvfp4=row_scaled_nvfp4, ) x_nvfp4_ref = ref_quantizer.quantize(x) @@ -102,6 +123,7 @@ def check_quantization_nvfp4_versus_reference( x_nvfp4_ref.scale_t.view(dtype=torch.uint8) if x_nvfp4_ref.scale_t is not None else None ) ref_amax = x_nvfp4_ref.global_amax_row + ref_amax_t = x_nvfp4_ref.global_amax_col qx = unpack_fp4(qx) qx_t = unpack_fp4(qx_t) if qx_t is not None else None @@ -121,6 +143,7 @@ def check_quantization_nvfp4_versus_reference( ref_sx_t_shape = sx_t_ref.shape sx_t_valid = sx_t[: ref_sx_t_shape[0], : ref_sx_t_shape[1]] torch.testing.assert_close(sx_t_valid, sx_t_ref, atol=0.0, rtol=0.0) + torch.testing.assert_close(qx_amax_t, ref_amax_t, atol=0.0, rtol=0.0) torch.testing.assert_close(qx_amax, ref_amax, atol=0.0, rtol=0.0) @@ -155,6 +178,7 @@ def check_quantization_nvfp4_versus_reference( @pytest.mark.parametrize( "with_2d_quantization", [True, False], ids=["2d_quantization", "1d_quantization"] ) +@pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"]) def test_quantization_block_tiling_versus_reference( x_dtype: torch.dtype, M: int, @@ -163,6 +187,7 @@ def test_quantization_block_tiling_versus_reference( swizzled_scale: bool, use_cpp_allocator: bool, with_2d_quantization: bool, + row_scaled_nvfp4: bool, ) -> None: check_quantization_nvfp4_versus_reference( x_dtype=x_dtype, @@ -172,6 +197,7 @@ def test_quantization_block_tiling_versus_reference( swizzled_scale=swizzled_scale, use_cpp_allocator=use_cpp_allocator, with_2d_quantization=with_2d_quantization, + row_scaled_nvfp4=row_scaled_nvfp4, ) @@ -188,6 +214,7 @@ def test_quantization_block_tiling_versus_reference( @pytest.mark.parametrize( "use_cpp_allocator", [True, False], ids=["cpp_allocator", "python_allocator"] ) +@pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"]) def test_nvfp4_quantization_extrema_versus_reference( x_dtype: torch.dtype, M: int, @@ -195,7 +222,10 @@ def test_nvfp4_quantization_extrema_versus_reference( extrema_high: bool, return_transpose: bool, use_cpp_allocator: bool, + row_scaled_nvfp4: bool, ): + maybe_skip_row_scaled_unsupported_quantization(row_scaled_nvfp4, return_transpose) + te_dtype = tex.DType.kFloat4E2M1 device = "cuda" @@ -216,6 +246,7 @@ def test_nvfp4_quantization_extrema_versus_reference( amax_reduction_group=None, with_rht=False, with_post_rht_amax=False, + row_scaled_nvfp4=row_scaled_nvfp4, ) if use_cpp_allocator: @@ -237,6 +268,7 @@ def test_nvfp4_quantization_extrema_versus_reference( ) sx_t = x_nvfp4_sut._columnwise_scale_inv qx_amax = x_nvfp4_sut._amax_rowwise + qx_amax_t = x_nvfp4_sut._amax_columnwise ref_quantizer = NVFP4QuantizerRef( dtype=utils.Fp4Formats.E2M1, @@ -245,6 +277,7 @@ def test_nvfp4_quantization_extrema_versus_reference( pow_2_scales=False, eps=0.0, quant_tile_shape=(1, 16), + row_scaled_nvfp4=row_scaled_nvfp4, ) x_nvfp4_ref = ref_quantizer.quantize(x) @@ -257,6 +290,7 @@ def test_nvfp4_quantization_extrema_versus_reference( x_nvfp4_ref.scale_t.view(dtype=torch.uint8) if x_nvfp4_ref.scale_t is not None else None ) ref_amax = x_nvfp4_ref.global_amax_row + ref_amax_t = x_nvfp4_ref.global_amax_col torch.testing.assert_close(qx, qx_ref, atol=0.0, rtol=0.0) @@ -269,6 +303,7 @@ def test_nvfp4_quantization_extrema_versus_reference( ref_sx_t_shape = sx_t_ref.shape sx_t_valid = sx_t[: ref_sx_t_shape[0], : ref_sx_t_shape[1]] torch.testing.assert_close(sx_t_valid, sx_t_ref, atol=0.0, rtol=0.0) + torch.testing.assert_close(qx_amax_t, ref_amax_t, atol=0.0, rtol=0.0) torch.testing.assert_close(qx_amax, ref_amax, atol=0.0, rtol=0.0) @@ -286,18 +321,22 @@ def test_nvfp4_quantization_extrema_versus_reference( @pytest.mark.parametrize( "use_cpp_allocator", [True, False], ids=["cpp_allocator", "python_allocator"] ) +@pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"]) def test_nvfp4_quantization_boundary_values( x_dtype: torch.dtype, M: int, N: int, return_transpose: bool, use_cpp_allocator: bool, + row_scaled_nvfp4: bool, ): """ Stress rounding/threshold behavior by placing values just below/above many potential bin edges within each 16-element microblock. Validates native vs reference byte-for-byte and scale parity. """ + maybe_skip_row_scaled_unsupported_quantization(row_scaled_nvfp4, return_transpose) + te_dtype = tex.DType.kFloat4E2M1 device = "cuda" @@ -327,6 +366,7 @@ def test_nvfp4_quantization_boundary_values( amax_reduction_group=None, with_rht=False, with_post_rht_amax=False, + row_scaled_nvfp4=row_scaled_nvfp4, ) if use_cpp_allocator: @@ -348,6 +388,7 @@ def test_nvfp4_quantization_boundary_values( ) sx_t = x_nvfp4_sut._columnwise_scale_inv qx_amax = x_nvfp4_sut._amax_rowwise + qx_amax_t = x_nvfp4_sut._amax_columnwise ref_quantizer = NVFP4QuantizerRef( dtype=utils.Fp4Formats.E2M1, @@ -356,6 +397,7 @@ def test_nvfp4_quantization_boundary_values( pow_2_scales=False, eps=0.0, quant_tile_shape=(1, 16), + row_scaled_nvfp4=row_scaled_nvfp4, ) x_nvfp4_ref = ref_quantizer.quantize(x) @@ -368,6 +410,7 @@ def test_nvfp4_quantization_boundary_values( x_nvfp4_ref.scale_t.view(dtype=torch.uint8) if x_nvfp4_ref.scale_t is not None else None ) ref_amax = x_nvfp4_ref.global_amax_row + ref_amax_t = x_nvfp4_ref.global_amax_col torch.testing.assert_close(qx, qx_ref, atol=0.0, rtol=0.0) @@ -381,6 +424,7 @@ def test_nvfp4_quantization_boundary_values( ref_sx_t_shape = sx_t_ref.shape sx_t_valid = sx_t[: ref_sx_t_shape[0], : ref_sx_t_shape[1]] torch.testing.assert_close(sx_t_valid, sx_t_ref, atol=0.0, rtol=0.0) + torch.testing.assert_close(qx_amax_t, ref_amax_t, atol=0.0, rtol=0.0) torch.testing.assert_close(qx_amax, ref_amax, atol=0.0, rtol=0.0) @@ -397,13 +441,17 @@ def test_nvfp4_quantization_boundary_values( @pytest.mark.parametrize( "use_cpp_allocator", [True, False], ids=["cpp_allocator", "python_allocator"] ) +@pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"]) def test_nvfp4_quantization_noncontiguous_inputs( x_dtype: torch.dtype, M: int, N: int, return_transpose: bool, use_cpp_allocator: bool, + row_scaled_nvfp4: bool, ): + maybe_skip_row_scaled_unsupported_quantization(row_scaled_nvfp4, return_transpose) + te_dtype = tex.DType.kFloat4E2M1 device = "cuda" @@ -424,6 +472,7 @@ def test_nvfp4_quantization_noncontiguous_inputs( amax_reduction_group=None, with_rht=False, with_post_rht_amax=False, + row_scaled_nvfp4=row_scaled_nvfp4, ) if use_cpp_allocator: @@ -445,6 +494,7 @@ def test_nvfp4_quantization_noncontiguous_inputs( ) sx_t = x_nvfp4_sut._columnwise_scale_inv qx_amax = x_nvfp4_sut._amax_rowwise + qx_amax_t = x_nvfp4_sut._amax_columnwise ref_quantizer = NVFP4QuantizerRef( dtype=utils.Fp4Formats.E2M1, @@ -453,6 +503,7 @@ def test_nvfp4_quantization_noncontiguous_inputs( pow_2_scales=False, eps=0.0, quant_tile_shape=(1, 16), + row_scaled_nvfp4=row_scaled_nvfp4, ) x_nvfp4_ref = ref_quantizer.quantize(x_nc) @@ -465,6 +516,7 @@ def test_nvfp4_quantization_noncontiguous_inputs( x_nvfp4_ref.scale_t.view(dtype=torch.uint8) if x_nvfp4_ref.scale_t is not None else None ) ref_amax = x_nvfp4_ref.global_amax_row + ref_amax_t = x_nvfp4_ref.global_amax_col # Quantized must match torch.testing.assert_close(qx, qx_ref, atol=0.0, rtol=0.0) @@ -479,5 +531,6 @@ def test_nvfp4_quantization_noncontiguous_inputs( ref_sx_t_shape = sx_t_ref.shape sx_t_valid = sx_t[: ref_sx_t_shape[0], : ref_sx_t_shape[1]] torch.testing.assert_close(sx_t_valid, sx_t_ref, atol=0.0, rtol=0.0) + torch.testing.assert_close(qx_amax_t, ref_amax_t, atol=0.0, rtol=0.0) torch.testing.assert_close(qx_amax, ref_amax, atol=0.0, rtol=0.0) diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py index ed4f73adb..c7c5a5b99 100644 --- a/tests/pytorch/test_backward_override.py +++ b/tests/pytorch/test_backward_override.py @@ -78,6 +78,11 @@ marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4), id="NVFP4BlockScaling", ), + pytest.param( + "nvfp4_row_scaled", + marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4), + id="NVFP4RowScaledBlockScaling", + ), ] @@ -165,7 +170,7 @@ def _maybe_skip_recipe_dtype( ) -> None: if dtype == torch.bfloat16 and not bf16_available: pytest.skip(reason_for_no_bf16) - if recipe_name == "nvfp4": + if recipe_name in ("nvfp4", "nvfp4_row_scaled"): if module_type in ("linear", "layernorm_linear") and dtype not in ( torch.bfloat16, torch.float32, @@ -178,6 +183,14 @@ def _maybe_skip_recipe_dtype( def _maybe_skip_unsupported_recipe_module_combo(recipe_name: str, module_type: str) -> None: if module_type == "ops_linear" and recipe_name == "fp8_block_scaling": pytest.skip("Fusible ops (te_ops.Linear) do not support Float8BlockScaling recipe") + if module_type == "ops_linear" and recipe_name == "nvfp4_row_scaled": + pytest.skip("Row-scaled NVFP4 currently does not support fused te_ops paths.") + + +def _make_quantized_forward_reference_recipe(recipe_name: str) -> recipe.Recipe: + if recipe_name == "nvfp4_row_scaled": + return make_recipe(recipe_name, backward_override="dequantized") + return make_recipe(recipe_name) def _maybe_skip_unsupported_recipe_shape( @@ -195,7 +208,9 @@ def _maybe_skip_unsupported_recipe_shape( " by 32." ) return - if recipe_name == "nvfp4" and (flat_first_dim % 16 != 0 or last_dim % 16 != 0): + if recipe_name in ("nvfp4", "nvfp4_row_scaled") and ( + flat_first_dim % 16 != 0 or last_dim % 16 != 0 + ): pytest.skip( "Linear/LayerNormLinear + NVFP4 requires prod(shape[:-1]) and shape[-1] divisible" " by 16." @@ -220,7 +235,9 @@ def _maybe_skip_unsupported_recipe_shape( pytest.skip( "te_ops.Linear + MXFP8 requires prod(shape[:-1]) and shape[-1] divisible by 32." ) - if recipe_name == "nvfp4" and (flat_first_dim % 16 != 0 or last_dim % 16 != 0): + if recipe_name in ("nvfp4", "nvfp4_row_scaled") and ( + flat_first_dim % 16 != 0 or last_dim % 16 != 0 + ): pytest.skip( "te_ops.Linear + NVFP4 requires prod(shape[:-1]) and shape[-1] divisible by 16." ) @@ -239,9 +256,9 @@ def _maybe_skip_unsupported_grouped_splits(recipe_name: str, m_splits: list[int] ) if recipe_name == "mxfp8" and any(m % 32 != 0 for m in non_empty_splits): pytest.skip("GroupedLinear + MXFP8 requires each non-empty m_split divisible by 32.") - if recipe_name == "nvfp4" and any(m % 16 != 0 for m in non_empty_splits): + if recipe_name in ("nvfp4", "nvfp4_row_scaled") and any(m % 16 != 0 for m in non_empty_splits): pytest.skip("GroupedLinear + NVFP4 requires each non-empty m_split divisible by 16.") - if recipe_name == "nvfp4" and any(m % 64 != 0 for m in non_empty_splits): + if recipe_name in ("nvfp4", "nvfp4_row_scaled") and any(m % 64 != 0 for m in non_empty_splits): pytest.skip( "GroupedLinear + NVFP4 grouped split_quantize currently requires each non-empty " "m_split divisible by 64 due to grouped amax kernel constraints." @@ -847,7 +864,7 @@ def test_linear_like_backward_override_matches_reference( _maybe_skip_unsupported_recipe_shape(recipe_name, input_shape, module_type) in_features = input_shape[-1] - quantized_ref_recipe = make_recipe(recipe_name) + quantized_ref_recipe = _make_quantized_forward_reference_recipe(recipe_name) mode_recipe = make_recipe(recipe_name, backward_override=backward_override) skip_unsupported_backward_override(module_type, mode_recipe, backward_override) @@ -1031,8 +1048,9 @@ def test_grouped_linear_backward_override_matches_reference( num_gemms = len(m_splits) num_tokens = sum(m_splits) - quantized_ref_recipe = make_recipe(recipe_name) + quantized_ref_recipe = _make_quantized_forward_reference_recipe(recipe_name) mode_recipe = make_recipe(recipe_name, backward_override=backward_override) + skip_unsupported_backward_override("grouped_linear", mode_recipe, backward_override) module_quantized_ref = te.GroupedLinear( num_gemms, @@ -1200,6 +1218,7 @@ def test_linear_like_runtime_backward_override_switch_updates_ctx( dy = torch.randn(*input_shape[:-1], out_features, dtype=dtype, device="cuda") default_recipe = make_recipe(recipe_name) + skip_unsupported_backward_override(module_type, default_recipe, None) mode_recipe = make_recipe(recipe_name, backward_override=backward_override) skip_unsupported_backward_override(module_type, mode_recipe, backward_override) @@ -1270,7 +1289,9 @@ def test_grouped_linear_runtime_backward_override_switch_updates_ctx( dy = torch.randn(num_tokens, out_features, dtype=dtype, device="cuda") default_recipe = make_recipe(recipe_name) + skip_unsupported_backward_override("grouped_linear", default_recipe, None) mode_recipe = make_recipe(recipe_name, backward_override=backward_override) + skip_unsupported_backward_override("grouped_linear", mode_recipe, backward_override) *_, default_ctx = _run_grouped_linear_single_step_with_ctx_state( module, @@ -1336,7 +1357,7 @@ def test_fused_linear_paths_match_backward_override_reference( reset_rng_states() - quantized_ref_recipe = make_recipe(recipe_name) + quantized_ref_recipe = _make_quantized_forward_reference_recipe(recipe_name) mode_recipe = make_recipe(recipe_name, backward_override=backward_override) skip_unsupported_backward_override("ops_linear", mode_recipe, backward_override) @@ -1476,7 +1497,7 @@ def test_fused_bias_activation_matches_masked_linear_backward( reset_rng_states() in_features = input_shape[-1] - quantized_ref_recipe = make_recipe(recipe_name) + quantized_ref_recipe = _make_quantized_forward_reference_recipe(recipe_name) mode_recipe = make_recipe(recipe_name, backward_override=backward_override) skip_unsupported_backward_override("ops_linear", mode_recipe, backward_override) @@ -1715,7 +1736,11 @@ def test_backward_override_memory_peak_report( x = torch.randn(*input_shape, dtype=dtype, device="cuda") dy = torch.randn(*input_shape[:-1], out_features, dtype=dtype, device="cuda") - modes = (None, "high_precision", "dequantized") + modes = ( + ("high_precision", "dequantized") + if recipe_name == "nvfp4_row_scaled" + else (None, "high_precision", "dequantized") + ) mode_results: dict[str, dict[str, float] | str] = {} for mode in modes: diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index a782dadc6..33ba65e0d 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -20,17 +20,19 @@ is_fp8_available, is_fp8_block_scaling_available, is_mxfp8_available, + is_nvfp4_available, is_bf16_available, ) from transformer_engine.pytorch.quantization import FP8GlobalStateManager import transformer_engine.pytorch.ops as te_ops from transformer_engine.common import recipe -from utils import ModelConfig, reset_rng_states, skip_unsupported_backward_override +from utils import ModelConfig, recipe_id, reset_rng_states, skip_unsupported_backward_override # Check if FP8 is supported. fp8_available = is_fp8_available() fp8_block_scaling_available = is_fp8_block_scaling_available() mxfp8_available = is_mxfp8_available() +nvfp4_available = is_nvfp4_available() # Reset RNG states. reset_rng_states() @@ -62,6 +64,14 @@ def nvfp4_rht_and_2d_quantization(): return nvfp4_recipe +def nvfp4_row_scaled(): + nvfp4_recipe = recipe.NVFP4BlockScaling(row_scaled_activation=True) + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + return nvfp4_recipe + + def check_rht_usage(recipe: recipe.Recipe) -> bool: # if using RHT, we can only support bf16 # check fp4_quant_fwd_inp, fp4_quant_fwd_weight, fp4_quant_bwd_grad @@ -88,7 +98,9 @@ def get_nvfp4_inp_supported_dtypes(recipe: recipe.Recipe, dtype: torch.dtype) -> fp8_recipes = [] if mxfp8_available: fp8_recipes.append(recipe.MXFP8BlockScaling()) +if nvfp4_available: fp8_recipes.append(nvfp4_rht_and_2d_quantization()) + fp8_recipes.append(nvfp4_row_scaled()) if fp8_block_scaling_available: fp8_recipes.append(recipe.Float8BlockScaling()) if fp8_available: @@ -360,7 +372,7 @@ def _test_cuda_graphs( @pytest.mark.parametrize("module", _test_cuda_graphs_modules) @pytest.mark.parametrize("dtype", dtypes) @pytest.mark.parametrize("fp8_params", (False, True)) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes + [None], ids=lambda r: type(r).__name__) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes + [None], ids=recipe_id) @pytest.mark.parametrize("backward_override", (None, "high_precision", "dequantized")) def test_make_graphed_callables( *, @@ -390,6 +402,8 @@ def test_make_graphed_callables( f"Module not yet supported for {fp8_recipe.__class__.__name__} with CUDA graphs" ) if fp8 and fp8_recipe.nvfp4(): + if getattr(fp8_recipe, "row_scaled_activation", False) and module == "mha": + pytest.skip("Row-scaled NVFP4 CUDA graph coverage applies to GEMM modules.") if dtype not in get_nvfp4_inp_supported_dtypes(fp8_recipe, dtype): pytest.skip( f"Input dtype {dtype} not supported for NVFP4 Recipe" @@ -448,7 +462,7 @@ def test_make_graphed_callables( ) @pytest.mark.parametrize("dtype", dtypes) @pytest.mark.parametrize("fp8_params", (False, True)) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=lambda r: type(r).__name__) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("backward_override", (None, "high_precision", "dequantized")) def test_make_graphed_callables_with_fp8_weight_caching( *, diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 91d4b8901..5f5221af7 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -25,10 +25,16 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.quantization import ( FP8GlobalStateManager, + NVFP4BlockScalingRecipeState, _amax_and_scale_update, ) import transformer_engine.pytorch.ops as te_ops -from transformer_engine.common.recipe import DelayedScaling, Float8BlockScaling, MXFP8BlockScaling +from transformer_engine.common.recipe import ( + DelayedScaling, + Float8BlockScaling, + MXFP8BlockScaling, + NVFP4BlockScaling, +) # Check if FP8 is supported fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) @@ -507,8 +513,30 @@ def test_quantizer_update(self, module_class): y = module(x) +@pytest.mark.skipif(not fp4_available, reason=reason_for_no_fp4) +def test_nvfp4_row_scaled_quantizer_roles(): + recipe = NVFP4BlockScaling(row_scaled_activation=True) + + forward_quantizers = NVFP4BlockScalingRecipeState( + recipe, + mode="forward", + num_quantizers=3, + ).make_quantizers() + assert [q.row_scaled_nvfp4 for q in forward_quantizers] == [True, False, True] + assert not forward_quantizers[0].is_quantizable(torch.empty(16, 16)) + assert forward_quantizers[1].is_quantizable(torch.empty(16, 16)) + + backward_quantizers = NVFP4BlockScalingRecipeState( + recipe, + mode="backward", + num_quantizers=2, + ).make_quantizers() + assert [q.row_scaled_nvfp4 for q in backward_quantizers] == [False, False] + + @pytest.mark.skipif(not fp4_available, reason=reason_for_no_fp4) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=str) +@pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"]) @pytest.mark.parametrize( "M, N", [ @@ -524,12 +552,19 @@ def test_quantizer_update(self, module_class): (8192, 8192), ], ) -def test_fp4_dequantize(dtype, M, N): - q = NVFP4Quantizer() +def test_fp4_dequantize(dtype, row_scaled_nvfp4, M, N): + q = NVFP4Quantizer( + columnwise=not row_scaled_nvfp4, + row_scaled_nvfp4=row_scaled_nvfp4, + ) a = torch.rand((M, N)).cuda().to(dtype=dtype) starting_tensor = q(a) + assert starting_tensor._row_scaled_nvfp4 == row_scaled_nvfp4 + assert starting_tensor._amax_rowwise.numel() == (M if row_scaled_nvfp4 else 1) dequantized_tensor = starting_tensor.dequantize() new_tensor = q(dequantized_tensor) + assert new_tensor._row_scaled_nvfp4 == row_scaled_nvfp4 + assert new_tensor._amax_rowwise.numel() == (M if row_scaled_nvfp4 else 1) torch.testing.assert_close( new_tensor._rowwise_data, starting_tensor._rowwise_data, diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 7f2f24fd6..c811342df 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -38,12 +38,13 @@ import transformer_engine_torch as tex from transformer_engine.pytorch.cpp_extensions import general_gemm from transformer_engine.pytorch.tensor.utils import replace_raw_data -from utils import ModelConfig, skip_unsupported_backward_override +from utils import ModelConfig, recipe_id, skip_unsupported_backward_override # Only run FP8 tests on supported devices. fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) fp8_block_scaling_available, _ = te.is_fp8_block_scaling_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +nvfp4_available, _ = te.is_nvfp4_available(return_reason=True) # Record initial RNG state from script run. seed = 1234 @@ -93,9 +94,18 @@ def nvfp4_vanilla(): return nvfp4_recipe +def nvfp4_row_scaled(): + nvfp4_recipe = recipe.NVFP4BlockScaling(row_scaled_activation=True) + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + return nvfp4_recipe + + fp8_recipes = [] if mxfp8_available: fp8_recipes.append(recipe.MXFP8BlockScaling()) +if nvfp4_available: fp8_recipes.append(nvfp4_vanilla()) # TODO: fix check for this if fp8_block_scaling_available: fp8_recipes.append(recipe.Float8BlockScaling()) @@ -103,6 +113,9 @@ def nvfp4_vanilla(): fp8_recipes.append(recipe.Float8CurrentScaling()) fp8_recipes.append(recipe.DelayedScaling()) fp8_recipes.append(None) +fp8_recipes_with_row_scaled = fp8_recipes.copy() +if nvfp4_available: + fp8_recipes_with_row_scaled.insert(-1, nvfp4_row_scaled()) param_types = [torch.float32, torch.float16] if is_bf16_available(): # bf16 requires sm_80 or higher @@ -402,7 +415,7 @@ def test_sanity_normalization_amp(dtype, model, skip_wgrad, skip_dgrad, normaliz @pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes_with_row_scaled, ids=recipe_id) @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) @pytest.mark.parametrize("model", ["small", "weird"]) @pytest.mark.parametrize("skip_wgrad", all_boolean) @@ -450,7 +463,7 @@ def test_sanity_layernorm_linear( @pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes_with_row_scaled, ids=recipe_id) @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) @pytest.mark.parametrize("model", ["small", "weird"]) @pytest.mark.parametrize("skip_wgrad", all_boolean) @@ -488,7 +501,7 @@ def test_sanity_linear( @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("bs", batch_sizes_with_zero) @pytest.mark.parametrize("model", ["small", "weird"]) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes_with_row_scaled, ids=recipe_id) @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) @pytest.mark.parametrize("fp8_model_params", all_boolean) @pytest.mark.parametrize("use_bias", all_boolean) @@ -529,7 +542,7 @@ def test_sanity_linear_with_zero_tokens( @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("bs", batch_sizes_with_zero) @pytest.mark.parametrize("model", ["small", "weird"]) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes_with_row_scaled, ids=recipe_id) @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) @pytest.mark.parametrize("fp8_model_params", all_boolean) @pytest.mark.parametrize("use_bias", all_boolean) @@ -563,7 +576,12 @@ def test_sanity_grouped_linear( if not is_fp8_supported(config): pytest.skip("Model config does not support FP8") if fp8_recipe.nvfp4(): - pytest.skip("NVFP4 not supported for grouped linear") + if not getattr(fp8_recipe, "row_scaled_activation", False): + pytest.skip("NVFP4 not supported for grouped linear") + if single_param: + pytest.skip("Row-scaled NVFP4 does not support GroupedTensor grouped linear") + if dtype == torch.float16: + pytest.skip("FP16 output for NVFP4 not supported") use_fp8 = fp8_recipe is not None with quantized_model_init(enabled=use_fp8 and fp8_model_params, recipe=fp8_recipe): diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 9d0ed7988..51f72b1e5 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -32,6 +32,7 @@ is_fp8_block_scaling_available, is_nvfp4_available, ) +from utils import recipe_id fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) @@ -47,6 +48,7 @@ _all_recipes.append(recipe.MXFP8BlockScaling()) if nvfp4_available: _all_recipes.append(recipe.NVFP4BlockScaling()) + _all_recipes.append(recipe.NVFP4BlockScaling(row_scaled_activation=True)) # --------------------------------------------------------------------------- @@ -303,7 +305,7 @@ def fn(inp): @pytest.mark.skipif(not fp8_available, reason=reason_for_no_fp8) -@pytest.mark.parametrize("fp8_recipe", _all_recipes, ids=lambda r: type(r).__name__) +@pytest.mark.parametrize("fp8_recipe", _all_recipes, ids=recipe_id) def test_autocast_sanity(fp8_recipe): """Smoke test: torch.nn.Linear inside a single te.autocast with each built-in recipe. Forward + backward under torch.compile(fullgraph=True).""" diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 3b2e50be3..32e44be2a 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -117,7 +117,7 @@ def quantization_tols(name: str) -> dict[str, float]: "mxfp8_block_scaling", ): return dtype_tols(tex.DType.kFloat8E4M3) - if name == "nvfp4": + if name in ("nvfp4", "nvfp4_row_scaled"): return dtype_tols(tex.DType.kFloat4E2M1) raise ValueError(f"Unsupported quantization scheme ({name})") @@ -151,15 +151,39 @@ def make_recipe(name: Optional[str], **recipe_kwargs: Any) -> Optional[Recipe]: disable_2d_quantization=True, **recipe_kwargs, ) + if name == "nvfp4_row_scaled": + return transformer_engine.common.recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + row_scaled_activation=True, + **recipe_kwargs, + ) raise ValueError(f"Unsupported quantization scheme ({name})") +def recipe_id(recipe: Optional[Recipe]) -> str: + """Readable pytest id for a quantization recipe.""" + if not isinstance(recipe, Recipe): + return "None" + if recipe.nvfp4() and recipe.row_scaled_activation: + return "NVFP4RowScaledBlockScaling" + return type(recipe).__name__ + + def skip_unsupported_backward_override( layer_type: str, quant_recipe: Optional[Recipe], backward_override: Optional[str], ) -> None: """Skip known unsupported layer/recipe/backward-override combinations used in tests.""" + if ( + quant_recipe is not None + and quant_recipe.nvfp4() + and getattr(quant_recipe, "row_scaled_activation", False) + and backward_override is None + ): + pytest.skip("Row-scaled NVFP4 does not support default quantized backward.") if backward_override is None: return if quant_recipe is None and backward_override is not None: diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 5d0d3c28e..123362ce1 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -100,6 +100,14 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, int32_t rows = input_tensor->flat_first_dim(); int32_t cols = input_tensor->flat_last_dim(); auto dtype = input_tensor->dtype(); + const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; + if (row_scaled_nvfp4) { + NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, + "Row-scaled NVFP4 quantization does not support 2D quantization."); + NVTE_CHECK(!output_tensor->has_columnwise_data(), + "Row-scaled NVFP4 quantization does not produce columnwise output."); + nvfp4::compute_rowwise_amax(*input_tensor, noop_tensor, output_tensor, stream); + } bool use_optimized_kernel = (dtype == DType::kBFloat16) && (rows % 32 == 0) && (cols % 32 == 0) && output_tensor->has_data(); @@ -126,7 +134,9 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, /*use_stochastic_rounding=*/quant_config_cpp.stochastic_rounding, /*rng_state=*/quant_config_cpp.rng_state, /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization, - /*noop_tensor=*/noop_tensor->data, /*stream=*/stream); + /*row_scaled_nvfp4=*/row_scaled_nvfp4, + /*noop_tensor=*/noop_tensor->data, + /*stream=*/stream); } break; } @@ -239,6 +249,8 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens int32_t rows = grad_tensor->flat_first_dim(); int32_t cols = grad_tensor->flat_last_dim(); auto dtype = grad_tensor->dtype(); + NVTE_CHECK(!output_tensor->row_scaled_nvfp4, + "Backward NVFP4 quantization does not support row-scaled outputs."); bool use_optimized_kernel = (dtype == DType::kBFloat16) && (rows % 32 == 0) && (cols % 32 == 0) && output_tensor->has_data(); @@ -265,7 +277,8 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens /*use_stochastic_rounding=*/quant_config_cpp.stochastic_rounding, /*rng_state=*/quant_config_cpp.rng_state, /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization, - /*noop_tensor=*/noop_tensor->data, /*stream=*/stream); + /*row_scaled_nvfp4=*/false, /*noop_tensor=*/noop_tensor->data, + /*stream=*/stream); } break; } diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index 414320815..d549a050e 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -34,8 +34,9 @@ namespace dequantize_kernel { template __global__ void __launch_bounds__(512) dequantize_fp4_kernel(const void *const input, OType *output, const fp8e4m3 *const scales, - const float *const tensor_amax, const size_t N, const size_t M, - const size_t scale_stride, const size_t num_scale_tiles_X) { + const float *const tensor_amax, const bool row_scaled_nvfp4, + const size_t N, const size_t M, const size_t scale_stride, + const size_t num_scale_tiles_X) { const size_t thread_idx = blockIdx.x * blockDim.x + threadIdx.x; const size_t x = thread_idx % M; const size_t y = thread_idx / M; @@ -63,7 +64,7 @@ __global__ void __launch_bounds__(512) fp4vec value; value.vec = input_vectorized[my_index]; fp8e4m3 scale = scales[my_scale_index]; - float amax = *tensor_amax; + float amax = row_scaled_nvfp4 ? tensor_amax[y] : tensor_amax[0]; constexpr float factor_inv = 1.0 / (6.0 * 448.0); float final_scale = static_cast(scale) * amax * factor_inv; #pragma unroll @@ -90,6 +91,7 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) NVTE_CHECK(output->data.shape == input.data.shape, "Input and output shapes need to match."); const bool with_gemm_swizzled_scales = input.with_gemm_swizzled_scales; + const bool row_scaled_nvfp4 = input.row_scaled_nvfp4; constexpr int FP4_BLOCK_SIZE = 16; const size_t N = input.flat_first_dim(); @@ -103,6 +105,8 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) const size_t threads = 512; const size_t blocks = DIVUP(total, threads); const size_t num_scale_tiles_X = DIVUP(Mread, static_cast(4)); + NVTE_CHECK(!row_scaled_nvfp4 || input.amax.numel() == N, + "Row-scaled NVFP4 dequantization requires one rowwise amax per row."); TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY( output->data.dtype, OType, @@ -112,7 +116,8 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) dequantize_fp4_kernel<<>>( input.data.dptr, reinterpret_cast(output->data.dptr), reinterpret_cast(input.scale_inv.dptr), - reinterpret_cast(input.amax.dptr), N, Mread, input.scale_inv.shape.back(), + reinterpret_cast(input.amax.dptr), row_scaled_nvfp4, N, Mread, + input.scale_inv.shape.back(), num_scale_tiles_X);); // NOLINT(*) ); // NOLINT(*) NVTE_CHECK_CUDA(cudaGetLastError()); diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index f164636e3..9e4aef5a1 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -16,6 +16,8 @@ #include #include +#include + #include "../../common.h" #include "../../util/math.h" #include "../../util/ptx.cuh" @@ -27,6 +29,132 @@ namespace transformer_engine { namespace dispatch { namespace nvfp4 { +namespace rowwise_amax_kernel { + +using namespace ptx; + +#if FP4_TYPE_SUPPORTED + +constexpr int ROWWISE_AMAX_BLOCK_SIZE = 256; +constexpr int ROWWISE_AMAX_SF_VEC_SIZE = 16; + +template +__device__ __forceinline__ void abs_max_2x_update(ptx::FPx2 &dst, + const ptx::FPx2 &val) { + if constexpr (std::is_same_v) { + dst.x = fmaxf(fabsf(dst.x), fabsf(val.x)); + dst.y = fmaxf(fabsf(dst.y), fabsf(val.y)); + } else { + ptx::abs_max_2x(dst, dst, val); + } +} + +template +__device__ __forceinline__ float abs_max_2x_to_float(const ptx::FPx2 &val) { + if constexpr (std::is_same_v) { + return fmaxf(fabsf(val.x), fabsf(val.y)); + } else { + return static_cast(__hmax(__habs(val.x), __habs(val.y))); + } +} + +template +__global__ void +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) +__launch_bounds__(BLOCK_SIZE) +#endif + compute_rowwise_amax_kernel(const int num_rows, const int num_cols, + const IType *__restrict__ input, + float *__restrict__ output_rowwise_amax, + const float *__restrict__ noop) { +#if !defined(__CUDA_ARCH__) || (__CUDA_ARCH__ < 1000) + NVTE_DEVICE_ERROR("SM 10.0+ is required."); +#else + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + using IType2 = typename ptx::FPx2; + + const int row_idx = blockIdx.x; + if (row_idx >= num_rows) return; + + const int num_vec2 = num_cols / 2; + const IType2 *input_row = reinterpret_cast(input + row_idx * num_cols); + + IType2 thread_amax_2x = {static_cast(0.0f), static_cast(0.0f)}; + for (int i = threadIdx.x; i < num_vec2; i += BLOCK_SIZE) { + const IType2 val = input_row[i]; + abs_max_2x_update(thread_amax_2x, val); + } + const float thread_max = abs_max_2x_to_float(thread_amax_2x); + + const float row_amax = + reduce_max(thread_max, threadIdx.x / THREADS_PER_WARP); + + if (threadIdx.x == 0) { + output_rowwise_amax[row_idx] = row_amax; + } +#endif +} + +template +void launch_compute_rowwise_amax(const int num_rows, const int num_cols, const IType *input, + float *output_rowwise_amax, cudaStream_t stream, + const float *noop = nullptr) { + if (num_rows == 0 || num_cols == 0) return; + + dim3 grid(num_rows); + dim3 block(ROWWISE_AMAX_BLOCK_SIZE); + + compute_rowwise_amax_kernel + <<>>(num_rows, num_cols, input, output_rowwise_amax, noop); + NVTE_CHECK_CUDA(cudaGetLastError()); +} + +#endif // FP4_TYPE_SUPPORTED + +} // namespace rowwise_amax_kernel + +inline void compute_rowwise_amax(const Tensor &input, const Tensor *noop, Tensor *output, + cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + using namespace rowwise_amax_kernel; + + const size_t rows = input.flat_first_dim(); + const size_t cols = input.flat_last_dim(); + NVTE_CHECK(cols % ROWWISE_AMAX_SF_VEC_SIZE == 0, + "Row-scaled NVFP4 quantization requires last dim divisible by ", + ROWWISE_AMAX_SF_VEC_SIZE, "."); + + auto *amax_ptr = reinterpret_cast(output->amax.dptr); + NVTE_CHECK(amax_ptr != nullptr, "Row-scaled rowwise amax tensor must be allocated."); + NVTE_CHECK(output->amax.numel() == rows, "Row-scaled rowwise amax must have ", rows, + " entries, got ", output->amax.shape, "."); + + const auto *noop_ptr = reinterpret_cast(noop->data.dptr); + if (input.dtype() == DType::kBFloat16) { + const auto *input_ptr = reinterpret_cast(input.data.dptr); + launch_compute_rowwise_amax<__nv_bfloat16>(static_cast(rows), static_cast(cols), + input_ptr, amax_ptr, stream, noop_ptr); + } else if (input.dtype() == DType::kFloat16) { + const auto *input_ptr = reinterpret_cast(input.data.dptr); + launch_compute_rowwise_amax(static_cast(rows), static_cast(cols), input_ptr, + amax_ptr, stream, noop_ptr); + } else if (input.dtype() == DType::kFloat32) { + const auto *input_ptr = reinterpret_cast(input.data.dptr); + launch_compute_rowwise_amax(static_cast(rows), static_cast(cols), input_ptr, + amax_ptr, stream, noop_ptr); + } else { + NVTE_ERROR( + "Unsupported input dtype for row-scaled NVFP4 quantization. " + "Expected BFloat16, Float16, or Float32."); + } +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED +} + namespace quantize_transpose_kernel { using namespace quantization_and_transposition_SF; @@ -108,7 +236,8 @@ constexpr size_t TOTAL_BANKS_WIDTH = (32 * 4 * 8) / 4; // 256 constexpr size_t THREADS_PER_BANK = TOTAL_BANKS_WIDTH / SCALE_DIM; // 8 = 128 / 16 template + typename IType, bool USE_STOCHASTIC_ROUNDING, bool RETURN_TRANSPOSE, + bool ROW_SCALED_NVFP4> __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_kernel(const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, @@ -508,27 +637,56 @@ __global__ void __launch_bounds__(THREADS_NUM) } } - // 2. Compute E4M3 scaling factor - const nvfp4_scale_t S_dec_b_fp8 = - compute_decoding_scaling_factor(block_amax, S_enc_rowwise); + float block_scale_inverse; + if constexpr (ROW_SCALED_NVFP4) { + // 2. Compute E4M3 scaling factor + const size_t scales_offset_Y = + scales_offset_Y_rowwise + stage * BUFF_DIM_Y + it * THREADS_Y_ROWWISE; + const float S_enc_rowwise_block = + scales_offset_Y < rows + ? compute_global_encode_scaling_factor_FP4(amax_rowwise_ptr[scales_offset_Y]) + : 1.0f; + const float S_dec_rowwise_block = 1.0f / S_enc_rowwise_block; + const nvfp4_scale_t S_dec_b_fp8 = + compute_decoding_scaling_factor(block_amax, S_enc_rowwise_block); + + // Check boundaries + const size_t scales_offset_X = scales_offset_X_rowwise; + const size_t scale_idx_global = scales_offset_Y * scale_stride + scales_offset_X; + + const bool rowwise_scale_is_within_bounds_Y = + (stage_rowwise_scales_offset_Y + it * THREADS_Y_ROWWISE + tid_Y_rowwise) < chunk_rows; + if (rowwise_scale_is_within_bounds_X && rowwise_scale_is_within_bounds_Y) { + scales_ptr[scale_idx_global] = S_dec_b_fp8; + } - // Check boundaries - const size_t scales_offset_Y = - scales_offset_Y_rowwise + stage * BUFF_DIM_Y + it * THREADS_Y_ROWWISE; - const size_t scales_offset_X = scales_offset_X_rowwise; - const size_t scale_idx_global = scales_offset_Y * scale_stride + scales_offset_X; + // Compute "correct" per-block encoding scaling factor + constexpr float float_max = detail::TypeExtrema::max; + block_scale_inverse = + fminf(1.0f / (static_cast(S_dec_b_fp8) * S_dec_rowwise_block), + float_max); // S_enc_b_fp8 + } else { + // 2. Compute E4M3 scaling factor + const nvfp4_scale_t S_dec_b_fp8 = + compute_decoding_scaling_factor(block_amax, S_enc_rowwise); + + // Check boundaries + const size_t scales_offset_Y = + scales_offset_Y_rowwise + stage * BUFF_DIM_Y + it * THREADS_Y_ROWWISE; + const size_t scales_offset_X = scales_offset_X_rowwise; + const size_t scale_idx_global = scales_offset_Y * scale_stride + scales_offset_X; + + const bool rowwise_scale_is_within_bounds_Y = + (stage_rowwise_scales_offset_Y + it * THREADS_Y_ROWWISE + tid_Y_rowwise) < chunk_rows; + if (rowwise_scale_is_within_bounds_X && rowwise_scale_is_within_bounds_Y) { + scales_ptr[scale_idx_global] = S_dec_b_fp8; + } - // const bool rowwise_scale_is_within_bounds_Y = scales_offset_Y < rows; - const bool rowwise_scale_is_within_bounds_Y = - (stage_rowwise_scales_offset_Y + it * THREADS_Y_ROWWISE + tid_Y_rowwise) < chunk_rows; - if (rowwise_scale_is_within_bounds_X && rowwise_scale_is_within_bounds_Y) { - scales_ptr[scale_idx_global] = S_dec_b_fp8; + // Compute "correct" per-block encoding scaling factor + constexpr float float_max = detail::TypeExtrema::max; + block_scale_inverse = fminf(1.0f / (static_cast(S_dec_b_fp8) * S_dec_rowwise), + float_max); // S_enc_b_fp8 } - - // Compute "correct" per-block encoding scaling factor - constexpr float float_max = detail::TypeExtrema::max; - const float block_scale_inverse = fminf( - 1.0f / (static_cast(S_dec_b_fp8) * S_dec_rowwise), float_max); // S_enc_b_fp8 const float2 block_scale_inverse_2x{block_scale_inverse, block_scale_inverse}; // 3. Scale elements @@ -1051,7 +1209,6 @@ __global__ void __launch_bounds__(THREADS_NUM) const size_t scales_offset_X = scales_offset_X_rowwise; const size_t scale_idx_global = scales_offset_Y * scale_stride + scales_offset_X; - // const bool rowwise_scale_is_within_bounds_Y = scales_offset_Y < rows; const bool rowwise_scale_is_within_bounds_Y = (stage_rowwise_scales_offset_Y + it * THREADS_Y_ROWWISE + tid_Y_rowwise) < chunk_rows; if (rowwise_scale_is_within_bounds_X && rowwise_scale_is_within_bounds_Y) { @@ -1162,6 +1319,9 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, using namespace ptx; bool use_stochastic_rounding = quant_config ? quant_config->stochastic_rounding : false; + const bool row_scaled_nvfp4 = output->row_scaled_nvfp4; + NVTE_CHECK(!row_scaled_nvfp4 || !use_2d_quantization, + "Row-scaled NVFP4 quantization does not support 2D quantization."); // If transposed output is allocated, return the transposed data. Otherwise, it's not necesary to // return the transposed data. @@ -1186,6 +1346,10 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, NVTE_CHECK(output->has_data(), "NVFP4 output tensor must be allocated."); NVTE_CHECK(is_fp4_dtype(output->data.dtype), "Output must have FP4 type."); NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); + NVTE_CHECK(!row_scaled_nvfp4 || output->amax.dptr != nullptr, + "Row-scaled NVFP4 quantization requires rowwise amax."); + NVTE_CHECK(!row_scaled_nvfp4 || !output->has_columnwise_data(), + "Row-scaled NVFP4 quantization does not produce columnwise output."); NVTE_CHECK(!output->with_gemm_swizzled_scales, "Output must have scales in compact format."); if (return_transpose) { NVTE_CHECK(output->has_columnwise_data(), "NVFP4 transposed output tensor must be allocated."); @@ -1268,20 +1432,23 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, TRANSFORMER_ENGINE_SWITCH_CONDITION( use_stochastic_rounding, USE_STOCHASTIC_ROUNDING, - TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { - auto kernel = quantize_transpose_nvfp4_kernel; + TRANSFORMER_ENGINE_SWITCH_CONDITION(row_scaled_nvfp4, ROW_SCALED_NVFP4, { + TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { + auto kernel = quantize_transpose_nvfp4_kernel; - if constexpr (use_2d_quantization) { - kernel = quantize_transpose_nvfp4_2D_kernel; - } + if constexpr (use_2d_quantization) { + kernel = quantize_transpose_nvfp4_2D_kernel; + } - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); - kernel<<>>( - tensor_map_input, tensor_map_output, tensor_map_output_transpose, scales_ptr, - scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, cols, - scale_stride, scale_stride_transpose, rng_state); + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); + kernel<<>>( + tensor_map_input, tensor_map_output, tensor_map_output_transpose, scales_ptr, + scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, cols, + scale_stride, scale_stride_transpose, rng_state); + }); });); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh index fc337f607..8adda8213 100644 --- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh @@ -261,14 +261,12 @@ __device__ __forceinline__ void colwise_scaling(const IType *__restrict__ sIn_pt } } -template -__device__ __forceinline__ void rowwise_scaling(const IType *__restrict__ sIn_ptr, - fp4e2m1x2 *__restrict__ sOut_ptr, - nvfp4_scale_t *__restrict__ sSFrowwise_ptr, - const float S_enc_rowwise, const int stage_Y, - const int stage_X, const int buff_in, - const int buff_out, RNG_t &rng, uint4 &random_uint4, - int &rnd_idx) { +template +__device__ __forceinline__ void rowwise_scaling( + const IType *__restrict__ sIn_ptr, fp4e2m1x2 *__restrict__ sOut_ptr, + nvfp4_scale_t *__restrict__ sSFrowwise_ptr, const float S_enc_rowwise, const int stage_Y, + const int stage_X, const int buff_in, const int buff_out, const float *amax_rowwise_ptr, + const size_t row_offset, const size_t rows, RNG_t &rng, uint4 &random_uint4, int &rnd_idx) { using scaling_coeff_type = typename SCALING_COEFFICIENT_TYPE::type; const auto &sIn = *reinterpret_cast(sIn_ptr); @@ -315,9 +313,21 @@ __device__ __forceinline__ void rowwise_scaling(const IType *__restrict__ sIn_pt } const float block_amax = get_amax_of_pair(thread_amax_2x); - const nvfp4_scale_t S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax, S_enc_rowwise); - const scaling_coeff_type SFcoefficient = - compute_nvfp4_scaling_coefficient(S_dec_b_fp8, S_enc_rowwise); + nvfp4_scale_t S_dec_b_fp8; + scaling_coeff_type SFcoefficient; + if constexpr (ROW_SCALED_NVFP4) { + const size_t row_idx = row_offset + stage_Y * TILE_DIM_Y + it_offset_Y_rowwise; + const float S_enc_rowwise_block = + row_idx < rows ? core::compute_global_encode_scaling_factor_FP4(amax_rowwise_ptr[row_idx]) + : 1.0f; + S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax, S_enc_rowwise_block); + SFcoefficient = + compute_nvfp4_scaling_coefficient(S_dec_b_fp8, S_enc_rowwise_block); + } else { + S_dec_b_fp8 = compute_decoding_scaling_factor(block_amax, S_enc_rowwise); + SFcoefficient = + compute_nvfp4_scaling_coefficient(S_dec_b_fp8, S_enc_rowwise); + } // Store scaling factors to SMEM buffer (R2S) if (SF_storing_thread) { @@ -350,7 +360,8 @@ __device__ __forceinline__ void rowwise_scaling(const IType *__restrict__ sIn_pt } } -template +template __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D_kernel( const __grid_constant__ CUtensorMap tensor_map_input, const __grid_constant__ CUtensorMap tensor_map_output, @@ -571,9 +582,9 @@ __global__ void __launch_bounds__(THREADS_NUM) quantize_transpose_nvfp4_tuned_1D ptx::cp_async_bulk_wait_group_read(); // NVFP4 Quantization - rowwise_scaling( + rowwise_scaling( sIn_ptr, sOut_ptr, sSFrowwise_ptr, S_enc_rowwise, stage_Y, stage_X, buff_in, buff_out, - rng, random_uint4, rnd_idx); + amax_rowwise_ptr, block_offset_Y, rows, rng, random_uint4, rnd_idx); if constexpr (RETURN_TRANSPOSE) { colwise_scaling( @@ -680,6 +691,7 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, const bool use_stochastic_rounding = quant_config ? quant_config->stochastic_rounding : false; const bool use_fast_math = quant_config ? quant_config->use_fast_math : false; + const bool row_scaled_nvfp4 = output->row_scaled_nvfp4; // If transposed output is allocated, return the transposed data // Otherwise, it's not necesary to return the transposed data. @@ -694,6 +706,10 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, NVTE_CHECK(output->has_data(), "NVFP4 output tensor must be allocated."); NVTE_CHECK(is_fp4_dtype(output->data.dtype), "Output must have FP4 type."); NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated"); + NVTE_CHECK(!row_scaled_nvfp4 || output->amax.dptr != nullptr, + "Row-scaled NVFP4 quantization requires rowwise amax."); + NVTE_CHECK(!row_scaled_nvfp4 || !output->has_columnwise_data(), + "Row-scaled NVFP4 quantization does not produce columnwise output."); if (return_transpose) { NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), @@ -783,16 +799,20 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, use_stochastic_rounding, USE_STOCHASTIC_ROUNDING, TRANSFORMER_ENGINE_SWITCH_CONDITION( use_fast_math, USE_FAST_MATH, - TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { - auto kernel = quantize_transpose_nvfp4_tuned_1D_kernel; - - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, dshmem_size); - kernel<<>>( - tensor_map_input, tensor_map_output, tensor_map_output_transpose, scales_ptr, - scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, cols, - scale_stride, scale_stride_transpose, rng_state); - }););); + TRANSFORMER_ENGINE_SWITCH_CONDITION( + row_scaled_nvfp4, ROW_SCALED_NVFP4, + TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { + auto kernel = + quantize_transpose_nvfp4_tuned_1D_kernel; + + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + dshmem_size); + kernel<<>>( + tensor_map_input, tensor_map_output, tensor_map_output_transpose, scales_ptr, + scales_transpose_ptr, noop_ptr, amax_rowwise_ptr, amax_colwise_ptr, rows, cols, + scale_stride, scale_stride_transpose, rng_state); + });););); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); #endif // FP4_TYPE_SUPPORTED diff --git a/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp b/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp index 133f1a09e..28218e2b4 100644 --- a/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp +++ b/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp @@ -222,6 +222,14 @@ TensorWrapper CommOverlapCore::get_tensor_chunk(const TensorWrapper &source, siz TensorWrapper chunk(scaling_mode); for (int param_id = 0; param_id < NVTETensorParam::kNVTENumTensorParams; param_id++) { auto param_type = static_cast(param_id); + if (param_type == NVTETensorParam::kNVTEWithGEMMSwizzledScales) { + chunk.set_with_gemm_swizzled_scales(source.get_with_gemm_swizzled_scales()); + continue; + } + if (param_type == NVTETensorParam::kNVTERowScaledNVFP4) { + chunk.set_row_scaled_nvfp4(source.get_row_scaled_nvfp4()); + continue; + } auto param = source.get_parameter(param_type); auto param_dptr = reinterpret_cast(param.data_ptr); auto param_dtype = static_cast(param.dtype); diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index c1b3f8f42..12479f2a9 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -173,6 +173,11 @@ struct Tensor { * Only meaningful for MXFP8 and NVFP4. */ bool with_gemm_swizzled_scales = false; + /*! \brief Whether NVFP4 rowwise amax metadata is row-scaled. + * + * Only meaningful for NVFP4 tensors. + */ + bool row_scaled_nvfp4 = false; /*! Map from NVTETensorParam to parameter sizes */ static constexpr size_t attr_sizes[] = { @@ -183,7 +188,8 @@ struct Tensor { sizeof(NVTEBasicTensor), // kNVTERowwiseScaleInv sizeof(NVTEBasicTensor), // kNVTEColumnwiseScaleInv sizeof(NVTEBasicTensor), // kNVTEColumnwiseAmax - sizeof(uint8_t) // kNVTEWithGEMMSwizzledScales + sizeof(uint8_t), // kNVTEWithGEMMSwizzledScales + sizeof(uint8_t) // kNVTERowScaledNVFP4 }; Tensor() : scaling_mode{NVTE_DELAYED_TENSOR_SCALING}, nvte_tensor{0} {} @@ -199,6 +205,7 @@ struct Tensor { columnwise_scale_inv.clear(); scaling_mode = NVTE_DELAYED_TENSOR_SCALING; with_gemm_swizzled_scales = false; + row_scaled_nvfp4 = false; } explicit operator NVTETensor() const noexcept { return nvte_tensor; } diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 144aea1a0..8589d7045 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -318,6 +318,9 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, const void *alpha, const void *beta, bool use_split_accumulator, int math_sm_count, int m_split, int n_split, bool gemm_producer, const Tensor *inputCounter, cudaStream_t stream) { + NVTE_CHECK(!inputA->row_scaled_nvfp4 && !inputB->row_scaled_nvfp4, + "cuBLAS GEMM does not support row-scaled NVFP4 inputs."); + // Tensor dims in row-major order const int A0 = inputA->flat_first_dim(); const int A1 = inputA->flat_last_dim(); diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index bf9394c98..9fe692dd2 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -440,7 +440,7 @@ void nvte_grouped_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTens /*! \brief Grouped Scaled Bias add for grouped GEMM outputs. * * output[row,col] += bias[col] * scale[row], where biases are per-group -* and scales are per-token (per-row across all groups). +* and scales are per-row across all groups. * Requires uniform last-dimension across all output tensors and bias tensors. */ void nvte_grouped_scaled_bias_add(const NVTEGroupedTensor output, const NVTEGroupedTensor bias, diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index b7461a85d..e9a6f4f73 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -72,6 +72,7 @@ enum NVTETensorParam { kNVTEColumnwiseScaleInv = 5, /*!< Scale inverse tensor for decoding Columnwise Data */ kNVTEColumnwiseAmax = 6, /*!< Columnwise Amax tensor */ kNVTEWithGEMMSwizzledScales = 7, /*!< Whether scaling factors are in format expected by GEMM */ + kNVTERowScaledNVFP4 = 8, /*!< Whether an NVFP4 tensor uses row scaling */ kNVTENumTensorParams }; @@ -765,6 +766,11 @@ class TensorWrapper { nvte_set_tensor_param_v2(tensor_, kNVTEWithGEMMSwizzledScales, &val, sizeof(val)); } + void set_row_scaled_nvfp4(bool row_scaled_nvfp4) { + const auto val = static_cast(row_scaled_nvfp4); + nvte_set_tensor_param_v2(tensor_, kNVTERowScaledNVFP4, &val, sizeof(val)); + } + // Parameter getters NVTEBasicTensor get_parameter(const NVTETensorParam param) const noexcept { @@ -801,6 +807,12 @@ class TensorWrapper { return static_cast(val); } + bool get_row_scaled_nvfp4() const { + uint8_t val = 0; + nvte_get_tensor_param_v2(tensor_, kNVTERowScaledNVFP4, &val, sizeof(val), nullptr); + return static_cast(val); + } + /*! \brief Get an underlying NVTETensor. * * \return NVTETensor held by this TensorWrapper. diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 67b6f8706..0d0b2fd37 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -478,6 +478,10 @@ class NVFP4BlockScaling(Recipe): If set to `True`, stochastic rounding is disabled during quantization for all tensors. disable_2d_quantization : bool, default = False If set to `True`, 1D block scaling with block size 16 is used for all tensors. + row_scaled_activation : bool, default = False + If set to `True`, forward activation quantizers emit row-scaled + NVFP4 tensors. In this mode, rowwise ``amax`` metadata is stored + as a vector with one FP32 value per tensor row. backward_override : {None, 'high_precision', 'dequantized'}, default = None Backward precision mode. None does not modify backward behavior, `high_precision` keeps original high-precision operands for backward, @@ -491,6 +495,7 @@ class NVFP4BlockScaling(Recipe): os.getenv("NVTE_NVFP4_DISABLE_STOCHASTIC_ROUNDING", "0") == "1" ) disable_2d_quantization: bool = os.getenv("NVTE_NVFP4_DISABLE_2D_QUANTIZATION", "0") == "1" + row_scaled_activation: bool = os.getenv("NVTE_NVFP4_ROW_SCALED_ACTIVATION", "0") == "1" fp4_format: Format = Format.E2M1 fp8_format: Format = Format.E4M3 @@ -534,6 +539,7 @@ def __repr__(self) -> str: f"fp8_dpa={self.fp8_dpa}, " f"fp8_mha={self.fp8_mha}, " f"backward_override={self.backward_override}, " + f"row_scaled_activation={self.row_scaled_activation}, " f"fp4_quant_fwd_inp={self.fp4_quant_fwd_inp}, " f"fp4_quant_fwd_weight={self.fp4_quant_fwd_weight}, " f"fp4_quant_bwd_grad={self.fp4_quant_bwd_grad}, " diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 1261879a8..1a52d7601 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -852,6 +852,9 @@ void nvte_set_tensor_param_v2(NVTETensor tensor, NVTETensorParam param, const vo case kNVTEWithGEMMSwizzledScales: t.with_gemm_swizzled_scales = static_cast(*reinterpret_cast(buf)); break; + case kNVTERowScaledNVFP4: + t.row_scaled_nvfp4 = static_cast(*reinterpret_cast(buf)); + break; default: NVTE_ERROR("Unsupported tensor parameter (", static_cast(param), ")"); } @@ -932,6 +935,9 @@ void nvte_get_tensor_param_v2(const NVTETensor tensor, NVTETensorParam param, vo case kNVTEWithGEMMSwizzledScales: *reinterpret_cast(buf) = static_cast(t->with_gemm_swizzled_scales); break; + case kNVTERowScaledNVFP4: + *reinterpret_cast(buf) = static_cast(t->row_scaled_nvfp4); + break; default: NVTE_ERROR("Unsupported tensor parameter (", static_cast(param), ")"); } diff --git a/transformer_engine/common/transpose/cast_transpose.h b/transformer_engine/common/transpose/cast_transpose.h index a5ec2306b..c462b3014 100644 --- a/transformer_engine/common/transpose/cast_transpose.h +++ b/transformer_engine/common/transpose/cast_transpose.h @@ -67,7 +67,7 @@ void quantize_transpose_vector_blockwise_fp4( SimpleTensor &scale_inv_t, SimpleTensor &output, SimpleTensor &output_t, const float epsilon, const bool return_identity, const bool return_transpose, const bool pow2_scale, const bool swizzled_scale, const bool use_stochastic_rounding, - const NVTETensor rng_state_tensor, const bool use_2d_quantization, + const NVTETensor rng_state_tensor, const bool use_2d_quantization, const bool row_scaled_nvfp4, const SimpleTensor &noop_tensor, cudaStream_t stream); } // namespace transformer_engine::detail diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu index d3d3dceca..cf9821f1a 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -316,7 +316,7 @@ __device__ __forceinline__ __nv_fp4x4_e2m1 cvt_fp32_to_fp4_4x(const float2 in01, template + bool kApplyStochasticRounding, bool kIs2DBlockScaling, bool kRowScaledNVFP4> __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpose_kernel( const IType* const input, const float* global_amax, OType* const output_c, OType* const output_t, ScaleType* const tile_scales_inv_c, ScaleType* const tile_scales_inv_t, @@ -509,8 +509,19 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo amax = amax_smem[data_row_idx / kFP4BlockScalingSize][tid_in_warp_x]; } // Step 2.4: Compute scale - ScaleType scale_inv = ComputeDecodeScaleFP4(amax, global_encode_scale_multiplier); - float encode_scale = ComputeEncodeScaleFP4(scale_inv, global_decode_scale); + const size_t row_idx = block_idx_y * kTileDim + r_s; + float row_global_encode_scale = global_encode_scale; + if constexpr (kRowScaledNVFP4) { + row_global_encode_scale = + row_idx < num_rows ? ComputeGlobalEncodeScaleFP4(global_amax[row_idx]) : 1.0f; + } + const float row_global_encode_scale_multiplier = + kRowScaledNVFP4 ? row_global_encode_scale * fp4_max_inv : global_encode_scale_multiplier; + const float row_global_decode_scale = + kRowScaledNVFP4 ? 1.0f / row_global_encode_scale : global_decode_scale; + ScaleType scale_inv = + ComputeDecodeScaleFP4(amax, row_global_encode_scale_multiplier); + float encode_scale = ComputeEncodeScaleFP4(scale_inv, row_global_decode_scale); // Step 2.5: Write scale_inv bool write_scale_inv = is_src_lane; if constexpr (!kAligned) { @@ -708,7 +719,7 @@ void quantize_transpose_vector_blockwise_fp4( SimpleTensor& scale_inv_t, SimpleTensor& output, SimpleTensor& output_t, const float epsilon, const bool return_identity, const bool return_transpose, const bool pow2_scale, const bool swizzled_scale, const bool use_stochastic_rounding, - const NVTETensor rng_state_tensor, const bool use_2d_quantization, + const NVTETensor rng_state_tensor, const bool use_2d_quantization, const bool row_scaled_nvfp4, const SimpleTensor& noop_tensor, cudaStream_t stream) { NVTE_API_CALL(quantize_transpose_vector_blockwise_fp4); #if CUDA_VERSION >= 12080 @@ -722,6 +733,10 @@ void quantize_transpose_vector_blockwise_fp4( NVTE_CHECK(return_identity || !use_2d_quantization, "2D block quantization is only supported when return_identity is true."); + NVTE_CHECK(!row_scaled_nvfp4 || (return_identity && !return_transpose), + "Row-scaled NVFP4 quantization only supports rowwise quantization."); + NVTE_CHECK(!row_scaled_nvfp4 || !use_2d_quantization, + "Row-scaled NVFP4 quantization does not support 2D quantization."); const size_t row_length = input.shape.size() > 0 ? input.shape.at(input.shape.size() - 1) : 1u; size_t num_elements = row_length; @@ -801,35 +816,41 @@ void quantize_transpose_vector_blockwise_fp4( TRANSFORMER_ENGINE_SWITCH_CONDITION( use_2d_quantization, kIs2DBlockScaling, - size_t smem_bytes = kSMemSize * sizeof(InputType); - auto kernel = block_scaled_1d_cast_transpose_kernel< - kReturnIdentity, kReturnTranspose, kPow2Scale, kAligned, - float, InputType, OutputType, ScaleType, kSwizzledScale, - kApplyStochasticRounding, kIs2DBlockScaling>; - if (smem_bytes >= 48 * 1024) { - cudaError_t err = cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - smem_bytes); - NVTE_CHECK(err == cudaSuccess, - "Failed to set dynamic shared memory size."); - } kernel<<>>( - reinterpret_cast(input.dptr), - reinterpret_cast(global_amax.dptr), - reinterpret_cast(output.dptr), - reinterpret_cast(output_t.dptr), - reinterpret_cast(scale_inv.dptr), - reinterpret_cast(scale_inv_t.dptr), row_length, - num_rows, scale_stride_x, scale_stride_y, scale_t_stride_x, - scale_t_stride_y, kScaleBlockDim, epsilon, rng_state, - noop_ptr);) // kIs2DBlockScaling - ) // kApplyStochasticRounding - ) // kSwizzledScale - ) // kAligned - ) // kReturnTranspose - ) // kReturnIdentity - ) // OutputType - ) // InputType + TRANSFORMER_ENGINE_SWITCH_CONDITION( + row_scaled_nvfp4, kRowScaledNVFP4, + + size_t smem_bytes = kSMemSize * sizeof(InputType); + auto kernel = block_scaled_1d_cast_transpose_kernel< + kReturnIdentity, kReturnTranspose, kPow2Scale, kAligned, + float, InputType, OutputType, ScaleType, kSwizzledScale, + kApplyStochasticRounding, kIs2DBlockScaling, + kRowScaledNVFP4>; + if (smem_bytes >= 48 * 1024) { + cudaError_t err = cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, + smem_bytes); + NVTE_CHECK(err == cudaSuccess, + "Failed to set dynamic shared memory size."); + } kernel<<>>( + reinterpret_cast(input.dptr), + reinterpret_cast(global_amax.dptr), + reinterpret_cast(output.dptr), + reinterpret_cast(output_t.dptr), + reinterpret_cast(scale_inv.dptr), + reinterpret_cast(scale_inv_t.dptr), + row_length, num_rows, scale_stride_x, scale_stride_y, + scale_t_stride_x, scale_t_stride_y, kScaleBlockDim, + epsilon, rng_state, + noop_ptr);) // kRowScaledNVFP4 + ) // kIs2DBlockScaling + ) // kApplyStochasticRounding + ) // kSwizzledScale + ) // kAligned + ) // kReturnTranspose + ) // kReturnIdentity + ) // OutputType + ) // InputType NVTE_CHECK_CUDA(cudaGetLastError()); #else diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index 6f3553bf9..edf2c1e1c 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -15,6 +15,8 @@ from ..quantized_tensor import Quantizer from ..tensor.storage.float8_blockwise_tensor_storage import Float8BlockwiseQTensorStorage +from ..tensor.storage.grouped_tensor_storage import GroupedTensorStorage +from ..tensor.storage.nvfp4_tensor_storage import NVFP4TensorStorage from ..tensor.utils import is_custom from ..custom_recipes.gemm import custom_gemm from ...debug.pytorch.debug_quantization import DebugQuantizer @@ -69,6 +71,38 @@ def validate_gemm_scale(scale: Optional[float], required: bool) -> float: return 0.0 +def _is_nvfp4_row_scaled_tensor(tensor: torch.Tensor) -> bool: + """Whether tensor carries row-scaled NVFP4 global amax metadata.""" + return isinstance(tensor, NVFP4TensorStorage) and tensor._row_scaled_nvfp4 + + +def _nvfp4_row_scaled_gemm_inputs( + A: NVFP4TensorStorage, + B: NVFP4TensorStorage, + *, + transa: bool, +) -> Tuple[NVFP4TensorStorage, NVFP4TensorStorage, torch.Tensor]: + """Return GEMM aliases and FP32 output scales for row-scaled NVFP4.""" + A_metadata = A.get_metadata() + weight_amax = A._amax_rowwise if transa else A._amax_columnwise + assert weight_amax is not None and weight_amax.numel() == 1 + A_metadata["amax_rowwise" if transa else "amax_columnwise"] = weight_amax.new_ones(1) + A_metadata["row_scaled_nvfp4"] = False + + B_metadata = B.get_metadata() + rhs_rowwise_amax = B._amax_rowwise + assert rhs_rowwise_amax is not None + B_metadata["amax_rowwise"] = rhs_rowwise_amax.new_ones(1) + B_metadata["row_scaled_nvfp4"] = False + + assert rhs_rowwise_amax.dtype == torch.float32 and weight_amax.dtype == torch.float32 + return ( + NVFP4TensorStorage(**A_metadata), + NVFP4TensorStorage(**B_metadata), + (rhs_rowwise_amax * weight_amax).view(-1, 1), + ) + + def general_gemm( A: torch.Tensor, B: torch.Tensor, @@ -174,7 +208,65 @@ def general_gemm( "beta": beta, } - out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) + if not _is_nvfp4_row_scaled_tensor(A) and not _is_nvfp4_row_scaled_tensor(B): + out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*args, **kwargs) + else: + if _is_nvfp4_row_scaled_tensor(A): + raise NotImplementedError("Row-scaled NVFP4 GEMM does not support row-scaled A.") + assert layout[1] == "N", "Row-scaled NVFP4 GEMM currently supports N-layout B only." + if grad: + raise RuntimeError( + "Row-scaled NVFP4 GEMM currently supports fprop only. " + "Backward NVFP4 gradient quantizers should use scalar global amax." + ) + assert not gelu, "Row-scaled NVFP4 GEMM currently does not support fused GELU." + assert not accumulate, "Row-scaled NVFP4 GEMM currently does not support accumulation." + assert ( + quantization_params is None + ), "Row-scaled NVFP4 GEMM currently does not support output quantization." + assert ub is None, "Row-scaled NVFP4 GEMM currently does not support CommOverlap." + assert ( + extra_output is None + ), "Row-scaled NVFP4 GEMM currently does not support extra output." + assert not bulk_overlap, "Row-scaled NVFP4 GEMM currently does not support bulk overlap." + assert out is None or ( + isinstance(out, torch.Tensor) and not is_custom(out) + ), "Row-scaled NVFP4 GEMM currently supports only plain torch.Tensor outputs." + assert isinstance( + A, NVFP4TensorStorage + ), "Row-scaled NVFP4 GEMM currently requires NVFP4 A." + # cuBLAS folds NVFP4 global amax values into GEMM alpha. Keep the row-scaled + # recipe's global scales out of alpha and apply them in FP32 below. + gemm_A, gemm_B, rowwise_global_scales = _nvfp4_row_scaled_gemm_inputs(A, B, transa=transa) + + requested_out, requested_out_dtype = out, out_dtype + fp32_out = ( + torch.empty_like(requested_out, dtype=torch.float32) + if requested_out is not None + else None + ) + gemm_args = list(args) + gemm_args[0] = gemm_A # A + gemm_args[2] = gemm_B # B + gemm_args[4] = fp32_out # out + gemm_args[5] = None # quantization_params + gemm_args[6] = TE_DType[torch.float32] # out_dtype + gemm_args[7] = None # bias + out, bias_grad, gelu_input, extra_output = tex.generic_gemm(*gemm_args, **kwargs) + out_2d = out.reshape(-1, out.shape[-1]) + + assert rowwise_global_scales.dtype == torch.float32 and out.dtype == torch.float32 + assert rowwise_global_scales.numel() == out_2d.shape[0] + + out_2d.mul_(rowwise_global_scales) + if bias is not None: + out_2d.add_(bias.to(dtype=torch.float32)) + + if requested_out is not None: + requested_out.copy_(out.to(dtype=requested_out.dtype)) + out = requested_out + elif requested_out_dtype is not None and requested_out_dtype != torch.float32: + out = out.to(dtype=requested_out_dtype) if debug_quantizer is not None: out = debug_quantizer.process_gemm_output(out) @@ -229,6 +321,44 @@ def general_grouped_gemm( else: bias_dtype = TE_DType[torch.bfloat16] + if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in A): + raise NotImplementedError("Row-scaled NVFP4 grouped GEMM does not support row-scaled A.") + if any(_is_nvfp4_row_scaled_tensor(tensor) for tensor in B): + assert D_dtype is None, "Row-scaled NVFP4 grouped GEMM currently does not support D_dtype." + if single_output: + assert ( + m_splits is not None + ), "Row-scaled NVFP4 grouped GEMM requires m_splits with single output." + out_init = out[0] if single_output else None + if single_output: + start_idx = 0 + out_views = [] + for i in range(num_gemms): + size = m_splits[i] + out_views.append(out_init[start_idx : start_idx + size]) + start_idx += size + else: + out_views = out + for i in range(num_gemms): + if out_views[i].numel() == 0: + continue + general_gemm( + A[i], + B[i], + quantization_params=quantization_params[i], + out_dtype=out_views[i].dtype, + out=out_views[i], + gelu=gelu, + accumulate=accumulate, + layout=layout, + bias=bias[i] if use_bias else None, + use_split_accumulator=use_split_accumulator, + grad=grad, + ) + if single_output: + out = out_init + return out, grad_bias, gelu_input + if isinstance(quantization_params[0], DebugQuantizer): assert not gelu, "GELU not supported in debug mode" if single_output: @@ -350,6 +480,13 @@ def general_grouped_gemm_for_grouped_tensor( if is_discrete_in and is_discrete_out: raise ValueError("Both A and out are discrete. This is not supported yet.") + if isinstance(A, GroupedTensorStorage) and A.row_scaled_nvfp4: + raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") + if isinstance(B, GroupedTensorStorage) and B.row_scaled_nvfp4: + raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") + if isinstance(out, GroupedTensorStorage) and out.row_scaled_nvfp4: + raise NotImplementedError("Row-scaled NVFP4 GroupedTensor GEMM is not supported yet.") + if is_discrete_out: # wgrad case. grouped_gemm_impl = tex.te_general_grouped_gemm_for_discrete_out diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 8e3bcdd5b..8f5b8294e 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -320,6 +320,8 @@ class NVFP4Quantizer : public Quantizer { // 2D block scaling bool with_2d_quantization; bool stochastic_rounding; + // Whether tensors emitted by this quantizer use row-scaled NVFP4 metadata. + bool row_scaled_nvfp4; int rht_matrix_random_sign_mask_t; at::Tensor rht_matrix; diff --git a/transformer_engine/pytorch/csrc/extensions/activation.cpp b/transformer_engine/pytorch/csrc/extensions/activation.cpp index 2df3b6655..cab9fab30 100644 --- a/transformer_engine/pytorch/csrc/extensions/activation.cpp +++ b/transformer_engine/pytorch/csrc/extensions/activation.cpp @@ -42,8 +42,9 @@ py::object activation_helper(const at::Tensor& input, py::handle quantizer, int } else if (detail::IsNVFP4Quantizers(quantizer.ptr())) { auto nvfp4_quantizer_cpp = dynamic_cast(quantizer_cpp.get()); NVTE_CHECK(nvfp4_quantizer_cpp != nullptr, "Could not cast to NVFP4 quantizer"); - if (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax) { - // Post-RHT amax is handled within NVFP4 quantizer + if (nvfp4_quantizer_cpp->row_scaled_nvfp4 || + (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax)) { + // Amax is handled within NVFP4 quantizer impl = Impl::UNFUSED; } else { impl = Impl::FUSED_ACTIVATION_AMAX_NVFP4; @@ -154,8 +155,9 @@ py::object dactivation_helper(const at::Tensor& grad_output, const at::Tensor& i } else if (detail::IsNVFP4Quantizers(quantizer.ptr())) { auto nvfp4_quantizer_cpp = dynamic_cast(quantizer_cpp.get()); NVTE_CHECK(nvfp4_quantizer_cpp != nullptr, "Could not cast to NVFP4 quantizer"); - if (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax) { - // Post-RHT amax is handled within NVFP4 quantizer + if (nvfp4_quantizer_cpp->row_scaled_nvfp4 || + (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax)) { + // Amax is handled within NVFP4 quantizer impl = Impl::UNFUSED; } else { impl = Impl::FUSED_ACTIVATION_AMAX_NVFP4; diff --git a/transformer_engine/pytorch/csrc/extensions/bias.cpp b/transformer_engine/pytorch/csrc/extensions/bias.cpp index 0cf2025f1..4a78dde38 100644 --- a/transformer_engine/pytorch/csrc/extensions/bias.cpp +++ b/transformer_engine/pytorch/csrc/extensions/bias.cpp @@ -152,8 +152,9 @@ std::vector dact_dbias( } else if (detail::IsNVFP4Quantizers(quantizer_py.ptr())) { auto nvfp4_quantizer_cpp = dynamic_cast(quantizer_cpp.get()); NVTE_CHECK(nvfp4_quantizer_cpp != nullptr, "Could not cast to NVFP4 quantizer"); - if (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax) { - // Post-RHT amax is handled within NVFP4 quantizer + if (nvfp4_quantizer_cpp->row_scaled_nvfp4 || + (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax)) { + // Amax is handled within NVFP4 quantizer impl = Impl::UNFUSED; } else { impl = Impl::FUSED_DACT_AMAX_NVFP4; diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 50fe4c109..9e1f381bf 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -798,7 +798,13 @@ std::tuple, std::vector, bool> bulk_alloc // Quantization parameters const auto rowwise_usage = quantizer_cpp_list[0]->rowwise_usage; + const bool row_scaled_nvfp4 = quantizer_cpp_list[0]->row_scaled_nvfp4; const auto columnwise_usage = quantizer_cpp_list[0]->columnwise_usage; + if (row_scaled_nvfp4) { + NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 bulk allocation requires rowwise usage."); + NVTE_CHECK(!columnwise_usage, + "Row-scaled NVFP4 bulk allocation does not support columnwise usage."); + } const auto scaling_mode = quantizer_cpp_list[0]->get_scaling_mode(); const auto fp4_dtype = quantizer_cpp_list[0]->dtype; const bool with_gemm_swizzled_scales = false; /// TODO (tmoon) Enable based on optimize_for_gemm; @@ -828,6 +834,16 @@ std::tuple, std::vector, bool> bulk_alloc } return fp4_shape; }; + auto flat_first_dim = [](const std::vector &shape) -> size_t { + if (shape.empty()) { + return 1; + } + size_t rows = 1; + for (size_t i = 0; i + 1 < shape.size(); ++i) { + rows *= shape[i]; + } + return rows; + }; // Allocate row-wise data std::vector rowwise_data_list, rowwise_scale_list, amax_rowwise_list; @@ -866,7 +882,11 @@ std::tuple, std::vector, bool> bulk_alloc // Note: Multi-quantize kernel does not require contiguous amaxes. const auto offset = roundup(buffer_size, 16); amax_offsets.push_back(offset); - buffer_size = offset + 4; + size_t amax_size = 4; + if (row_scaled_nvfp4) { + amax_size *= flat_first_dim(rowwise_data_shapes[i]); + } + buffer_size = offset + amax_size; } // Allocate full buffer @@ -879,8 +899,12 @@ std::tuple, std::vector, bool> bulk_alloc data_offsets[i], torch::kUInt8)); rowwise_scale_list.emplace_back( make_torch_view(buffer, rowwise_scale_shapes[i], scale_offsets[i], torch::kUInt8)); + std::vector amax_shape{1}; + if (row_scaled_nvfp4) { + amax_shape = {flat_first_dim(rowwise_data_shapes[i])}; + } amax_rowwise_list.emplace_back( - make_torch_view(buffer, std::vector{1}, amax_offsets[i], torch::kFloat32)); + make_torch_view(buffer, amax_shape, amax_offsets[i], torch::kFloat32)); } } @@ -960,9 +984,10 @@ std::tuple, std::vector, bool> bulk_alloc py::object amax_columnwise = columnwise_usage ? py::cast(amax_columnwise_list[i]) : py::none(); // Construct Python tensor - tensor_py_list.emplace_back(NVFP4TensorClass( - rowwise_data, rowwise_scale, columnwise_data, columnwise_scale, amax_rowwise, - amax_columnwise, fp4_dtype, quantizer_py_list[i], with_gemm_swizzled_scales)); + tensor_py_list.emplace_back(NVFP4TensorClass(rowwise_data, rowwise_scale, columnwise_data, + columnwise_scale, amax_rowwise, amax_columnwise, + fp4_dtype, quantizer_py_list[i], + with_gemm_swizzled_scales, row_scaled_nvfp4)); // Construct C++ tensor // Use a TensorWrapper variable to hold the output of makeTransformerEngineTensor, @@ -979,11 +1004,12 @@ std::tuple, std::vector, bool> bulk_alloc rowwise_usage ? rowwise_scale_shapes[i] : std::vector{0}, columnwise_usage ? columnwise_scale_shapes[i] : std::vector{0}, scaling_mode); tensor_wrapper.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); + tensor_wrapper.set_row_scaled_nvfp4(row_scaled_nvfp4); // Set the amax rowwise and amax columnwise if available if (rowwise_usage) { tensor_wrapper.set_amax(amax_rowwise_list[i].data_ptr(), DType::kFloat32, - std::vector{1}); + getTensorShape(amax_rowwise_list[i])); } if (columnwise_usage) { tensor_wrapper.set_columnwise_amax(amax_columnwise_list[i].data_ptr(), DType::kFloat32, @@ -1455,7 +1481,16 @@ std::vector split_quantize(const at::Tensor &tensor, return detail::IsNVFP4Quantizers(quantizer.ptr()); })) { allocation_method = AllocationMethod::BULK_NVFP4; - quantization_method = QuantizationMethod::FUSED_NVFP4; + const bool has_row_scaled_nvfp4 = + std::any_of(quantizer_cpp_list.begin(), quantizer_cpp_list.end(), + [](const std::unique_ptr &quantizer) { + return static_cast(quantizer.get())->row_scaled_nvfp4; + }); + if (has_row_scaled_nvfp4) { + quantization_method = QuantizationMethod::UNFUSED; + } else { + quantization_method = QuantizationMethod::FUSED_NVFP4; + } } } @@ -1492,7 +1527,8 @@ std::vector split_quantize(const at::Tensor &tensor, bool contiguous_data_and_scale = false; std::tie(output_py_list, output_cpp_list, contiguous_data_and_scale) = bulk_allocate_nvfp4_tensors(split_shapes, quantizer_list, nvfp4_quantizers); - if (!input_shape.empty() && input_shape.back() % 128 != 0) { + if (quantization_method == QuantizationMethod::FUSED_NVFP4 && !input_shape.empty() && + input_shape.back() % 128 != 0) { static std::once_flag once_unfused_nvfp4_fallback_warning; std::call_once(once_unfused_nvfp4_fallback_warning, []() { NVTE_WARN( diff --git a/transformer_engine/pytorch/csrc/extensions/normalization.cpp b/transformer_engine/pytorch/csrc/extensions/normalization.cpp index fb4c7aa1c..4887b59c2 100644 --- a/transformer_engine/pytorch/csrc/extensions/normalization.cpp +++ b/transformer_engine/pytorch/csrc/extensions/normalization.cpp @@ -120,8 +120,9 @@ std::vector layernorm_fwd(py::handle input, py::handle weight, Maybe } else if (detail::IsNVFP4Quantizers(quantizer.ptr())) { auto nvfp4_quantizer_cpp = dynamic_cast(quantizer_cpp.get()); NVTE_CHECK(nvfp4_quantizer_cpp != nullptr, "Could not cast to NVFP4 quantizer"); - if (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax) { - // Post-RHT amax is handled within NVFP4 quantizer + if (nvfp4_quantizer_cpp->row_scaled_nvfp4 || + (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax)) { + // Amax is handled within NVFP4 quantizer impl = Impl::UNFUSED; } else if (!transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")) { // TE kernel supports amax output @@ -357,8 +358,9 @@ std::vector rmsnorm_fwd(const py::handle &input, const py::handle &w } else if (detail::IsNVFP4Quantizers(quantizer.ptr())) { auto nvfp4_quantizer_cpp = dynamic_cast(quantizer_cpp.get()); NVTE_CHECK(nvfp4_quantizer_cpp != nullptr, "Could not cast to NVFP4 quantizer"); - if (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax) { - // Post-RHT amax is handled within NVFP4 quantizer + if (nvfp4_quantizer_cpp->row_scaled_nvfp4 || + (nvfp4_quantizer_cpp->with_rht && nvfp4_quantizer_cpp->with_post_rht_amax)) { + // Amax is handled within NVFP4 quantizer impl = Impl::UNFUSED; } else if (!transformer_engine::getenv("NVTE_NORM_FWD_USE_CUDNN")) { // TE kernel supports amax output diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index da91e5c17..8f2de325a 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1696,6 +1696,7 @@ NVFP4Quantizer::NVFP4Quantizer(const py::handle& quantizer) : Quantizer(quantize this->with_post_rht_amax = quantizer.attr("with_post_rht_amax").cast(); this->with_2d_quantization = quantizer.attr("with_2d_quantization").cast(); this->stochastic_rounding = quantizer.attr("stochastic_rounding").cast(); + this->row_scaled_nvfp4 = quantizer.attr("row_scaled_nvfp4").cast(); // Get amax reduction group if needed for NVFP4 AG const bool with_amax_reduction = quantizer.attr("with_amax_reduction").cast(); @@ -1747,6 +1748,12 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve NVTE_CHECK(flat_last_dim % NVFP4_BLOCK_SIZE == 0, "NVFP4 requires tensor dims that are divisible by ", NVFP4_BLOCK_SIZE, " (got shape=", shape, ")"); + const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; + if (row_scaled_nvfp4) { + NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 quantization requires rowwise usage."); + NVTE_CHECK(!columnwise_usage, + "Row-scaled NVFP4 quantization does not support columnwise usage."); + } const auto rowwise_scale_inv_shape = get_scale_shape(shape, false); const auto columnwise_scale_inv_shape = get_scale_shape(shape, true); @@ -1760,9 +1767,10 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve rowwise_scale_inv_shape.end()); rowwise_data_tensor = at::empty(convert_shape_for_fp4(shape_int64), bit8_tensor_opts); rowwise_scale_inv_tensor = at::empty(scale_inv_shape_int64, bit8_tensor_opts); + const int64_t amax_rows = row_scaled_nvfp4 ? static_cast(flat_first_dim) : 1; // hadamard amax kernel will zero out pointer with ZeroAmaxKernel // nvte_compute_amax_with_config will zero out the pointer if needed - amax_rowwise = at::empty({1}, bit32_tensor_opts); + amax_rowwise = at::empty({amax_rows}, bit32_tensor_opts); } if (columnwise_usage) { const std::vector scale_inv_shape_int64(columnwise_scale_inv_shape.begin(), @@ -1805,6 +1813,7 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve kwargs["fp4_dtype"] = py::cast(this->dtype); kwargs["quantizer"] = this->quantizer; kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); + kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); kwargs["fake_dtype"] = GetATenDType(dtype); py::tuple args(0); @@ -1833,6 +1842,7 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve kwargs["fp4_dtype"] = py::cast(this->dtype); kwargs["quantizer"] = this->quantizer; kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); + kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); py::tuple args(0); PyObject* result = PyObject_Call(reinterpret_cast(NVFP4TensorPythonClass), args.ptr(), kwargs.ptr()); @@ -1850,7 +1860,7 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve out_cpp.set_rowwise_data(rowwise_data_tensor.data_ptr(), DType::kFloat4E2M1, shape); out_cpp.set_rowwise_scale_inv(rowwise_scale_inv_tensor.data_ptr(), DType::kFloat8E4M3, rowwise_scale_inv_shape); - out_cpp.set_amax(amax_rowwise.data_ptr(), DType::kFloat32, std::vector{1}); + out_cpp.set_amax(amax_rowwise.data_ptr(), DType::kFloat32, getTensorShape(amax_rowwise)); } if (columnwise_usage) { // enforce 2D shape to avoid [S, B, H] shape and B and be 1 @@ -1865,6 +1875,7 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve std::vector{1}); } out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); + out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); this->set_quantization_params(&out_cpp); return {std::move(out_cpp), std::move(out_py)}; @@ -1892,6 +1903,12 @@ std::pair NVFP4Quantizer::create_grouped_tenso std::optional rowwise_amax; std::optional columnwise_amax; const std::vector logical_shape_vec = {logical_first_dim, logical_last_dim}; + const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; + if (row_scaled_nvfp4) { + NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 grouped quantization requires rowwise usage."); + NVTE_CHECK(!columnwise_usage, + "Row-scaled NVFP4 grouped quantization does not support columnwise usage."); + } const int64_t total_data_elements = total_elements / 2; @@ -1900,7 +1917,9 @@ std::pair NVFP4Quantizer::create_grouped_tenso const auto scale_shape = get_scale_shape(logical_shape_vec, false); const int64_t total_scale_elements = static_cast(product(scale_shape)); rowwise_scale_inv = at::empty({total_scale_elements}, uint8_opts); - rowwise_amax = at::empty({static_cast(num_tensors)}, float_opts); + const int64_t amax_elements = row_scaled_nvfp4 ? static_cast(logical_first_dim) + : static_cast(num_tensors); + rowwise_amax = at::empty({amax_elements}, float_opts); } if (columnwise_usage) { @@ -1958,6 +1977,7 @@ std::pair NVFP4Quantizer::create_grouped_tenso kwargs["last_dims"] = py::none(); kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); kwargs["with_gemm_swizzled_scales"] = this->optimize_for_gemm; + kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); if (result == nullptr) { PyErr_Print(); @@ -1975,15 +1995,22 @@ std::pair NVFP4Quantizer::create_unquantized_tensor_w auto [out_cpp, out_py] = NoneQuantizer(py::none()).create_tensor(shape, dtype); // Register amax pointer from quantized tensor - void* amax_ptr = quantized_tensor.amax(); + auto rowwise_amax = quantized_tensor.get_amax(); + auto columnwise_amax = quantized_tensor.get_columnwise_amax(); + + void* amax_ptr = rowwise_amax.data_ptr; + std::vector amax_shape = convertShape(rowwise_amax.shape); if (amax_ptr == nullptr) { - amax_ptr = quantized_tensor.get_columnwise_amax().data_ptr; + amax_ptr = columnwise_amax.data_ptr; + amax_shape = convertShape(columnwise_amax.shape); } NVTE_CHECK(amax_ptr != nullptr, "Could not extract amax pointer from NVFP4 tensor."); - out_cpp.set_amax(amax_ptr, DType::kFloat32, std::vector{1}); + out_cpp.set_amax(amax_ptr, DType::kFloat32, amax_shape); // Zero out amax - NVTE_CHECK_CUDA(cudaMemsetAsync(amax_ptr, 0, sizeof(float), at::cuda::getCurrentCUDAStream())); + const size_t amax_numel = product(amax_shape); + NVTE_CHECK_CUDA( + cudaMemsetAsync(amax_ptr, 0, amax_numel * sizeof(float), at::cuda::getCurrentCUDAStream())); return {std::move(out_cpp), std::move(out_py)}; } @@ -2031,6 +2058,13 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( } } const size_t flat_last_dim = shape.size() > 0 ? shape.back() : 1; + const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; + if (row_scaled_nvfp4) { + NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 quantization requires rowwise usage."); + NVTE_CHECK(!columnwise_usage, + "Row-scaled NVFP4 quantization does not support columnwise usage."); + } + tensor.attr("_row_scaled_nvfp4") = py::cast(row_scaled_nvfp4); // Coerce row-wise data if (rowwise_usage) { @@ -2048,11 +2082,12 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( rowwise_scale_inv = at::empty(scale_inv_shape_int64, opts); tensor.attr("_rowwise_scale_inv") = *rowwise_scale_inv; } - if (!amax_rowwise) { + const int64_t amax_rows = row_scaled_nvfp4 ? static_cast(flat_first_dim) : 1; + if (!amax_rowwise || amax_rowwise->numel() != amax_rows) { const auto opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); // hadamard amax kernel will zero out pointer with ZeroAmaxKernel // nvte_compute_amax_with_config will zero out the pointer if needed - amax_rowwise = at::empty({1}, opts); + amax_rowwise = at::empty({amax_rows}, opts); tensor.attr("_amax_rowwise") = *amax_rowwise; } } else { // rowwise_usage == false @@ -2118,7 +2153,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( out_cpp.set_rowwise_data(rowwise_data->data_ptr(), DType::kFloat4E2M1, shape); out_cpp.set_rowwise_scale_inv(rowwise_scale_inv->data_ptr(), DType::kFloat8E4M3, getTensorShape(*rowwise_scale_inv)); - out_cpp.set_amax(amax_rowwise->data_ptr(), DType::kFloat32, std::vector{1}); + out_cpp.set_amax(amax_rowwise->data_ptr(), DType::kFloat32, getTensorShape(*amax_rowwise)); } if (columnwise_usage) { // enforce 2D shape to avoid [S, B, H] shape and B and be 1 @@ -2133,6 +2168,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( std::vector{1}); } out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); + out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); this->set_quantization_params(&out_cpp); return {std::move(out_cpp), std::move(tensor)}; @@ -2241,6 +2277,18 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou } size_t cols = input.size(input.ndim() - 1); + const bool row_scaled_nvfp4 = out.get_row_scaled_nvfp4(); + if (row_scaled_nvfp4) { + NVTE_CHECK(!this->with_rht, "Row-scaled NVFP4 quantization does not support RHT."); + NVTE_CHECK(!this->with_2d_quantization, + "Row-scaled NVFP4 quantization does not support 2D quantization."); + NVTE_CHECK(!this->stochastic_rounding, + "Row-scaled NVFP4 quantization does not support stochastic rounding."); + NVTE_CHECK(!this->with_amax_reduction, + "Row-scaled NVFP4 quantization does not support amax reduction."); + NVTE_CHECK(cols % 16 == 0, "Row-scaled NVFP4 quantization requires last dim divisible by 16."); + } + // Restriction for the RHT cast fusion kernel because we are using MMA hardware for computing RHT bool eligible_for_rht_cast_fusion = input.dtype() == DType::kBFloat16 && rows % 64 == 0 && cols % 128 == 0; @@ -2307,7 +2355,7 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou "Use with_post_rht_amax=true instead."); } } else { // Without RHT - if (compute_amax) { + if (compute_amax && !row_scaled_nvfp4) { // Amax pointers auto rowwise_amax_ptr = out.get_amax().data_ptr; auto columnwise_amax_ptr = out.get_columnwise_amax().data_ptr; @@ -2408,6 +2456,8 @@ void NVFP4Quantizer::quantize(const TensorWrapper& input, TensorWrapper& out, } void NVFP4Quantizer::quantize_with_amax(TensorWrapper& input, TensorWrapper& out) { + NVTE_CHECK(!out.get_row_scaled_nvfp4(), + "quantize_with_amax is not supported for row-scaled NVFP4 quantization."); // Update output tensor amaxes with input tensor amax auto input_amax_ptr = input.amax(); auto output_rowwise_amax_ptr = out.get_amax().data_ptr; diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index e13554a98..37ab0b053 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -134,6 +134,7 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) const bool rowwise_usage = !(tensor.attr("_rowwise_data").is_none()); const bool columnwise_usage = !(tensor.attr("_columnwise_data").is_none()); const bool with_gemm_swizzled_scales = tensor.attr("_with_gemm_swizzled_scales").cast(); + const bool row_scaled_nvfp4 = tensor.attr("_row_scaled_nvfp4").cast(); NVTE_CHECK(rowwise_usage || columnwise_usage, "No data found for NVFP4 Tensor."); @@ -163,6 +164,7 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) // Scale layout ret.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); + ret.set_row_scaled_nvfp4(row_scaled_nvfp4); // Quantizer state quantizer->set_quantization_params(&ret); diff --git a/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py b/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py index dd01ae05d..12f8ef8f5 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py @@ -350,9 +350,17 @@ def __init__( pow_2_scales: bool = False, eps: float = 0.0, quant_tile_shape: Tuple[int, int] = (1, 16), + row_scaled_nvfp4: bool = False, with_rht: bool = False, with_random_sign_mask: bool = True, ): + if row_scaled_nvfp4: + if not rowwise: + raise ValueError("Row-scaled NVFP4 reference quantization requires rowwise usage.") + if columnwise: + raise ValueError( + "Row-scaled NVFP4 reference quantization does not support columnwise usage." + ) super().__init__(rowwise=rowwise, columnwise=columnwise) self.internal = True @@ -360,6 +368,7 @@ def __init__( self.pow_2_scales = pow_2_scales self.eps = eps self.quant_tile_shape = quant_tile_shape + self.row_scaled_nvfp4 = row_scaled_nvfp4 self.with_rht = with_rht self.with_random_sign_mask = with_random_sign_mask @@ -447,6 +456,7 @@ def _quantize_blockwise_reference( tile_len_y: int, *, pow_2_scales: bool, + row_scaled_nvfp4: bool = False, eps: float, # pylint: disable=unused-argument ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -488,6 +498,9 @@ def _quantize_blockwise_reference( decode_scale.to(torch.float32), ) else: + if row_scaled_nvfp4: + global_amax = global_amax.to(torch.float32).view(m, 1, 1) + global_encode_scale = torch.div(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX, global_amax) global_encode_scale = torch.min( global_encode_scale, @@ -497,8 +510,15 @@ def _quantize_blockwise_reference( dtype=torch.float32, ), ) - if global_encode_scale == torch.tensor(0.0, device=x.device, dtype=torch.float32): - global_encode_scale = torch.tensor(1.0, device=x.device, dtype=torch.float32) + if global_encode_scale.numel() == 1: + if global_encode_scale == torch.tensor(0.0, device=x.device, dtype=torch.float32): + global_encode_scale = torch.tensor(1.0, device=x.device, dtype=torch.float32) + else: + global_encode_scale = torch.where( + global_encode_scale == 0.0, + torch.ones_like(global_encode_scale), + global_encode_scale, + ) global_decode_scale = torch.div(1.0, global_encode_scale) global_encode_scale_multiplier = global_encode_scale * torch.reciprocal(FLOAT4_E2M1_MAX) @@ -609,6 +629,8 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ raise ValueError( f"MXFP4 only supports 1x32 tile shape, got {self.quant_tile_shape}" ) + if self.row_scaled_nvfp4: + raise ValueError("Row-scaled NVFP4 is only supported for NVFP4 (non-pow2) mode.") # TODO(etsykunov): Fix bug where global_amax_row and # global_amax_col are not defined # global_amax = torch.empty(0, device=tensor.device, dtype=torch.float32) @@ -625,13 +647,22 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ if self.with_rht else tensor.t().contiguous() ) - # Compute amax for rowwise and columnwise paths separately - global_amax_row = torch.max(torch.abs(row_input)).to(torch.float32).view(1) - global_amax_col = ( - torch.max(torch.abs(col_input)).to(torch.float32).view(1) - if self.columnwise_usage - else global_amax_row - ) + if self.row_scaled_nvfp4: + if self.quant_tile_shape != (1, 16): + raise ValueError( + "Row-scaled NVFP4 only supports NVFP4 1x16 tile shape, " + f"got {self.quant_tile_shape}" + ) + global_amax_row = torch.max(torch.abs(row_input), dim=1).values.to(torch.float32) + global_amax_col = global_amax_row + else: + # Compute amax for rowwise and columnwise paths separately + global_amax_row = torch.max(torch.abs(row_input)).to(torch.float32).view(1) + global_amax_col = ( + torch.max(torch.abs(col_input)).to(torch.float32).view(1) + if self.columnwise_usage + else global_amax_row + ) transpose_scales = False @@ -648,6 +679,7 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ self.quant_tile_shape[1], self.quant_tile_shape[0], pow_2_scales=self.pow_2_scales, + row_scaled_nvfp4=self.row_scaled_nvfp4, eps=self.eps, ) if transpose_scales: @@ -868,7 +900,11 @@ def qgemm( partial_alpha = qresult_x.global_amax_col * qresult_w.global_amax_col else: partial_alpha = qresult_x.global_amax_row * qresult_w.global_amax_row - alpha = torch.div(partial_alpha, factor).squeeze(-1) + if partial_alpha.numel() > 1 and partial_alpha.numel() == high_precision_x.shape[0]: + partial_alpha = partial_alpha.view(-1, 1) + else: + partial_alpha = partial_alpha.squeeze(-1) + alpha = torch.div(partial_alpha, factor) M, K = high_precision_x.shape N, K_w = high_precision_w.shape diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 9956fb77e..e9f009d93 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1375,6 +1375,7 @@ def _make_quantizer(idx: int) -> NVFP4Quantizer: with_post_rht_amax=qparams.random_hadamard_transform, with_2d_quantization=qparams.fp4_2d_quantization, stochastic_rounding=qparams.stochastic_rounding, + row_scaled_nvfp4=self.recipe.row_scaled_activation and idx % 3 != 1, ) return [_make_quantizer(idx) for idx in range(self.num_quantizers)] @@ -1389,6 +1390,7 @@ def _make_quantizer(idx: int) -> NVFP4Quantizer: with_post_rht_amax=self.recipe.fp4_quant_bwd_grad.random_hadamard_transform, with_2d_quantization=self.recipe.fp4_quant_bwd_grad.fp4_2d_quantization, stochastic_rounding=self.recipe.fp4_quant_bwd_grad.stochastic_rounding, + row_scaled_nvfp4=False, ) for _ in range(self.num_quantizers) ] diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index ab0c7484f..f28f972b5 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -92,6 +92,7 @@ def __new__( requires_grad: bool = False, stride: Optional[List[int]] = None, with_gemm_swizzled_scales: bool = False, + row_scaled_nvfp4: bool = False, ): if ( shapes is not None @@ -164,6 +165,7 @@ def __new__( scale_inv_offsets=scale_inv_offsets, columnwise_scale_inv_offsets=columnwise_scale_inv_offsets, with_gemm_swizzled_scales=with_gemm_swizzled_scales, + row_scaled_nvfp4=row_scaled_nvfp4, ) return instance @@ -195,6 +197,7 @@ def copy_grouped_storage_metadata(dst: GroupedTensor, src: GroupedTensor) -> Non dst.logical_shape = src.logical_shape dst.quantized_tensors = src.quantized_tensors dst._with_gemm_swizzled_scales = src._with_gemm_swizzled_scales + dst.row_scaled_nvfp4 = src.row_scaled_nvfp4 def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor: """Create a wrapper of the same type and tensor metadata as src.""" diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 65678aa34..285a7f030 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -128,6 +128,9 @@ class NVFP4Quantizer(Quantizer): """Stochastic rounding, only applicable for gradients.""" stochastic_rounding: bool + """Whether emitted NVFP4 tensors store one FP32 amax per row.""" + row_scaled_nvfp4: bool + """RHT matrix random sign mask""" rht_matrix_random_sign_mask_t: int rht_matrix: torch.Tensor @@ -143,6 +146,7 @@ def __init__( with_post_rht_amax: bool = False, with_2d_quantization: bool = False, stochastic_rounding: bool = False, + row_scaled_nvfp4: bool = False, with_random_sign_mask: bool = True, ) -> None: super().__init__(rowwise=rowwise, columnwise=columnwise) @@ -153,6 +157,7 @@ def __init__( self.amax_reduction_group = amax_reduction_group self.with_2d_quantization = with_2d_quantization self.stochastic_rounding = stochastic_rounding + self.row_scaled_nvfp4 = row_scaled_nvfp4 self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( with_random_sign_mask, torch.cuda.current_device() ) @@ -198,6 +203,7 @@ def copy(self) -> NVFP4Quantizer: with_post_rht_amax=self.with_post_rht_amax, with_2d_quantization=self.with_2d_quantization, stochastic_rounding=self.stochastic_rounding, + row_scaled_nvfp4=self.row_scaled_nvfp4, ) quantizer.internal = self.internal quantizer.optimize_for_gemm = self.optimize_for_gemm @@ -212,6 +218,8 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: def is_quantizable(self, inp: torch.Tensor) -> bool: """Returns whether or not given inp can be quantized""" + if self.row_scaled_nvfp4: + return False if inp.ndim < 2: return False if inp.shape[-1] % NVFP4_BLOCK_SCALING_SIZE != 0: @@ -313,6 +321,11 @@ def make_empty( f"Incorrect shape {shape} for NVFP4. Tensor dims must be divisible by" f" {NVFP4_BLOCK_SCALING_SIZE}" ) + if self.row_scaled_nvfp4: + if not self.rowwise_usage: + raise ValueError("Row-scaled NVFP4 quantization requires rowwise usage.") + if self.columnwise_usage: + raise ValueError("Row-scaled NVFP4 quantization does not support columnwise usage.") # Allocate FP4 data data = None @@ -329,8 +342,11 @@ def make_empty( scale_inv = torch.empty( scale_shape, dtype=torch.uint8, device=device, pin_memory=pin_memory ) - # Allocate per tensor scale inverse. FP32 format. - amax_rowwise = torch.zeros(1, dtype=torch.float32, device=device, pin_memory=pin_memory) + # Allocate global amax metadata. Row-scaled NVFP4 stores one value per row. + amax_rows = flat_first_dim if self.row_scaled_nvfp4 else 1 + amax_rowwise = torch.zeros( + amax_rows, dtype=torch.float32, device=device, pin_memory=pin_memory + ) # Allocate FP8 data transpose if needed columnwise_data = None @@ -371,6 +387,7 @@ def make_empty( quantizer=self, requires_grad=requires_grad, with_gemm_swizzled_scales=False, + row_scaled_nvfp4=self.row_scaled_nvfp4, ) def calibrate(self, tensor: torch.Tensor) -> None: @@ -431,6 +448,7 @@ def __new__( fp4_dtype: TE_DType, quantizer: Quantizer, with_gemm_swizzled_scales: bool, + row_scaled_nvfp4: bool = False, **kwargs, ): instance = super().__new__( @@ -445,6 +463,7 @@ def __new__( quantizer, with_gemm_swizzled_scales, *args, + row_scaled_nvfp4=row_scaled_nvfp4, **kwargs, ) return instance diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index 485b32328..ac56d334b 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -72,6 +72,7 @@ def _initialize_storage_fields( requires_grad: bool = False, stride: Optional[List[int]] = None, with_gemm_swizzled_scales: bool = False, + row_scaled_nvfp4: bool = False, ) -> None: """ Initialize a GroupedTensor. @@ -147,6 +148,7 @@ def _initialize_storage_fields( # Used as a convenience. instance.quantized_tensors = None instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales + instance.row_scaled_nvfp4 = row_scaled_nvfp4 def __new__( cls, @@ -172,6 +174,7 @@ def __new__( requires_grad: bool = False, stride: Optional[List[int]] = None, with_gemm_swizzled_scales: bool = False, + row_scaled_nvfp4: bool = False, ): instance = object.__new__(cls) cls._initialize_storage_fields( @@ -197,6 +200,7 @@ def __new__( requires_grad=requires_grad, stride=stride, with_gemm_swizzled_scales=with_gemm_swizzled_scales, + row_scaled_nvfp4=row_scaled_nvfp4, ) return instance @@ -371,6 +375,7 @@ def clear(self) -> None: self.columnwise_scale_inv_offsets = None self.tensor_shapes = [] self.fake_dtype = torch.float32 + self.row_scaled_nvfp4 = False def __repr__(self) -> str: """String representation of the GroupedTensorStorage.""" @@ -539,6 +544,7 @@ def copy(self) -> "GroupedTensorStorage": scale_inv_offsets=self.scale_inv_offsets, columnwise_scale_inv_offsets=self.columnwise_scale_inv_offsets, with_gemm_swizzled_scales=self._with_gemm_swizzled_scales, + row_scaled_nvfp4=self.row_scaled_nvfp4, ) @staticmethod @@ -649,6 +655,7 @@ def make_grouped_tensor( scale = None scale_inv_offsets = None columnwise_scale_inv_offsets = None + row_scaled_nvfp4 = False if no_quantization: assert dtype is not None, "dtype must be provided for unquantized GroupedTensor" if rowwise_usage: @@ -707,6 +714,19 @@ def make_grouped_tensor( # Amax buffer for delayed scaling - one per tensor amax = torch.empty(num_tensors, dtype=torch.float32, device=device) elif quantizer._get_compatible_recipe().nvfp4(): + row_scaled_nvfp4 = quantizer.row_scaled_nvfp4 + if row_scaled_nvfp4: + if not rowwise_usage: + raise ValueError( + "Row-scaled NVFP4 grouped quantization requires rowwise usage." + ) + if columnwise_usage: + raise ValueError( + "Row-scaled NVFP4 grouped quantization does not support columnwise usage." + ) + total_amax_elements = ( + sum(math.prod(s[:-1]) for s in shape) if row_scaled_nvfp4 else num_tensors + ) if rowwise_usage: # Allocate rowwise data buffer (1D flattened, uint8, but FP4 packs 2 values per byte) @@ -720,8 +740,7 @@ def make_grouped_tensor( total_scale_elements += math.prod(scale_inv_shape) scale_inv_offsets.append(total_scale_elements) scale_inv = torch.empty(total_scale_elements, dtype=torch.uint8, device=device) - # Amax buffer - one per tensor - amax = torch.empty(num_tensors, dtype=torch.float32, device=device) + amax = torch.empty(total_amax_elements, dtype=torch.float32, device=device) if columnwise_usage: # Allocate columnwise data buffer (1D flattened, uint8, FP4 packed) @@ -738,7 +757,6 @@ def make_grouped_tensor( columnwise_scale_inv = torch.empty( total_columnwise_scale_elements, dtype=torch.uint8, device=device ) - # Columnwise amax buffer - one per tensor columnwise_amax = torch.empty(num_tensors, dtype=torch.float32, device=device) elif quantizer._get_compatible_recipe().float8_block_scaling(): if rowwise_usage: @@ -824,6 +842,7 @@ def make_grouped_tensor( with_gemm_swizzled_scales=( quantizer.optimize_for_gemm if quantizer is not None else False ), + row_scaled_nvfp4=row_scaled_nvfp4, ) grouped_tensor.quantized_tensors = grouped_tensor.split_into_quantized_tensors() return grouped_tensor @@ -936,6 +955,14 @@ def split_into_quantized_tensors( cum += math.prod(scale_shape) columnwise_scale_inv_offsets.append(cum) self.columnwise_scale_inv_offsets = columnwise_scale_inv_offsets + nvfp4_rowwise_amax_offsets = None + row_scaled_nvfp4 = self.row_scaled_nvfp4 + if recipe.nvfp4() and row_scaled_nvfp4: + cum = 0 + nvfp4_rowwise_amax_offsets = [0] + for i in range(self.num_tensors): + cum += math.prod(self.tensor_shapes[i][:-1]) + nvfp4_rowwise_amax_offsets.append(cum) for i in range(self.num_tensors): quantizer = self.quantizer @@ -1128,9 +1155,13 @@ def split_into_quantized_tensors( cscale_shape ) - # Extract amax - one per tensor if self.amax is not None: - amax_rowwise = self.amax[i : i + 1] + if nvfp4_rowwise_amax_offsets is not None: + amax_start = nvfp4_rowwise_amax_offsets[i] + amax_end = nvfp4_rowwise_amax_offsets[i + 1] + amax_rowwise = self.amax[amax_start:amax_end] + else: + amax_rowwise = self.amax[i : i + 1] if self.columnwise_amax is not None: amax_columnwise = self.columnwise_amax[i : i + 1] @@ -1152,6 +1183,7 @@ def split_into_quantized_tensors( fp4_dtype=quantizer.dtype, quantizer=quantizer, with_gemm_swizzled_scales=quantizer.optimize_for_gemm, + row_scaled_nvfp4=row_scaled_nvfp4, ) result.append(tensor) diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 70699ad71..e51acb71e 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -97,6 +97,8 @@ class NVFP4TensorStorage(QuantizedTensorStorage): # Whether scaling factors are in the swizzled format expected by # GEMM _with_gemm_swizzled_scales: bool + # Whether this NVFP4 tensor uses row-scaled amax metadata + _row_scaled_nvfp4: bool def __new__( cls, @@ -111,6 +113,7 @@ def __new__( with_gemm_swizzled_scales: bool, *args, fake_dtype: Optional[torch.dtype] = None, + row_scaled_nvfp4: bool = False, **kwargs, ): if cls is NVFP4TensorStorage: @@ -128,6 +131,7 @@ def __new__( instance._amax_rowwise = amax_rowwise instance._amax_columnwise = amax_columnwise instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales + instance._row_scaled_nvfp4 = row_scaled_nvfp4 return instance @@ -152,6 +156,8 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None: raise RuntimeError("FP4 dtype mismatch in copy_from_storage") if self._with_gemm_swizzled_scales != src._with_gemm_swizzled_scales: raise RuntimeError("Scale layout mismatch in copy_from_storage") + if self._row_scaled_nvfp4 != src._row_scaled_nvfp4: + raise RuntimeError("Rowwise amax scaling mode mismatch in copy_from_storage") def _copy_optional(dst: Optional[torch.Tensor], src_tensor: Optional[torch.Tensor]): if dst is not None and src_tensor is not None: @@ -176,6 +182,7 @@ def get_metadata(self) -> Dict[str, Any]: "fp4_dtype": self._fp4_dtype, "quantizer": self._quantizer, "with_gemm_swizzled_scales": self._with_gemm_swizzled_scales, + "row_scaled_nvfp4": self._row_scaled_nvfp4, "fake_dtype": self._dtype, } @@ -308,6 +315,7 @@ def view(self, shape: torch.Size): quantizer=self._quantizer, fp4_dtype=self._fp4_dtype, with_gemm_swizzled_scales=self._with_gemm_swizzled_scales, + row_scaled_nvfp4=self._row_scaled_nvfp4, fake_dtype=self._dtype, ) From b1b302613f2e1f1e5738c2e8363ce5e6452ab241 Mon Sep 17 00:00:00 2001 From: Carlos Gomes Date: Fri, 8 May 2026 23:08:37 +0200 Subject: [PATCH 043/180] guard fuser grad checks on non-leaf nodes (#2919) * guard fuser grad checks on non-leaf nodes Signed-off-by: CarlosGomes98 * rely on set_output_requires_grad flag, update docstring Signed-off-by: CarlosGomes98 * make code clearer Signed-off-by: CarlosGomes98 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: CarlosGomes98 * rely on set_output_requires_grad flag, update docstring Signed-off-by: CarlosGomes98 * make code clearer Signed-off-by: CarlosGomes98 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Revert cudnn-frontend submodule bump Signed-off-by: CarlosGomes98 Made-with: Cursor * Tweak comment Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --------- Signed-off-by: CarlosGomes98 Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- transformer_engine/pytorch/ops/fuser.py | 28 ++++++++++++++++--------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/transformer_engine/pytorch/ops/fuser.py b/transformer_engine/pytorch/ops/fuser.py index a3c7e1bac..5283af814 100644 --- a/transformer_engine/pytorch/ops/fuser.py +++ b/transformer_engine/pytorch/ops/fuser.py @@ -65,6 +65,7 @@ def forward( input_: torch.Tensor, fuser: OperationFuser, basic_op_kwargs: list[dict[str, Any]], + set_output_requires_grad: bool, *params_and_extra_inputs: torch.Tensor, ) -> torch.Tensor | tuple[torch.Tensor, ...]: """Forward pass @@ -79,6 +80,8 @@ def forward( Container for the pipeline of operations to run basic_op_kwargs: list of dict Keyword arguments to BasicOperation + set_output_requires_grad: bool + Whether to set ``requires_grad`` flags on returned tensors *params_and_extra_inputs: torch.Tensor Other tensor inputs to include in autograd graph. Consists of parameter tensors, followed by extra operation inputs. @@ -138,7 +141,8 @@ def forward( ) for idx, ys in zip(basic_op_idxs, fused_op_extra_outputs): for y in ys: - y.requires_grad_(idx >= fuser.first_op_requiring_backward) + if set_output_requires_grad: + y.requires_grad_(idx >= fuser.first_op_requiring_backward) extra_outputs[idx] = ys # Flatten list of extra outputs @@ -190,7 +194,8 @@ def forward( for tensor in [x] + extra_outputs_flat: tensor._do_not_clear = True - x.requires_grad_(fuser.first_op_requiring_backward < fuser._num_basic_ops) + if set_output_requires_grad: + x.requires_grad_(fuser.first_op_requiring_backward < fuser._num_basic_ops) if extra_outputs_flat: return x, *extra_outputs_flat @@ -293,6 +298,7 @@ def backward( dx, # input_ None, # fuser None, # basic_op_kwargs + None, # set_output_requires_grad *grad_params_flat, *grad_extra_inputs_flat, ) @@ -501,20 +507,22 @@ def __call__( op.pre_fuser_forward(requires_grad=idx >= self.first_op_requiring_backward) # Fuser forward pass - if is_grad_enabled: - forward_func = _OperationFuserAutogradFunction.apply - args = [] - else: - forward_func = _OperationFuserAutogradFunction.forward - args = [None] - args += ( + # Note: We call forward directly when is_grad_enabled=False, + # which can expose non-leaf tensors to the inner ops. Avoid + # problems in this case by passing set_output_requires_grad=False. + args = ( input, self, basic_op_kwargs, + is_grad_enabled, # set_output_requires_grad *self._flat_basic_op_params, *extra_inputs, ) - return forward_func(*args) + + if not is_grad_enabled: + return _OperationFuserAutogradFunction.forward(None, *args) + + return _OperationFuserAutogradFunction.apply(*args) def register_forward_fusion( From 56ff4c6331d0cd8dd1ea86c9dff38ea06b6599ed Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Fri, 8 May 2026 17:16:26 -0700 Subject: [PATCH 044/180] [PyTorch] Remove internal PyTorch testing helper (#2969) * Remove internal PyTorch testing helper Signed-off-by: Tim Moon * Review suggestion from @greptile-apps 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> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Tim Moon Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/test_fused_optimizer.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/pytorch/test_fused_optimizer.py b/tests/pytorch/test_fused_optimizer.py index e72cad9db..a2863cba9 100644 --- a/tests/pytorch/test_fused_optimizer.py +++ b/tests/pytorch/test_fused_optimizer.py @@ -8,7 +8,6 @@ import pytest import torch from torch import nn -from torch.testing._internal.common_device_type import largeTensorTest import transformer_engine.pytorch as te from transformer_engine.common.recipe import DelayedScaling, MXFP8BlockScaling, Float8BlockScaling from transformer_engine.pytorch import MultiheadAttention, quantized_model_init, is_bf16_available @@ -1053,8 +1052,13 @@ def test_native(self): self.model_.load_state_dict(copy.deepcopy(self.model.state_dict())) - @largeTensorTest("60GB", "cuda") def test_large_tensor(self): + import gc + + gc.collect() + torch.cuda.empty_cache() + if torch.cuda.memory.mem_get_info()[0] < 60 * 1024**3: + pytest.skip("Insufficient available memory") t = torch.zeros(2359332864, dtype=torch.half, device="cuda") t2 = torch.zeros(2359332864, dtype=torch.half, device="cuda") grad = torch.randn_like(t) From 0e289534985c95192d2e48e6d2447f7c53feff2f Mon Sep 17 00:00:00 2001 From: Zhang Haitao Date: Sat, 9 May 2026 08:44:08 +0800 Subject: [PATCH 045/180] Fix nvfp4 convert_and_update_tensor shape check (#2670) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix nvfp4 convert_and_update_tensor shape check Signed-off-by: 乙划 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add headers and check 2D shapes Signed-off-by: 乙划 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Przemyslaw Tredak * add unittest and doctring Signed-off-by: 乙划 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Fix NVFP4 shape check for N-D tensors in convert_and_update_tensor Introduce get_2d_dims() in common.h/cpp to flatten an N-D shape to 2D dims (flat_first, flat_last), replacing the ad-hoc compressShapeTo2D helper from the contributor PR. The helper takes NVTEShape as its core argument (stack-allocated) with a header-only vector overload, and supports a transpose flag for the shape[1:] flattening direction. Use get_2d_dims in NVFP4Quantizer::convert_and_update_tensor to compare row-wise and column-wise shapes under 2D equivalence — fixing a false mismatch when the logical shape is 3D (columnwise data is always stored 2D). Also restructure the if-block to treat row-wise data as the ground truth when present. Fixes #2607 Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Tim Moon * [PyTorch] Add test for updating N-D quantized tensors via copy_ Replaces the NVFP4-only test_nvfp4_3d_shape_quantization in test_nvfp4_quantize_exact.py with a broader test_update_nd_tensor in TestQuantizedTensor that covers all quantization formats. The test constructs an N-D quantized tensor, updates it with copy_, and checks both shape preservation and numerical accuracy. The "nvfp4_2d" variant is appended to the parametrize list inline to cover both NVFP4 quantization modes without affecting the shared _quantization_list. Also adds "fp8_blockwise" to quantization_tols in utils.py. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Tim Moon * [PyTorch] Propagate get_2d_dims helper across C++ extensions Replace ad-hoc loops computing flat_first_dim/flat_last_dim and equivalent product(shape)/shape.back() patterns in quantizer.cpp, cast.cpp, gemm.cpp, normalization.cpp, swizzle.cpp, and transpose.cpp. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Tim Moon * [PyTorch] Drop intermediate variable in cast.cpp get_2d_dims call Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Check that tensor shape is not too large Suggestion from @greptile-apps. Signed-off-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 --------- Signed-off-by: 乙划 Signed-off-by: Przemyslaw Tredak Signed-off-by: Tim Moon Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: 乙划 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Przemyslaw Tredak Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Tim Moon Co-authored-by: Claude Sonnet 4.6 Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- tests/pytorch/test_quantized_tensor.py | 58 ++++++++++++++++ tests/pytorch/utils.py | 1 + transformer_engine/pytorch/csrc/common.cpp | 14 ++++ transformer_engine/pytorch/csrc/common.h | 16 +++++ .../pytorch/csrc/extensions/cast.cpp | 10 +-- .../pytorch/csrc/extensions/gemm.cpp | 6 +- .../pytorch/csrc/extensions/normalization.cpp | 6 +- .../pytorch/csrc/extensions/swizzle.cpp | 18 +---- .../pytorch/csrc/extensions/transpose.cpp | 3 +- transformer_engine/pytorch/csrc/quantizer.cpp | 69 +++++-------------- 10 files changed, 116 insertions(+), 85 deletions(-) diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index 23ce93319..526045e43 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -28,6 +28,7 @@ import transformer_engine_torch as tex from references.ref_per_tensor_cs import ref_per_tensor_cs_cast +from utils import assert_close, quantization_tols # PyTorch tensor dtypes _dtypes: List[torch.dtype] = [torch.float32, torch.float16, torch.bfloat16] @@ -702,6 +703,63 @@ def test_shape_with_none_data( f"after setting data to None on {type(x_test).__name__}" ) + @pytest.mark.parametrize( + "quantization", + _quantization_list + (["nvfp4_2d"] if nvfp4_available else []), + ) + def test_update_nd_tensor( + self, + *, + quantization: str, + shape: Iterable[int] = (32, 4, 128), + dtype: torch.dtype = torch.bfloat16, + device: str = "cuda", + ) -> None: + """Check that an N-D quantized tensor can be updated.""" + + # Construct quantizer + if quantization == "fp8": + quantizer = Float8Quantizer( + scale=torch.ones(1, dtype=torch.float32, device=device).squeeze(), + amax=torch.zeros(1, dtype=torch.float32, device=device), + fp8_dtype=tex.DType.kFloat8E4M3, + ) + elif quantization == "fp8_blockwise": + quantizer = Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + force_pow_2_scales=True, + amax_epsilon=0.0, + block_scaling_dim=1, + ) + elif quantization == "mxfp8": + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) + elif quantization in ("nvfp4", "nvfp4_2d"): + quantizer = NVFP4Quantizer( + rowwise=True, + columnwise=True, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=(quantization == "nvfp4_2d"), + ) + quantization = "nvfp4" + else: + raise ValueError(f"Unknown quantization: {quantization}") + + # Construct quantized tensor + x = torch.randn(list(shape), dtype=dtype, device=device) + q_x = quantizer(x) + + # Update tensor + x_new = torch.randn(list(shape), dtype=dtype, device=device) + q_x.copy_(x_new) + + # Check results + assert q_x.shape == torch.Size(shape) + tols = quantization_tols(quantization) + assert_close(q_x, x_new, **tols) + @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) class TestMXFP8Tensor: diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 32e44be2a..2ee18aaf5 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -113,6 +113,7 @@ def quantization_tols(name: str) -> dict[str, float]: "fp8", "fp8_delayed_scaling", "fp8_current_scaling", + "fp8_blockwise", "mxfp8", "mxfp8_block_scaling", ): diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp index b06f6f561..66bb2dc40 100644 --- a/transformer_engine/pytorch/csrc/common.cpp +++ b/transformer_engine/pytorch/csrc/common.cpp @@ -26,6 +26,20 @@ std::vector convert_shape_back_from_fp4(const std::vector& shape return ret; } +std::array get_2d_dims(NVTEShape shape, bool transpose) { + if (!transpose) { + size_t flat_first = 1; + for (size_t i = 0; i + 1 < shape.ndim; ++i) flat_first *= shape.data[i]; + const size_t flat_last = shape.ndim > 0 ? shape.data[shape.ndim - 1] : 1; + return {flat_first, flat_last}; + } else { + const size_t flat_first = shape.ndim > 0 ? shape.data[0] : 1; + size_t flat_last = 1; + for (size_t i = 1; i < shape.ndim; ++i) flat_last *= shape.data[i]; + return {flat_first, flat_last}; + } +} + std::vector getTensorShape(const at::Tensor& t) { std::vector shape; for (auto s : t.sizes()) { diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 8f5b8294e..35a459351 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -45,6 +45,7 @@ #include #include +#include #include #include #include @@ -523,6 +524,21 @@ NVTEShape convertTorchShape(const c10::IntArrayRef torch_shape); std::vector convert_shape_back_from_fp4(const std::vector& shape, bool transpose); +// Flatten an N-D shape to 2D: {product(shape[:-1]), shape[-1]}. +// With transpose=true: {shape[0], product(shape[1:])}. +std::array get_2d_dims(NVTEShape shape, bool transpose = false); + +template +inline std::array get_2d_dims(const std::vector& shape, bool transpose = false) { + NVTEShape s{}; + s.ndim = shape.size(); + constexpr size_t max_ndim = sizeof(s.data) / sizeof(size_t); + NVTE_CHECK(s.ndim <= max_ndim, "Shape has too many dimensions (got ", s.ndim, ", max ", max_ndim, + ")."); + for (size_t i = 0; i < shape.size(); ++i) s.data[i] = static_cast(shape[i]); + return get_2d_dims(s, transpose); +} + // unpack the PhiloxCudaState into CUDA tensor void philox_unpack(at::PhiloxCudaState arg, int64_t* rng_state_ptr); diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 9e1f381bf..00f4383ab 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -1243,14 +1243,8 @@ void split_quantize_nvfp4_impl_with_rht_helper(const TensorWrapper &input, // Create a wrapper for the columnwise output, as the rowwise output. Input is in transposed layout. TensorWrapper out_transpose(output_list[i].scaling_mode()); if (!is_empty_split) { - auto colwise_data_shape = out_columnwise_data.shape; - std::vector colwise_data_shape_2d; - colwise_data_shape_2d.push_back(colwise_data_shape.data[0]); - size_t last_dim = 1; - for (size_t j = 1; j < colwise_data_shape.ndim; ++j) { - last_dim *= colwise_data_shape.data[j]; - } - colwise_data_shape_2d.push_back(last_dim); + auto [cw_first, cw_last] = get_2d_dims(out_columnwise_data.shape, true); + std::vector colwise_data_shape_2d = {cw_first, cw_last}; out_transpose.set_rowwise_data(out_columnwise_data.data_ptr, static_cast(out_columnwise_data.dtype), diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index 427eb7934..9cb1fb7f5 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -41,10 +41,8 @@ bool is_low_precision(const DType type) { std::vector getGemmOutputShape(const NVTEShape& A_shape, const bool transa, const NVTEShape& B_shape, const bool transb) { // Flatten outer dims to get 2D matrices - const size_t A0 = A_shape.ndim > 0 ? product(A_shape, 0, A_shape.ndim - 1) : 1; - const size_t A1 = A_shape.ndim > 0 ? A_shape.data[A_shape.ndim - 1] : 1; - const size_t B0 = B_shape.ndim > 0 ? product(B_shape, 0, B_shape.ndim - 1) : 1; - const size_t B1 = B_shape.ndim > 0 ? B_shape.data[B_shape.ndim - 1] : 1; + const auto [A0, A1] = get_2d_dims(A_shape); + const auto [B0, B1] = get_2d_dims(B_shape); // Check matrix dims NVTE_CHECK((transa ? A1 : A0) == (transb ? B0 : B1), "Invalid matrix dimensions for GEMM (A=(", diff --git a/transformer_engine/pytorch/csrc/extensions/normalization.cpp b/transformer_engine/pytorch/csrc/extensions/normalization.cpp index 4887b59c2..c3dec944e 100644 --- a/transformer_engine/pytorch/csrc/extensions/normalization.cpp +++ b/transformer_engine/pytorch/csrc/extensions/normalization.cpp @@ -80,8 +80,7 @@ std::vector layernorm_fwd(py::handle input, py::handle weight, Maybe // Tensor dimensions const auto shape = nvte_shape_to_vector(input_nvte.shape()); - const auto outer_size = product(shape) / shape.back(); - const auto inner_size = shape.back(); + const auto [outer_size, inner_size] = get_2d_dims(shape); // Tensors to save for backward pass at::Tensor mu_py = at::empty({static_cast(outer_size)}, at::CUDA(at::kFloat)); @@ -320,8 +319,7 @@ std::vector rmsnorm_fwd(const py::handle &input, const py::handle &w // Tensor dimensions const auto shape = nvte_shape_to_vector(input_nvte.shape()); - const auto outer_size = product(shape) / shape.back(); - const auto inner_size = shape.back(); + const auto [outer_size, inner_size] = get_2d_dims(shape); // Tensors to save for backward pass at::Tensor rsigma_py = at::empty({static_cast(outer_size)}, at::CUDA(at::kFloat)); diff --git a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp index 7f7f8a435..193aed29e 100644 --- a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp +++ b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp @@ -301,22 +301,8 @@ at::Tensor convert_block_scaling_to_mxfp8_tensor(transformer_engine::TensorWrapp "Input tensor must be a block scaling tensor"); // Get tensor data - NVTEBasicTensor data; - size_t data_flat_first_dim = 1; - size_t data_flat_last_dim = 1; - if (rowwise) { - data = input.get_rowwise_data(); - for (size_t i = 0; i < data.shape.ndim - 1; ++i) { - data_flat_first_dim *= data.shape.data[i]; - } - data_flat_last_dim = data.shape.data[data.shape.ndim - 1]; - } else { - data = input.get_columnwise_data(); - data_flat_first_dim = data.shape.data[0]; - for (size_t i = 1; i < data.shape.ndim; ++i) { - data_flat_last_dim *= data.shape.data[i]; - } - } + NVTEBasicTensor data = rowwise ? input.get_rowwise_data() : input.get_columnwise_data(); + const auto [data_flat_first_dim, data_flat_last_dim] = get_2d_dims(data.shape, !rowwise); NVTEShape data_shape{}; data_shape.data[0] = data_flat_first_dim; data_shape.data[1] = data_flat_last_dim; diff --git a/transformer_engine/pytorch/csrc/extensions/transpose.cpp b/transformer_engine/pytorch/csrc/extensions/transpose.cpp index aaa27a104..031897819 100644 --- a/transformer_engine/pytorch/csrc/extensions/transpose.cpp +++ b/transformer_engine/pytorch/csrc/extensions/transpose.cpp @@ -29,8 +29,7 @@ at::Tensor fp8_transpose(at::Tensor input, DType otype, std::optional 0 ? product(shape) / shape.back() : 1; - const size_t N = shape.size() > 0 ? shape.back() : 1; + const auto [M, N] = get_2d_dims(shape); // Output tensor at::Tensor out; diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 8f2de325a..2b29f260e 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1342,13 +1342,7 @@ std::pair MXFP8Quantizer::create_tensor(const std::ve // Tensor dimensions const std::vector shape_int64(shape.begin(), shape.end()); - size_t flat_first_dim = 1; - if (shape.size() > 0) { - for (size_t i = 0; i < shape.size() - 1; ++i) { - flat_first_dim *= shape[i]; - } - } - const size_t flat_last_dim = shape.size() > 0 ? shape.back() : 1; + const auto [flat_first_dim, flat_last_dim] = get_2d_dims(shape); NVTE_CHECK(flat_first_dim % MXFP8_BLOCK_SIZE == 0 && flat_last_dim % MXFP8_BLOCK_SIZE == 0, "MXFP8 requires tensor dims that are divisible by ", MXFP8_BLOCK_SIZE, " (got shape=", shape, ")"); @@ -1736,13 +1730,7 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve // Tensor dimensions const std::vector shape_int64(shape.begin(), shape.end()); - size_t flat_first_dim = 1; - if (shape.size() > 0) { - for (size_t i = 0; i < shape.size() - 1; ++i) { - flat_first_dim *= shape[i]; - } - } - const size_t flat_last_dim = shape.size() > 0 ? shape.back() : 1; + const auto [flat_first_dim, flat_last_dim] = get_2d_dims(shape); NVTE_CHECK(flat_first_dim % NVFP4_BLOCK_SIZE == 0, "First dim for NVFP4 must be divisible by ", NVFP4_BLOCK_SIZE, " (got shape=", shape, ")"); NVTE_CHECK(flat_last_dim % NVFP4_BLOCK_SIZE == 0, @@ -2040,24 +2028,18 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( // Tensor dimensions, shape means original shape std::vector shape; - if (columnwise_data) { - shape = convert_shape_back_from_fp4(getTensorShape(*columnwise_data), true); - if (rowwise_data) { - auto expected_shape = convert_shape_back_from_fp4(getTensorShape(*rowwise_data), false); - NVTE_CHECK(shape == expected_shape, "NVFP4 row-wise data (shape=", expected_shape, - ") and column-wise data (shape=", shape, ") do not match"); - } - } else { // Already checked columnwise_data_tensor == true + if (rowwise_data) { shape = convert_shape_back_from_fp4(getTensorShape(*rowwise_data), false); - } - - size_t flat_first_dim = 1; - if (shape.size() > 0) { - for (size_t i = 0; i < shape.size() - 1; ++i) { - flat_first_dim *= shape[i]; + if (columnwise_data) { + auto col_shape = convert_shape_back_from_fp4(getTensorShape(*columnwise_data), true); + NVTE_CHECK(get_2d_dims(shape) == get_2d_dims(col_shape), "NVFP4 row-wise data (shape=", shape, + ") and column-wise data (shape=", col_shape, ") do not match"); } + } else { + shape = convert_shape_back_from_fp4(getTensorShape(*columnwise_data), true); } - const size_t flat_last_dim = shape.size() > 0 ? shape.back() : 1; + + const auto [flat_first_dim, flat_last_dim] = get_2d_dims(shape); const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 quantization requires rowwise usage."); @@ -2205,24 +2187,16 @@ void NVFP4Quantizer::quantize_with_rht_unfused_helper( // NOTE: should already be populated. auto out_columnwise_amax = out.get_columnwise_amax(); + // Flatten column-wise data shape to 2D to avoid problems when + // converting between FP4 tensor shape and byte tensor shape + // (involves dividing last dim by 2). + auto [flat_first_dim, flat_last_dim] = get_2d_dims(out_columnwise_data.shape, true); + std::vector colwise_data_shape_2d = {flat_first_dim, flat_last_dim}; + // Create a wrapper for the columnwise output, as the rowwise output. // The reason is due to the input `rht_output_t` is already in the transposed layout. // Thus, we only need a rowwise quantization to generate the columnwise output. TensorWrapper out_transpose(out.scaling_mode()); - // Note: since we are faking columnwise tensor into rowwise, the flat first dim check will fail - // need to convert the shape to 2D here - auto colwise_data_shape = out_columnwise_data.shape; - std::vector colwise_data_shape_2d; - // shape could be [512, 32, 64], that's actually 512, 32, 128 because 2 FP4 take 1 byte - // the 2D shape should be [512, 32*128], but columnwise data shape expect last dim to be halved again - // so the multiple 2 get cancelled out - colwise_data_shape_2d.push_back(colwise_data_shape.data[0]); - size_t last_dim = 1; - for (size_t i = 1; i < colwise_data_shape.ndim; ++i) { - last_dim *= colwise_data_shape.data[i]; - } - colwise_data_shape_2d.push_back(last_dim); - out_transpose.set_rowwise_data(out_columnwise_data.data_ptr, static_cast(out_columnwise_data.dtype), colwise_data_shape_2d); @@ -2234,7 +2208,6 @@ void NVFP4Quantizer::quantize_with_rht_unfused_helper( out_columnwise_amax.shape); // Invoking fallback RHT kernel unfused. - NVTE_SCOPED_GIL_RELEASE({ // Perform the RHT(input.t), and write to rht_output_cpp.columnwise. nvte_hadamard_transform(input.data(), rht_output_t_cpp.data(), 0, @@ -2483,13 +2456,7 @@ void NVFP4Quantizer::quantize_with_amax(TensorWrapper& input, TensorWrapper& out std::vector NVFP4Quantizer::get_scale_shape(const std::vector& shape, bool columnwise) const { - size_t numel = 1; - for (auto s : shape) { - numel *= s; - } - - auto last_dim = shape.back(); - auto flat_first_dim = numel / last_dim; + const auto [flat_first_dim, last_dim] = get_2d_dims(shape); NVTE_CHECK(last_dim % NVFP4_BLOCK_SIZE == 0, "Last dim for NVFP4 must be divisible by ", NVFP4_BLOCK_SIZE, " (got dim=", last_dim, ")"); From 25934ac33702dc44d26b6f3ee514d57a9e08dcfa Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Mon, 11 May 2026 02:33:48 -0700 Subject: [PATCH 046/180] Refactor tensor class in C++ unit tests (#2962) * Fix bug in NVFP4 quantize test where we set scale instead of amax Refactor test tensor wrapper by removing recipe-specific logic whenever possible. Signed-off-by: Tim Moon * Only get fp32 scale when tensor is expected to have fp32 scale Signed-off-by: Tim Moon * Create dedicated class for managing GPU/CPU buffers Signed-off-by: Tim Moon * Fix bugs in C++ test tensor infrastructure - Fix syntax error in switch case (:: -> :) - Fix double-underscore typo in variable name - Fix wrong buffer passed to set_amax_columnwise - Fix unique_ptr assignment from raw pointer (use reset()) - Remove dead duplicate NVTE_MXFP8_1D_SCALING branch in get_scales() - Rename cpu_data -> cpu_buffer to match Buffer class API - Remove const from Tensor::to_cpu/from_cpu and their callers, since both methods write to the CPU buffer Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Tim Moon * Debug compilation errors Signed-off-by: Tim Moon * Remove type check when accessing raw pointers CPU and GPU types are inconsistent, so the type checks cause too many problems. Signed-off-by: Tim Moon * Debug distributed C++ tests Also adopt review suggestions from @greptile-apps. Signed-off-by: Tim Moon * Remove unused header Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Copy-paste error Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Use shared buffer for FP8 row-wise scale-inv and col-wise scale-inv Signed-off-by: Tim Moon * Typo Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Debug merge conflicts with #2931 Also do some cleanup and improve documentation. Signed-off-by: Tim Moon * Address code review feedback - Restore amax buffer size assertion in compare_rowwise_amax - Remove set_tensor_amax alias in favor of set_amax - Extract fill_uniform_buffer helper to anonymous namespace, eliminating duplication in fill_uniform_{rowwise,columnwise}_scale_inv Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Tim Moon --------- Signed-off-by: Tim Moon Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 --- tests/cpp/operator/test_act.cu | 10 +- tests/cpp/operator/test_cast.cu | 5 +- .../cpp/operator/test_cast_current_scaling.cu | 7 +- tests/cpp/operator/test_cast_dbias.cu | 5 +- tests/cpp/operator/test_cast_dbias_dgelu.cu | 5 +- .../cpp/operator/test_cast_float8blockwise.cu | 28 +- tests/cpp/operator/test_cast_gated_swiglu.cu | 5 +- .../cpp/operator/test_cast_nvfp4_transpose.cu | 219 +++---- tests/cpp/operator/test_cast_transpose.cu | 6 +- .../cpp/operator/test_cast_transpose_dbias.cu | 5 +- .../test_cast_transpose_dbias_dgelu.cu | 5 +- .../operator/test_cast_transpose_dgeglu.cu | 5 +- tests/cpp/operator/test_dequantize_nvfp4.cu | 80 +-- .../cpp/operator/test_multi_cast_transpose.cu | 4 +- tests/cpp/operator/test_normalization.cu | 2 +- tests/cpp/operator/test_qdq.cu | 3 +- tests/cpp/test_common.cu | 560 ++++++++---------- tests/cpp/test_common.h | 319 +++++----- tests/cpp_distributed/test_comm_gemm.cu | 12 +- .../transformer_engine/transformer_engine.h | 12 +- 20 files changed, 614 insertions(+), 683 deletions(-) diff --git a/tests/cpp/operator/test_act.cu b/tests/cpp/operator/test_act.cu index ca5ccdc4c..6edc6bd63 100644 --- a/tests/cpp/operator/test_act.cu +++ b/tests/cpp/operator/test_act.cu @@ -124,6 +124,7 @@ void performTest(const size_t N, const size_t H) { fillUniform(&input); fillUniform(&ograd); setRandomScale(&output); + const float ref_scale = isFp8Type(otype) ? output.scale() : 1.0f; std::unique_ptr ref_output = std::make_unique(N*H); std::unique_ptr ref_igrad = std::make_unique(N*H); @@ -132,7 +133,7 @@ void performTest(const size_t N, const size_t H) { float ref_amax; compute_ref_act_cast(input.rowwise_cpu_dptr(), ref_output.get(), - output.scale(), &ref_amax, N, H); + ref_scale, &ref_amax, N, H); cudaDeviceSynchronize(); auto err = cudaGetLastError(); @@ -179,6 +180,7 @@ void performTestGLU(const size_t N, const size_t H) { fillUniform(&input); fillUniform(&ograd); setRandomScale(&output); + const float ref_scale = isFp8Type(otype) ? output.scale() : 1.0f; std::unique_ptr ref_output = std::make_unique(N * H); std::unique_ptr ref_igrad = std::make_unique(2 * N * H); @@ -187,7 +189,7 @@ void performTestGLU(const size_t N, const size_t H) { float ref_amax; compute_ref_glu_act_cast(input.rowwise_cpu_dptr(), ref_output.get(), - output.scale(), &ref_amax, N, H); + ref_scale, &ref_amax, N, H); cudaDeviceSynchronize(); auto err = cudaGetLastError(); @@ -197,8 +199,8 @@ void performTestGLU(const size_t N, const size_t H) { auto [atol, rtol] = getTolerances(DType::kFloat32); compareResults("amax", output.amax(), ref_amax, atol, rtol); if (output.scaling_mode() == NVTE_DELAYED_TENSOR_SCALING) { - const float ref_scale = 1.f / output.scale(); - compareResults("scale_inv", *output.rowwise_cpu_scale_inv_ptr(), ref_scale, atol, rtol); + const float ref_scale_inv = 1.f / ref_scale; + compareResults("scale_inv", *output.rowwise_cpu_scale_inv_ptr(), ref_scale_inv, atol, rtol); } } auto [atol, rtol] = getTolerances(otype); diff --git a/tests/cpp/operator/test_cast.cu b/tests/cpp/operator/test_cast.cu index 35d9dd2ef..e8f48feef 100644 --- a/tests/cpp/operator/test_cast.cu +++ b/tests/cpp/operator/test_cast.cu @@ -53,13 +53,14 @@ void performTest(const std::vector& shape) { fillUniform(&input); setRandomScale(&output_c); + const float ref_scale = isFp8Type(otype) ? output_c.scale() : 1.0f; nvte_quantize(input.data(), output_c.data(), 0); float ref_amax; compute_ref(input.rowwise_cpu_dptr(), ref_output_c.get(), - full_size, &ref_amax, output_c.scale()); + full_size, &ref_amax, ref_scale); cudaDeviceSynchronize(); auto err = cudaGetLastError(); @@ -67,7 +68,7 @@ void performTest(const std::vector& shape) { if (isFp8Type(otype)) { auto [atol_amax, rtol_amax] = getTolerances(DType::kFloat32); compareResults("amax", output_c.amax(), ref_amax, atol_amax, rtol_amax); - float ref_scale_inv = 1.f / output_c.scale(); + float ref_scale_inv = 1.f / ref_scale; compareResults("scale_inv", output_c.rowwise_scale_inv(), ref_scale_inv, atol_amax, rtol_amax); } auto [atol, rtol] = getTolerances(otype); diff --git a/tests/cpp/operator/test_cast_current_scaling.cu b/tests/cpp/operator/test_cast_current_scaling.cu index 4dd6cd2d5..7cca0d72e 100644 --- a/tests/cpp/operator/test_cast_current_scaling.cu +++ b/tests/cpp/operator/test_cast_current_scaling.cu @@ -123,6 +123,7 @@ void performTest(const std::vector& shape) { nvte_compute_amax(input.data(), output_c.data(), 0); QuantizationConfigWrapper config; nvte_compute_scale_from_amax(output_c.data(), config, 0); + // avoid atomic amax update in cuda cast kernels because of current per-tensor scaling amax_to_check = output_c.amax(); output_c.set_tensor_amax_nullptr(); @@ -130,7 +131,7 @@ void performTest(const std::vector& shape) { nvte_quantize(input.data(), output_c.data(), 0); float ref_amax; - float ref_scale; + float ref_scale = 1.0; float ref_scale_inv; if (is_out_fp8){ compute_amax_scale_ref(input.rowwise_cpu_dptr(), @@ -138,13 +139,13 @@ void performTest(const std::vector& shape) { } compute_ref(input.rowwise_cpu_dptr(), ref_output_c.get(), - full_size, nullptr, is_out_fp8 ? output_c.scale() : 1.0f ); + full_size, nullptr, ref_scale); cudaDeviceSynchronize(); auto err = cudaGetLastError(); ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); - if (isFp8Type(otype)) { + if (is_out_fp8) { auto [atol_fp32, rtol_fp32] = getTolerances(DType::kFloat32); compareResults("amax", amax_to_check, ref_amax, 0.0f, rtol_fp32); compareResults("scale", output_c.scale(), ref_scale, 0.0f, rtol_fp32); diff --git a/tests/cpp/operator/test_cast_dbias.cu b/tests/cpp/operator/test_cast_dbias.cu index 18f07153c..b7b5db48c 100644 --- a/tests/cpp/operator/test_cast_dbias.cu +++ b/tests/cpp/operator/test_cast_dbias.cu @@ -74,13 +74,14 @@ void performTest(const std::vector& shape) { fillUniform(&input); setRandomScale(&output_c); + const float ref_scale = isFp8Type(otype) ? output_c.scale() : 1.0f; std::unique_ptr ref_output_c = std::make_unique(N*H); std::unique_ptr ref_output_dbias = std::make_unique(H); CType ref_amax; compute_ref_cast_dbias(input.rowwise_cpu_dptr(), - output_c.scale(), + ref_scale, ref_output_c.get(), &ref_amax, ref_output_dbias.get(), @@ -109,7 +110,7 @@ void performTest(const std::vector& shape) { if (isFp8Type(otype)) { auto [atol_amax, rtol_amax] = getTolerances(DType::kFloat32); compareResults("amax", output_c.amax(), ref_amax, atol_amax, rtol_amax); - float ref_scale_inv = 1.f / output_c.scale(); + float ref_scale_inv = 1.f / ref_scale; compareResults("scale_inv", output_c.rowwise_scale_inv(), ref_scale_inv, atol_amax, rtol_amax); } auto [atol, rtol] = getTolerances(otype); diff --git a/tests/cpp/operator/test_cast_dbias_dgelu.cu b/tests/cpp/operator/test_cast_dbias_dgelu.cu index 8213e5665..d8b8a20e6 100644 --- a/tests/cpp/operator/test_cast_dbias_dgelu.cu +++ b/tests/cpp/operator/test_cast_dbias_dgelu.cu @@ -84,6 +84,7 @@ void performTest(const std::vector& shape) { fillUniform(&input); fillUniform(&grad); setRandomScale(&output_c); + const float ref_scale = isFp8Type(otype) ? output_c.scale() : 1.0f; std::unique_ptr ref_output_c = std::make_unique(N*H); std::unique_ptr ref_output_dbias = std::make_unique(H); @@ -91,7 +92,7 @@ void performTest(const std::vector& shape) { CType ref_amax; compute_ref_cast_dbias_dgelu(input.rowwise_cpu_dptr(), grad.rowwise_cpu_dptr(), - output_c.scale(), + ref_scale, ref_output_c.get(), &ref_amax, ref_output_dbias.get(), @@ -123,7 +124,7 @@ void performTest(const std::vector& shape) { if (isFp8Type(otype)) { auto [atol_amax, rtol_amax] = getTolerances(DType::kFloat32); compareResults("amax", output_c.amax(), ref_amax, atol_amax, rtol_amax); - float ref_scale_inv = 1.f / output_c.scale(); + float ref_scale_inv = 1.f / ref_scale; compareResults("scale_inv", output_c.rowwise_scale_inv(), ref_scale_inv, atol_amax, rtol_amax); } diff --git a/tests/cpp/operator/test_cast_float8blockwise.cu b/tests/cpp/operator/test_cast_float8blockwise.cu index 8e9da91d0..d50589ca4 100644 --- a/tests/cpp/operator/test_cast_float8blockwise.cu +++ b/tests/cpp/operator/test_cast_float8blockwise.cu @@ -524,14 +524,12 @@ TEST_P(FusedCastFloat8BlockwiseTestSuite, TestFusedCastFloat8Blockwise) { // GTEST_SKIP(); // } - DACT_FUNC_SWITCH( - Act_type, OP, - TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY( - input_type, InputType, - TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY( - output_type, OutputType, - runTestCase(processing_method, matrix_size, rowwise, colwise, - fill_case, q_opts);););); + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY( + input_type, InputType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY( + output_type, OutputType, + runTestCase(processing_method, matrix_size, rowwise, colwise, + fill_case, q_opts););); } TEST_P(FusedCastFloat8VectorwiseTestSuite, TestFusedCastFloat8Vectorwise) { @@ -581,14 +579,12 @@ TEST_P(FusedCastFloat8VectorwiseTestSuite, TestFusedCastFloat8Vectorwise) { // GTEST_SKIP(); // } - DACT_FUNC_SWITCH( - Act_type, OP, - TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY( - input_type, InputType, - TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY( - output_type, OutputType, - runTestCaseOneDimensionalBlocks( - processing_method, matrix_size, rowwise, colwise, fill_case, q_opts);););); + TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY( + input_type, InputType, + TRANSFORMER_ENGINE_TYPE_SWITCH_FP8_ONLY( + output_type, OutputType, + runTestCaseOneDimensionalBlocks( + processing_method, matrix_size, rowwise, colwise, fill_case, q_opts););); } std::string to_string(const ProcessingMethod method) { diff --git a/tests/cpp/operator/test_cast_gated_swiglu.cu b/tests/cpp/operator/test_cast_gated_swiglu.cu index 298b978f2..5298cc757 100644 --- a/tests/cpp/operator/test_cast_gated_swiglu.cu +++ b/tests/cpp/operator/test_cast_gated_swiglu.cu @@ -79,6 +79,7 @@ void performTest(const std::vector& shape) { fillUniform(&grad); fillUniform(&input); setRandomScale(&output_c); + const float ref_scale = isFp8Type(otype) ? output_c.scale() : 1.0f; std::unique_ptr ref_output_c = std::make_unique(input_size); @@ -91,7 +92,7 @@ void performTest(const std::vector& shape) { float ref_amax; compute_ref_cast_dgated_swiglu(grad.rowwise_cpu_dptr(), input.rowwise_cpu_dptr(), - output_c.scale(), + ref_scale, ref_output_c.get(), &ref_amax, rows, @@ -100,7 +101,7 @@ void performTest(const std::vector& shape) { if (isFp8Type(otype)) { auto [atol_amax, rtol_amax] = getTolerances(DType::kFloat32); compareResults("amax", output_c.amax(), ref_amax, atol_amax, rtol_amax); - float ref_scale_inv = 1.f / output_c.scale(); + float ref_scale_inv = 1.f / ref_scale; compareResults("scale_inv", output_c.rowwise_scale_inv(), ref_scale_inv, atol_amax, rtol_amax); } diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index 1f37520bc..a8f58f859 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -4,6 +4,15 @@ * See LICENSE for license information. ************************************************************************/ +#include +#include +#include +#include +#include +#include +#include +#include + #include #include #include @@ -14,7 +23,6 @@ #include #include "../test_common.h" #include "transformer_engine/transformer_engine.h" -#include using namespace transformer_engine; using namespace test; @@ -309,40 +317,24 @@ void compute_ref(float (*OP)(const float), fp4e2m1x2* output_t, fp8e4m3* scales, fp8e4m3* scales_t, - const float global_amax, + const float* amax, const size_t rows, const size_t cols, const size_t scales_stride, const size_t scales_stride_t, const bool use_fast_math, const bool use_2d_quantization = false, - std::vector *rowwise_amax = nullptr) + const bool row_scaled_nvfp4 = false) { std::vector input_t = create_transpose(input, rows, cols); + NVTE_CHECK(!(use_2d_quantization && row_scaled_nvfp4), + "2D quantization and row-scaling are not supported together."); - if (rowwise_amax != nullptr) { - rowwise_amax->resize(rows, 0.0f); - for (size_t row = 0; row < rows; ++row) { - float row_amax = 0.0f; - for (size_t col = 0; col < cols; ++col) { - row_amax = fmaxf(row_amax, fabsf(static_cast(input[row * cols + col]))); - } - (*rowwise_amax)[row] = row_amax; - quantize_nvfp4(OP, - input + row * cols, - output + row * (cols / 2), - scales + row * scales_stride, - 1, - cols, - scales_stride, - row_amax, - use_fast_math, - use_2d_quantization); - } - } else if (use_2d_quantization) { + // Ref impl for 2D quantization + if (use_2d_quantization) { // Step 1: Compute mathematical 8×8 scaling factors std::vector> math_scales; - compute_2d_mathematical_scales(OP, input, rows, cols, global_amax, math_scales, use_fast_math); + compute_2d_mathematical_scales(OP, input, rows, cols, *amax, math_scales, use_fast_math); constexpr size_t block_size_Y = 16; constexpr size_t block_size_X = 16; @@ -369,17 +361,36 @@ void compute_ref(float (*OP)(const float), // Step 4: Process quantized outputs using the same algorithm as quantize_nvfp4_2d // (This part processes the actual FP4 data using the mathematical scaling factors) - quantize_nvfp4_2d(OP, input, output, nullptr, rows, cols, scales_stride, global_amax, + quantize_nvfp4_2d(OP, input, output, nullptr, rows, cols, scales_stride, *amax, use_fast_math); // scales already filled - quantize_nvfp4_2d(OP, input_t.data(), output_t, nullptr, cols, rows, scales_stride_t, global_amax, + quantize_nvfp4_2d(OP, input_t.data(), output_t, nullptr, cols, rows, scales_stride_t, *amax, use_fast_math); // scales_t already filled - } else { - quantize_nvfp4(OP, input, output, scales, rows, cols, scales_stride, global_amax, - use_fast_math, use_2d_quantization); - quantize_nvfp4(OP, input_t.data(), output_t, scales_t, cols, rows, scales_stride_t, global_amax, - use_fast_math, use_2d_quantization); + return; + } + + // Ref impl for row-scaling + if (row_scaled_nvfp4) { + for (size_t row = 0; row < rows; ++row) { + quantize_nvfp4(OP, + input + row * cols, + output + row * (cols / 2), + scales + row * scales_stride, + 1, + cols, + scales_stride, + amax[row], + use_fast_math, + use_2d_quantization); + } + return; } + + // Ref impl for basic NVFP4 + quantize_nvfp4(OP, input, output, scales, rows, cols, scales_stride, *amax, + use_fast_math, use_2d_quantization); + quantize_nvfp4(OP, input_t.data(), output_t, scales_t, cols, rows, scales_stride_t, *amax, + use_fast_math, use_2d_quantization); } void compare_nvfp4_tensors(const std::string& name, @@ -479,48 +490,7 @@ void dump_nvfp4_tensor_data(const std::string& prefix, } } -void print_detailed_tensor_comparison(const std::string& name, - const fp4e2m1 *test_data, const fp4e2m1 *ref_data, - const int rows, const int cols) { - printf("\n=== DETAILED COMPARISON for %s (%d×%d = %d elements) ===\n", - name.c_str(), rows, cols, rows * cols); - - const int total_elements = rows * cols; - const int check_count = 128; - - printf("--- FIRST %d ELEMENTS ---\n", check_count); - printf("Index | Test_Value | Ref_Value | Match\n"); - printf("------|---------------|---------------|-------\n"); - for (int i = 0; i < std::min(check_count, total_elements); ++i) { - double2 test_pair = cvt_fp4x2_to_double2(*reinterpret_cast(&test_data[i/2])); - double2 ref_pair = cvt_fp4x2_to_double2(*reinterpret_cast(&ref_data[i/2])); - - double t = (i % 2 == 0) ? test_pair.x : test_pair.y; - double r = (i % 2 == 0) ? ref_pair.x : ref_pair.y; - bool match = (fabs(t - r) < 1e-6); - - printf("%5d | %13.6f | %13.6f | %s\n", i, t, r, match ? "✓" : "✗"); - } - - if (total_elements > 2 * check_count) { - printf("\n--- LAST %d ELEMENTS ---\n", check_count); - printf("Index | Test_Value | Ref_Value | Match\n"); - printf("------|---------------|---------------|-------\n"); - for (int i = total_elements - check_count; i < total_elements; ++i) { - double2 test_pair = cvt_fp4x2_to_double2(*reinterpret_cast(&test_data[i/2])); - double2 ref_pair = cvt_fp4x2_to_double2(*reinterpret_cast(&ref_data[i/2])); - - double t = (i % 2 == 0) ? test_pair.x : test_pair.y; - double r = (i % 2 == 0) ? ref_pair.x : ref_pair.y; - bool match = (fabs(t - r) < 1e-6); - - printf("%5d | %13.6f | %13.6f | %s\n", i, t, r, match ? "✓" : "✗"); - } - } - printf("==================================\n"); -} - -void compareResults_nvfp4(const Tensor &test, +void compareResults_nvfp4(Tensor &test, const void *ref, const void *ref_t, const int rows, const int cols, double atol = 1e-5, double rtol = 1e-8, bool if_on_gpus = true, bool dump_data = false, bool compare_columnwise = true) { @@ -529,10 +499,6 @@ void compareResults_nvfp4(const Tensor &test, const fp4e2m1 *test_data = test.rowwise_cpu_dptr(); const fp4e2m1 *ref_data = reinterpret_cast(ref); - // Print detailed element-by-element comparison - // print_detailed_tensor_comparison("output", test_data, ref_data, rows, cols); - // print_detailed_tensor_comparison("output_t", test_data_t, ref_data_t, cols, rows); - // Optionally dump tensor data to files for detailed analysis if (dump_data) { dump_nvfp4_tensor_data("output", test_data, ref_data, rows, cols); @@ -549,9 +515,10 @@ void compareResults_nvfp4(const Tensor &test, } } -void compare_rowwise_amax(const Tensor &output, const std::vector &ref_amax) { - const std::vector test_amax_data = output.tensor_amax_values(); - ASSERT_EQ(test_amax_data.size(), ref_amax.size()); +void compare_rowwise_amax(Tensor &output, const std::vector &ref_amax) { + ASSERT_EQ(output.rowwise_amax_size(), ref_amax.size()); + const auto *amax_ptr = output.cpu_rowwise_amax_ptr(); + const std::vector test_amax_data(amax_ptr, amax_ptr + ref_amax.size()); for (size_t row = 0; row < ref_amax.size(); ++row) { ASSERT_EQ(test_amax_data[row], ref_amax[row]) << "Row-scaled amax mismatch at row " << row; @@ -568,6 +535,9 @@ void performTest(float (*OP)(const float), DType itype = TypeInfo::dtype; DType otype = DType::kFloat4E2M1; + const bool rowwise = true; + const bool columnwise = !row_scaled_nvfp4; + const size_t rows = first_dimension(shape); const size_t cols = last_dimension(shape); @@ -589,7 +559,7 @@ void performTest(float (*OP)(const float), const size_t scales_stride_t = blocks_X_t; Tensor input("input", shape, itype); - Tensor output("output", shape, otype, true, !row_scaled_nvfp4, NVTE_NVFP4_1D_SCALING); + Tensor output("output", shape, otype, rowwise, columnwise, NVTE_NVFP4_1D_SCALING); std::unique_ptr ref_output = std::make_unique(rows * (cols / 2)); std::unique_ptr ref_output_t = std::make_unique(cols * (rows / 2)); @@ -598,58 +568,65 @@ void performTest(float (*OP)(const float), fillCase(&input, InputsFillCase::uniform); - // Golden value of amax chosen to make the 2nd-stage scaling mantissa zero and avoid rounding issues - const float amax = 448.0f * 6.0f * 8.0f; - std::vector ref_rowwise_amax; - bool use_2d_quantization = false; + // Compute 2nd stage NVFP4 scaling factor + std::vector ref_amax; if (row_scaled_nvfp4) { - output.set_tensor_amax_shape({rows}); - output.set_row_scaled_nvfp4(true); - compute_ref(OP, - input.rowwise_cpu_dptr(), - ref_output.get(), - ref_output_t.get(), - ref_scales.get(), - ref_scales_t.get(), - 0.0f, - rows, - cols, - scales_stride, - scales_stride_t, - use_fast_math, - use_2d_quantization, - &ref_rowwise_amax); + // Compute per-row amaxes + const auto *input_vals = input.rowwise_cpu_dptr(); + for (size_t row = 0; row < rows; ++row){ + float row_amax = 0.0f; + for (size_t col = 0; col < cols; ++col) { + row_amax = fmaxf(row_amax, fabsf(static_cast(input_vals[row * cols + col]))); + } + ref_amax.push_back(row_amax); + } + + // Update tensor + // Note: No need to update amax like standard NVFP4, amaxes + // are computed during quantization. + output.set_row_scaled_nvfp4(row_scaled_nvfp4); } else { - // Set 2nd stage NVFP4 scaling factor - output.set_tensor_amax(amax); - output.set_tensor_amax_columnwise(amax); - compute_ref(OP, - input.rowwise_cpu_dptr(), - ref_output.get(), - ref_output_t.get(), - ref_scales.get(), - ref_scales_t.get(), - amax, - rows, - cols, - scales_stride, - scales_stride_t, - use_fast_math, - use_2d_quantization); + // Golden value of amax chosen to make the 2nd-stage scaling mantissa zero and avoid rounding issues + ref_amax.assign(1, 448.0f * 6.0f * 8.0f); + + // Update tensor + if (rowwise) { + std::copy(ref_amax.begin(), ref_amax.end(), output.cpu_rowwise_amax_ptr()); + } + if (columnwise) { + std::copy(ref_amax.begin(), ref_amax.end(), output.cpu_columnwise_amax_ptr()); + } + output.from_cpu(); } + // Reference implementation + bool use_2d_quantization = false; + compute_ref(OP, + input.rowwise_cpu_dptr(), + ref_output.get(), + ref_output_t.get(), + ref_scales.get(), + ref_scales_t.get(), + ref_amax.data(), + rows, + cols, + scales_stride, + scales_stride_t, + use_fast_math, + use_2d_quantization, + row_scaled_nvfp4); + // Initialize stochastic rounding Tensor rng_state("rng_state", std::vector{2}, DType::kInt64); rng_state.rowwise_cpu_dptr()[0] = 123; // rng_seed rng_state.rowwise_cpu_dptr()[1] = 321; // rng_sequence rng_state.from_cpu(); + // Quantization options QuantizationConfigWrapper quant_config; quant_config.set_use_fast_math(use_fast_math); quant_config.set_stochastic_rounding(false); quant_config.set_rng_state(rng_state.data()); - - // Set 2D quantization based on compile-time flag quant_config.set_nvfp4_2d_quantization(use_2d_quantization); // Call appropriate function based on operation type @@ -696,9 +673,7 @@ void performTest(float (*OP)(const float), scale_mismatches_num); } - if (row_scaled_nvfp4) { - compare_rowwise_amax(output, ref_rowwise_amax); - } + compare_rowwise_amax(output, ref_amax); } std::vector> tensor_dims = { diff --git a/tests/cpp/operator/test_cast_transpose.cu b/tests/cpp/operator/test_cast_transpose.cu index 44c78e4a0..9a5dc959d 100644 --- a/tests/cpp/operator/test_cast_transpose.cu +++ b/tests/cpp/operator/test_cast_transpose.cu @@ -55,13 +55,13 @@ void performTest(const size_t N, const size_t H) { fillUniform(&input); setRandomScale(&output); + const float ref_scale = isFp8Type(otype) ? output.scale() : 1.0f; nvte_quantize(input.data(), output.data(), 0); float ref_amax; compute_ref(input.rowwise_cpu_dptr(), ref_output_c.get(), - ref_output_t.get(), N, H, &ref_amax, - output.scale()); + ref_output_t.get(), N, H, &ref_amax, ref_scale); cudaDeviceSynchronize(); auto err = cudaGetLastError(); @@ -69,7 +69,7 @@ void performTest(const size_t N, const size_t H) { if (isFp8Type(otype)) { auto [atol_amax, rtol_amax] = getTolerances(DType::kFloat32); compareResults("amax", output.amax(), ref_amax, atol_amax, rtol_amax); - float ref_scale_inv = 1.f / output.scale(); + float ref_scale_inv = 1.f / ref_scale; compareResults("scale_inv", output.rowwise_scale_inv(), ref_scale_inv, atol_amax, rtol_amax); } auto [atol, rtol] = getTolerances(otype); diff --git a/tests/cpp/operator/test_cast_transpose_dbias.cu b/tests/cpp/operator/test_cast_transpose_dbias.cu index 5b06b2832..f9303d34f 100644 --- a/tests/cpp/operator/test_cast_transpose_dbias.cu +++ b/tests/cpp/operator/test_cast_transpose_dbias.cu @@ -73,6 +73,7 @@ void performTest(const size_t N, const size_t H) { fillUniform(&input); setRandomScale(&output); + const float ref_scale = isFp8Type(otype) ? output.scale() : 1.0f; std::unique_ptr ref_output_c = std::make_unique(N*H); std::unique_ptr ref_output_t = std::make_unique(N*H); @@ -80,7 +81,7 @@ void performTest(const size_t N, const size_t H) { CType ref_amax; compute_ref_cast_transpose_dbias(input.rowwise_cpu_dptr(), - output.scale(), + ref_scale, ref_output_c.get(), ref_output_t.get(), &ref_amax, @@ -111,7 +112,7 @@ void performTest(const size_t N, const size_t H) { if (isFp8Type(otype)) { auto [atol_amax, rtol_amax] = getTolerances(DType::kFloat32); compareResults("amax", output.amax(), ref_amax, atol_amax, rtol_amax); - float ref_scale_inv = 1.f / output.scale(); + float ref_scale_inv = 1.f / ref_scale; compareResults("scale_inv", output.rowwise_scale_inv(), ref_scale_inv, atol_amax, rtol_amax); } auto [atol, rtol] = getTolerances(otype); diff --git a/tests/cpp/operator/test_cast_transpose_dbias_dgelu.cu b/tests/cpp/operator/test_cast_transpose_dbias_dgelu.cu index 9a4a2fa08..31eafff80 100644 --- a/tests/cpp/operator/test_cast_transpose_dbias_dgelu.cu +++ b/tests/cpp/operator/test_cast_transpose_dbias_dgelu.cu @@ -86,6 +86,7 @@ void performTest(const size_t N, const size_t H) { fillUniform(&input); fillUniform(&gelu_input); setRandomScale(&output); + const float ref_scale = isFp8Type(otype) ? output.scale() : 1.0f; std::unique_ptr ref_output_c = std::make_unique(N*H); std::unique_ptr ref_output_t = std::make_unique(N*H); @@ -94,7 +95,7 @@ void performTest(const size_t N, const size_t H) { CType ref_amax; compute_ref_cast_transpose_dbias_dgelu(input.rowwise_cpu_dptr(), gelu_input.rowwise_cpu_dptr(), - output.scale(), + ref_scale, ref_output_c.get(), ref_output_t.get(), &ref_amax, @@ -127,7 +128,7 @@ void performTest(const size_t N, const size_t H) { if (isFp8Type(otype)) { auto [atol_amax, rtol_amax] = getTolerances(DType::kFloat32); compareResults("amax", output.amax(), ref_amax, atol_amax, rtol_amax); - float ref_scale_inv = 1.f / output.scale(); + float ref_scale_inv = 1.f / ref_scale; compareResults("scale_inv", output.rowwise_scale_inv(), ref_scale_inv, atol_amax, rtol_amax); } diff --git a/tests/cpp/operator/test_cast_transpose_dgeglu.cu b/tests/cpp/operator/test_cast_transpose_dgeglu.cu index a87c0c5a4..15ecd3ab6 100644 --- a/tests/cpp/operator/test_cast_transpose_dgeglu.cu +++ b/tests/cpp/operator/test_cast_transpose_dgeglu.cu @@ -81,6 +81,7 @@ void performTest(const size_t N, const size_t H) { fillUniform(&grad); fillUniform(&input); setRandomScale(&output); + const float ref_scale = isFp8Type(otype) ? output.scale() : 1.0f; std::unique_ptr ref_output_c = std::make_unique(N * H * 2); std::unique_ptr ref_output_t = std::make_unique(N * H * 2); @@ -89,7 +90,7 @@ void performTest(const size_t N, const size_t H) { CType ref_amax; compute_ref_cast_transpose_dgated_gelu(grad.rowwise_cpu_dptr(), input.rowwise_cpu_dptr(), - output.scale(), ref_output_c.get(), ref_output_t.get(), + ref_scale, ref_output_c.get(), ref_output_t.get(), &ref_amax, N, H); cudaDeviceSynchronize(); @@ -99,7 +100,7 @@ void performTest(const size_t N, const size_t H) { if (isFp8Type(otype)) { auto [atol_amax, rtol_amax] = getTolerances(DType::kFloat32); compareResults("amax", output.amax(), ref_amax, atol_amax, rtol_amax); - float ref_scale_inv = 1.f / output.scale(); + float ref_scale_inv = 1.f / ref_scale; compareResults("scale_inv", output.rowwise_scale_inv(), ref_scale_inv, atol_amax, rtol_amax); } diff --git a/tests/cpp/operator/test_dequantize_nvfp4.cu b/tests/cpp/operator/test_dequantize_nvfp4.cu index ec405b1d9..eb9e8bce2 100644 --- a/tests/cpp/operator/test_dequantize_nvfp4.cu +++ b/tests/cpp/operator/test_dequantize_nvfp4.cu @@ -76,7 +76,7 @@ void compute_ref_dequantize_nvfp4(const uint8_t *packed_data, } template -float compute_amax(const test::Tensor &t, size_t rows, size_t cols) { +float compute_amax(test::Tensor &t, size_t rows, size_t cols) { t.to_cpu(); const auto *data = t.rowwise_cpu_dptr(); float amax = 0.0f; @@ -94,50 +94,63 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols, using namespace test; DType otype = TypeInfo::dtype; + // Tensors Tensor input("input", std::vector{rows, cols}, otype); - fillCase(&input, InputsFillCase::uniform); - Tensor quantized("quantized", std::vector{rows, cols}, DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING); + Tensor output("output", std::vector{rows, cols}, otype, true, false); + + // Fill input with random data + fillCase(&input, InputsFillCase::uniform); + + // Configure quantized tensor amax + size_t amax_size = 1; if (row_scaled_nvfp4) { - quantized.set_tensor_amax_shape({rows}); - quantized.set_row_scaled_nvfp4(true); + quantized.set_row_scaled_nvfp4(true); + amax_size = rows; } else if (rows > 0 && cols > 0) { - quantized.set_tensor_amax(compute_amax(input, rows, cols)); + quantized.set_amax(compute_amax(input, rows, cols)); } else { - quantized.set_tensor_amax(0.0f); + quantized.set_amax(0.0f); } + // Quantize if (rows > 0 && cols > 0) { nvte_quantize(input.data(), quantized.data(), 0); cudaDeviceSynchronize(); + auto err = cudaGetLastError(); + ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); } - Tensor output("output", std::vector{rows, cols}, otype, true, false); + // Dequantize nvte_dequantize(quantized.data(), output.data(), 0); cudaDeviceSynchronize(); - auto err = cudaGetLastError(); ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); - if (rows > 0 && cols > 0) { - quantized.to_cpu(); - const uint8_t *fp4_data = - reinterpret_cast(quantized.rowwise_cpu_dptr()); - const fp8e4m3 *scales = quantized.rowwise_cpu_scale_inv_ptr(); - const std::vector amax_val = quantized.tensor_amax_values(); - const NVTEShape scale_shape = quantized.rowwise_scale_inv_shape(); - const size_t scale_stride = scale_shape.data[scale_shape.ndim - 1]; - - std::unique_ptr ref_output = - std::make_unique(rows * cols); - compute_ref_dequantize_nvfp4( - fp4_data, scales, amax_val, ref_output.get(), - rows, cols, scale_stride); - - auto [atol, rtol] = getTolerances(otype); - compareResults("output_nvfp4", output, ref_output.get(), true, atol, rtol); + // Nothing to be done if tensor is empty + if (rows == 0 && cols == 0) { + return; } + + // Dequantize reference implementation + quantized.to_cpu(); + const uint8_t *fp4_data = + reinterpret_cast(quantized.rowwise_cpu_dptr()); + const fp8e4m3 *scales = quantized.rowwise_cpu_scale_inv_ptr(); + const auto *amax = quantized.cpu_rowwise_amax_ptr(); + const std::vector amax_vals(amax, amax + amax_size); + const NVTEShape scale_shape = quantized.rowwise_scale_inv_shape(); + const size_t scale_stride = scale_shape.data[scale_shape.ndim - 1]; + std::unique_ptr ref_output = + std::make_unique(rows * cols); + compute_ref_dequantize_nvfp4( + fp4_data, scales, amax_vals, ref_output.get(), + rows, cols, scale_stride); + + // Compare results from TE and reference impls + auto [atol, rtol] = getTolerances(otype); + compareResults("output_nvfp4", output, ref_output.get(), true, atol, rtol); } // Dequantize NVFP4 with GEMM-swizzled scales and compare against compact path. @@ -153,12 +166,11 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols, Tensor quantized_compact("quantized_compact", std::vector{rows, cols}, DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING); if (row_scaled_nvfp4) { - quantized_compact.set_tensor_amax_shape({rows}); quantized_compact.set_row_scaled_nvfp4(true); } else if (rows > 0 && cols > 0) { - quantized_compact.set_tensor_amax(compute_amax(input, rows, cols)); + quantized_compact.set_amax(compute_amax(input, rows, cols)); } else { - quantized_compact.set_tensor_amax(0.0f); + quantized_compact.set_amax(0.0f); } if (rows > 0 && cols > 0) { @@ -175,10 +187,9 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols, Tensor quantized_swizzled("quantized_swizzled", std::vector{rows, cols}, DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING); if (row_scaled_nvfp4) { - quantized_swizzled.set_tensor_amax_shape({rows}); quantized_swizzled.set_row_scaled_nvfp4(true); } else { - quantized_swizzled.set_tensor_amax(0.0f); + quantized_swizzled.set_amax(0.0f); } quantized_swizzled.set_with_gemm_swizzled_scales(true); @@ -186,9 +197,12 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols, // since from_cpu() uploads all CPU buffers (including zero-init data). quantized_compact.to_cpu(); if (row_scaled_nvfp4) { - quantized_swizzled.copy_tensor_amax_from(quantized_compact); + const auto *src = quantized_compact.cpu_rowwise_amax_ptr(); + auto *dst = quantized_swizzled.cpu_rowwise_amax_ptr(); + std::copy(src, src + rows, dst); + quantized_swizzled.from_cpu(); } else { - quantized_swizzled.set_tensor_amax(quantized_compact.amax()); + quantized_swizzled.set_amax(quantized_compact.amax()); } // Copy FP4 data after from_cpu() to avoid being overwritten diff --git a/tests/cpp/operator/test_multi_cast_transpose.cu b/tests/cpp/operator/test_multi_cast_transpose.cu index 2bb35c4b8..0271c9dc6 100644 --- a/tests/cpp/operator/test_multi_cast_transpose.cu +++ b/tests/cpp/operator/test_multi_cast_transpose.cu @@ -97,7 +97,7 @@ void performTest() { std::copy(input.rowwise_cpu_dptr(), input.rowwise_cpu_dptr() + height * width, ref_input_list.back().begin()); - ref_scale_list[tensor_id] = output.scale(); + ref_scale_list[tensor_id] = isFp8Type(otype) ? output.scale() : 1.0f; ref_height_list[tensor_id] = height; ref_width_list[tensor_id] = width; } @@ -138,7 +138,7 @@ void performTest() { atol_amax, rtol_amax); compareResults("scale_inv", output_list[tensor_id].rowwise_scale_inv(), - 1.f / output_list[tensor_id].scale(), + 1.f / ref_scale_list[tensor_id], atol_amax, rtol_amax); } auto [atol, rtol] = getTolerances(otype); diff --git a/tests/cpp/operator/test_normalization.cu b/tests/cpp/operator/test_normalization.cu index f737005e2..ea6692dba 100644 --- a/tests/cpp/operator/test_normalization.cu +++ b/tests/cpp/operator/test_normalization.cu @@ -208,7 +208,7 @@ void performTest(const size_t N, const size_t H, const bool zero_centered_gamma, auto [atol_amax, rtol_amax] = getTolerances(DType::kFloat32); if (isFp8Type(otype)) { compareResults("amax", z.amax(), ref_amax, atol_amax, rtol_amax); - float ref_scale_inv = 1.f / z.scale(); + float ref_scale_inv = 1.f / ref_scale; compareResults("scale_inv", z.rowwise_scale_inv(), ref_scale_inv, atol_amax, rtol_amax); } diff --git a/tests/cpp/operator/test_qdq.cu b/tests/cpp/operator/test_qdq.cu index 4e364fffa..034280aa9 100644 --- a/tests/cpp/operator/test_qdq.cu +++ b/tests/cpp/operator/test_qdq.cu @@ -65,12 +65,13 @@ void performTestQ(const size_t N) { fillUniform(&input); setRandomScale(&output); + const float ref_scale = output.scale(); nvte_quantize(input.data(), output.data(), 0); float ref_amax; compute_ref_q(input.rowwise_cpu_dptr(), ref_output.get(), - N, &ref_amax, output.scale()); + N, &ref_amax, ref_scale); cudaDeviceSynchronize(); auto err = cudaGetLastError(); diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index 96e71f951..4fd75bb92 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -8,12 +8,14 @@ #include "test_common.h" #include +#include +#include +#include +#include #include #include +#include #include -#include -#include -#include #include #include @@ -193,33 +195,6 @@ std::pair get_scales(const NVTEShape& shape, return {ret_rowwise, ret_colwise}; } - if (scaling_mode == NVTE_MXFP8_1D_SCALING) { - std::vector shape_vec; - for (size_t i = 0; i < shape.ndim; ++i) { - shape_vec.push_back(shape.data[i]); - } - size_t first_dim = first_dimension(shape_vec); - size_t last_dim = last_dimension(shape_vec); - - scale_inv_meta ret_rowwise, ret_colwise; - - const size_t block_size_X_rowwise = 32; - size_t scale_dim_Y_rowwise = DIVUP_TO_MULTIPLE(first_dim, scale_tensor_alignment_Y_rowwise); - size_t scale_dim_X_rowwise = DIVUP_TO_MULTIPLE(DIVUP(last_dim, block_size_X_rowwise), scale_tensor_alignment_X_rowwise); - ret_rowwise.shape = {scale_dim_Y_rowwise, scale_dim_X_rowwise}; - - const size_t block_size_Y_colwise = 32; - size_t scale_dim_Y_colwise = DIVUP_TO_MULTIPLE(DIVUP(first_dim, block_size_Y_colwise), scale_tensor_alignment_Y_colwise); - size_t scale_dim_X_colwise = DIVUP_TO_MULTIPLE(last_dim, scale_tensor_alignment_X_colwise); - ret_colwise.shape = {scale_dim_Y_colwise, scale_dim_X_colwise}; - - ret_rowwise.type = DType::kFloat8E8M0; - ret_colwise.type = DType::kFloat8E8M0; - ret_rowwise.type_size_bits = typeToNumBits(DType::kFloat8E8M0); - ret_colwise.type_size_bits = typeToNumBits(DType::kFloat8E8M0); - - return {ret_rowwise, ret_colwise}; - } if (scaling_mode == NVTE_BLOCK_SCALING_2D) { std::vector shape_vec; for (size_t i = 0; i < shape.ndim; ++i) { @@ -276,6 +251,30 @@ std::pair get_scales(const NVTEShape& shape, NVTE_ERROR("Invalid scaling mode!"); } +Tensor::Buffer::Buffer(size_t size, DType dtype) + : size_{size}, dtype_{dtype}, bytes_{size * typeToNumBits(dtype) / 8} { + if (bytes_ > 0) { + cpu_buffer_.reset(new unsigned char[bytes_]); + std::memset(cpu_buffer_.get(), 0, bytes_); + unsigned char *gpu_buffer = nullptr; + NVTE_CHECK_CUDA(cudaMalloc(&gpu_buffer, bytes_)); + gpu_buffer_.reset(gpu_buffer); + NVTE_CHECK_CUDA(cudaMemset(gpu_buffer_.get(), 0, bytes_)); + } +} + +void Tensor::Buffer::to_cpu() { + if (bytes_ > 0) { + NVTE_CHECK_CUDA(cudaMemcpy(cpu_buffer_.get(), gpu_buffer_.get(), bytes_, cudaMemcpyDeviceToHost)); + } +} + +void Tensor::Buffer::from_cpu() { + if (bytes_ > 0) { + NVTE_CHECK_CUDA(cudaMemcpy(gpu_buffer_.get(), cpu_buffer_.get(), bytes_, cudaMemcpyHostToDevice)); + } +} + Tensor::Tensor(const std::string& name, const NVTEShape &shape, const DType type, const bool rowwise, const bool columnwise, @@ -303,31 +302,13 @@ Tensor::Tensor(const std::string& name, flattened_shape = convertShape(flattened_shape_vec); } - // Allocate and initialize data - void *dptr_rowwise = nullptr, *dptr_columnwise = nullptr; - const size_t total_size = bytes(shape, type); - if (total_size != 0) { - if (rowwise) { - cudaMalloc((void**)&dptr_rowwise, total_size); // NOLINT(*) - cudaMemset(dptr_rowwise, 0, total_size); - cpu_data_rowwise_ = std::make_unique(total_size); - std::fill_n(cpu_data_rowwise_.get(), total_size, 0); - } - if (columnwise) { - cudaMalloc((void**)&dptr_columnwise, total_size); // NOLINT(*) - cudaMemset(dptr_columnwise, 0, total_size); - cpu_data_columnwise_ = std::make_unique(total_size); - std::fill_n(cpu_data_columnwise_.get(), total_size, 0); - } - } - - // Set tensor row-wise data + // Allocate row-wise data if (rowwise) { - const DType rowwise_type = (scaling_mode == NVTE_NVFP4_1D_SCALING) ? DType::kFloat4E2M1 : type; - tensor_.set_rowwise_data(dptr_rowwise, rowwise_type, shape); + data_rowwise_.emplace(product(shape), type); + tensor_.set_rowwise_data(data_rowwise_->gpu_buffer(), type, shape); } - // Set tensor column-wise data + // Allocate column-wise data if (columnwise) { // Determine shape of column-wise data std::vector columnwise_shape_vec; @@ -358,310 +339,224 @@ Tensor::Tensor(const std::string& name, const auto columnwise_shape = nvte_make_shape(columnwise_shape_vec.data(), columnwise_shape_vec.size()); - // Set column-wise data buffer - const DType colwise_type = (scaling_mode == NVTE_NVFP4_1D_SCALING) ? DType::kFloat4E2M1 : type; - tensor_.set_columnwise_data(dptr_columnwise, colwise_type, columnwise_shape); + // Allocate buffer + data_columnwise_.emplace(product(columnwise_shape), type); + + // Configure TE tensor + tensor_.set_columnwise_data(data_columnwise_->gpu_buffer(), type, columnwise_shape); } - // Configure scales, amaxes, and other tensor buffers - float *amax = nullptr; - float *amax_columnwise = nullptr; - float *scale = nullptr; - float *rowwise_scale_inv = nullptr; - float *columnwise_scale_inv = nullptr; - if (isFp8Type(type) || isFp4Type(type)) { - if (scaling_mode == NVTE_DELAYED_TENSOR_SCALING) { - cudaMalloc((void**)&amax, sizeof(float)); // NOLINT(*) - cudaMemset(amax, 0, sizeof(float)); - cudaMalloc((void**)&scale, sizeof(float)); // NOLINT(*) - cudaMemset(scale, 0, sizeof(float)); - amax_cpu_data_ = std::make_shared(0); - scale_cpu_data_ = std::make_shared(0); - tensor_.set_amax(amax, DType::kFloat32, std::vector{1}); - tensor_.set_scale(scale, DType::kFloat32, std::vector{1}); - cudaMalloc((void**)&rowwise_scale_inv, sizeof(float)); // NOLINT(*) + // Allocate recipe-specific buffers + switch (scaling_mode) { + case NVTE_DELAYED_TENSOR_SCALING: + if (isFp8Type(type)) { + amax_rowwise_.emplace(1, DType::kFloat32); + scale_.emplace(1, DType::kFloat32); + tensor_.set_amax(amax_rowwise_->gpu_buffer(), DType::kFloat32, std::vector{1}); + tensor_.set_scale(scale_->gpu_buffer(), DType::kFloat32, std::vector{1}); + + // Use same buffer for row-wise and column-wise scale-inverse + auto scale_inv = std::make_shared(1, DType::kFloat32); if (rowwise) { - tensor_.set_rowwise_scale_inv(rowwise_scale_inv, DType::kFloat32, - std::vector{1}); - rowwise_scale_inv_cpu_data_ = std::make_unique(sizeof(float)); - std::fill_n(rowwise_scale_inv_cpu_data_.get(), sizeof(float), 0); + scale_inv_rowwise_ = scale_inv; + tensor_.set_rowwise_scale_inv(scale_inv_rowwise_->gpu_buffer(), DType::kFloat32, std::vector{1}); } if (columnwise) { - tensor_.set_columnwise_scale_inv(rowwise_scale_inv, DType::kFloat32, - std::vector{1}); - columnwise_scale_inv_cpu_data_ = std::make_unique(sizeof(float)); - std::fill_n(columnwise_scale_inv_cpu_data_.get(), sizeof(float), 0); - } - } else { - if (scaling_mode == NVTE_NVFP4_1D_SCALING) { - // Used for NVFP4 second stage scaling - amax_cpu_data_ = std::make_shared(0); - amax_cpu_data_columnwise_ = std::make_shared(0); - cudaMalloc((void**)&amax, sizeof(float)); // NOLINT(*) - cudaMalloc((void**)&amax_columnwise, sizeof(float)); // NOLINT(*) - cudaMemset(amax, 0, sizeof(float)); - cudaMemset(amax_columnwise, 0, sizeof(float)); - tensor_.set_amax(amax, DType::kFloat32, std::vector{1}); - tensor_.set_columnwise_amax(amax_columnwise, DType::kFloat32, std::vector{1}); + scale_inv_columnwise_ = scale_inv; + tensor_.set_columnwise_scale_inv(scale_inv_columnwise_->gpu_buffer(), DType::kFloat32, std::vector{1}); } + } + break; + case NVTE_MXFP8_1D_SCALING: + case NVTE_BLOCK_SCALING_1D: + case NVTE_BLOCK_SCALING_2D: + case NVTE_NVFP4_1D_SCALING: + { + // Block scaling factors auto [rowwise_scale_meta, colwise_scale_meta] = get_scales(flattened_shape, tensor_.scaling_mode()); - auto rowwise_scale_size = rowwise_scale_meta.bytes(); - auto columnwise_scale_size = colwise_scale_meta.bytes(); - auto scale_shape = rowwise_scale_meta.shape; - auto columnwise_scale_shape = colwise_scale_meta.shape; if (rowwise) { - cudaMalloc((void **)&rowwise_scale_inv, rowwise_scale_size); // NOLINT(*) - cudaMemset(rowwise_scale_inv, 0, rowwise_scale_size); - rowwise_scale_inv_cpu_data_ = std::make_unique(rowwise_scale_size); - std::fill_n(rowwise_scale_inv_cpu_data_.get(), rowwise_scale_size, 0); - auto scale_dtype = rowwise_scale_meta.type; - tensor_.set_rowwise_scale_inv(rowwise_scale_inv, scale_dtype, scale_shape); + const auto scale_shape = rowwise_scale_meta.shape; + const auto scale_dtype = rowwise_scale_meta.type; + scale_inv_rowwise_ = std::make_shared(product(scale_shape), scale_dtype); + tensor_.set_rowwise_scale_inv(scale_inv_rowwise_->gpu_buffer(), scale_dtype, scale_shape); } if (columnwise) { - cudaMalloc((void**)&columnwise_scale_inv, columnwise_scale_size); // NOLINT(*) - cudaMemset(columnwise_scale_inv, 0, columnwise_scale_size); - columnwise_scale_inv_cpu_data_ = std::make_unique(columnwise_scale_size); - std::fill_n(columnwise_scale_inv_cpu_data_.get(), columnwise_scale_size, 0); - auto scale_dtype = colwise_scale_meta.type; - tensor_.set_columnwise_scale_inv(columnwise_scale_inv, scale_dtype, columnwise_scale_shape); + const auto scale_shape = colwise_scale_meta.shape; + const auto scale_dtype = colwise_scale_meta.type; + scale_inv_columnwise_ = std::make_shared(product(scale_shape), scale_dtype); + tensor_.set_columnwise_scale_inv(scale_inv_columnwise_->gpu_buffer(), scale_dtype, scale_shape); } - } - } -} -void Tensor::to_cpu() const { - const NVTEShape s = tensor_.shape(); - const size_t size = bytes(s, tensor_.dtype()); - if (rowwise_) { - cudaMemcpy(cpu_data_rowwise_.get(), - tensor_.get_rowwise_data().data_ptr, - size, - cudaMemcpyDeviceToHost); - } - if (columnwise_) { - const DType colwise_type = tensor_.dtype(); - - const size_t colwise_size = bytes(s, colwise_type); - cudaMemcpy(cpu_data_columnwise_.get(), - tensor_.get_columnwise_data().data_ptr, - colwise_size, - cudaMemcpyDeviceToHost); - } - if (isFp8Type(dtype()) || isFp4Type(dtype())) { - if (tensor_.scaling_mode() == NVTE_DELAYED_TENSOR_SCALING) { - if (tensor_.amax() != nullptr){ - cudaMemcpy(amax_cpu_data_.get(), - tensor_.amax(), - sizeof(float), - cudaMemcpyDeviceToHost); - } - cudaMemcpy(scale_cpu_data_.get(), - tensor_.scale(), - sizeof(float), - cudaMemcpyDeviceToHost); - } else if (tensor_.scaling_mode() == NVTE_NVFP4_1D_SCALING) { - if (rowwise_ && (tensor_.amax() != nullptr)){ - cudaMemcpy(amax_cpu_data_.get(), - tensor_.amax(), - sizeof(float), - cudaMemcpyDeviceToHost); - } - if (columnwise_ && (tensor_.get_columnwise_amax().data_ptr != nullptr)){ - cudaMemcpy(amax_cpu_data_columnwise_.get(), - tensor_.get_columnwise_amax().data_ptr, - sizeof(float), - cudaMemcpyDeviceToHost); + // NVFP4 uses amax for tensor scaling + if (scaling_mode == NVTE_NVFP4_1D_SCALING) { + if (rowwise) { + amax_rowwise_.emplace(1, DType::kFloat32); + tensor_.set_amax(amax_rowwise_->gpu_buffer(), DType::kFloat32, std::vector{1}); + } + if (columnwise) { + amax_columnwise_.emplace(1, DType::kFloat32); + tensor_.set_columnwise_amax(amax_columnwise_->gpu_buffer(), DType::kFloat32, std::vector{1}); + } } } - auto [rowwise_scale_meta, colwise_scale_meta] = get_scales(s, tensor_.scaling_mode()); - if (rowwise_) { - auto scale_size = rowwise_scale_meta.bytes(); - cudaMemcpy(rowwise_scale_inv_cpu_data_.get(), - tensor_.get_rowwise_scale_inv().data_ptr, - scale_size, - cudaMemcpyDeviceToHost); - } - if (columnwise_) { - auto scale_size = colwise_scale_meta.bytes(); - cudaMemcpy(columnwise_scale_inv_cpu_data_.get(), - tensor_.get_columnwise_scale_inv().data_ptr, - scale_size, - cudaMemcpyDeviceToHost); - } + break; + default: + NVTE_ERROR("Unsupported tensor format (", static_cast(scaling_mode), ")"); } } -void Tensor::from_cpu() const { - const NVTEShape s = tensor_.shape(); - const size_t size = bytes(s, tensor_.dtype()); - if (rowwise_) { - cudaMemcpy(tensor_.get_rowwise_data().data_ptr, cpu_data_rowwise_.get(), size, - cudaMemcpyHostToDevice); - } - if (columnwise_) { - cudaMemcpy(tensor_.get_columnwise_data().data_ptr, cpu_data_columnwise_.get(), size, - cudaMemcpyHostToDevice); - } - if (isFp8Type(dtype()) || isFp4Type(dtype())) { - if (tensor_.scaling_mode() == NVTE_DELAYED_TENSOR_SCALING) { - if (tensor_.amax() != nullptr){ - cudaMemcpy(tensor_.amax(), amax_cpu_data_.get(), sizeof(float), cudaMemcpyHostToDevice); - } - cudaMemcpy(tensor_.scale(), scale_cpu_data_.get(), sizeof(float), cudaMemcpyHostToDevice); - } else if (tensor_.scaling_mode() == NVTE_NVFP4_1D_SCALING) { - if (rowwise_ && (tensor_.amax() != nullptr)) { - cudaMemcpy(tensor_.amax(), amax_cpu_data_.get(), sizeof(float), cudaMemcpyHostToDevice); - } - if (columnwise_ && (tensor_.get_columnwise_amax().data_ptr != nullptr)) { - cudaMemcpy(tensor_.get_columnwise_amax().data_ptr, amax_cpu_data_columnwise_.get(), - sizeof(float), cudaMemcpyHostToDevice); - } - } - auto [rowwise_scale_meta, colwise_scale_meta] = get_scales(s, tensor_.scaling_mode()); +void Tensor::set_tensor_amax_nullptr() { + tensor_.set_amax(nullptr, DType::kFloat32, tensor_.defaultShape); +} + +void Tensor::set_with_gemm_swizzled_scales(bool with_gemm_swizzled_scales) { + tensor_.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); +} + +void Tensor::set_row_scaled_nvfp4(bool row_scaled_nvfp4) { + NVTE_CHECK(tensor_.scaling_mode() == NVTE_NVFP4_1D_SCALING, + "Row-scaled NVFP4 is only supported for NVFP4 tensors."); + tensor_.set_row_scaled_nvfp4(row_scaled_nvfp4); + + // Update amax tensor + if (row_scaled_nvfp4) { + // Row-scaled NVFP4 has amax matching number of rows + NVTE_CHECK(rowwise_, "Row-scaled NVFP4 requires row-wise data."); + NVTE_CHECK(!columnwise_, "Row-scaled NVFP4 does not support column-wise data."); + auto shape = tensor_.shape(); + const size_t rows = product(shape, 0, shape.ndim - 1); + amax_rowwise_.emplace(rows, DType::kFloat32); + tensor_.set_amax(amax_rowwise_->gpu_buffer(), DType::kFloat32, std::vector{rows}); + } else { + // Tensor-scaled NVFP4 has single amax if (rowwise_) { - auto scale_size = rowwise_scale_meta.bytes(); - cudaMemcpy(tensor_.get_rowwise_scale_inv().data_ptr, - rowwise_scale_inv_cpu_data_.get(), scale_size, - cudaMemcpyHostToDevice); + amax_rowwise_.emplace(1, DType::kFloat32); + tensor_.set_amax(amax_rowwise_->gpu_buffer(), DType::kFloat32, std::vector{1}); } if (columnwise_) { - auto scale_size = colwise_scale_meta.bytes(); - cudaMemcpy(tensor_.get_columnwise_scale_inv().data_ptr, - columnwise_scale_inv_cpu_data_.get(), scale_size, - cudaMemcpyHostToDevice); + amax_columnwise_.emplace(1, DType::kFloat32); + tensor_.set_columnwise_amax(amax_columnwise_->gpu_buffer(), DType::kFloat32, std::vector{1}); } } } -void Tensor::set_scale(float scale) { - if (isFp8Type(dtype()) || isFp4Type(dtype())) { - NVTE_CHECK(scale_cpu_data_); - if (tensor_.scaling_mode() == NVTE_DELAYED_TENSOR_SCALING) { - *scale_cpu_data_ = scale; - from_cpu(); - } - } +void Tensor::to_cpu() { + if (data_rowwise_) { data_rowwise_->to_cpu(); } + if (data_columnwise_) { data_columnwise_->to_cpu(); } + if (scale_inv_rowwise_) { scale_inv_rowwise_->to_cpu(); } + if (scale_inv_columnwise_) { scale_inv_columnwise_->to_cpu(); } + if (amax_rowwise_) { amax_rowwise_->to_cpu(); } + if (amax_columnwise_) { amax_columnwise_->to_cpu(); } + if (scale_) { scale_->to_cpu(); } } -void Tensor::set_tensor_amax_shape(const std::vector &shape) { - const size_t numel = product(shape); - NVTE_CHECK(tensor_.scaling_mode() == NVTE_NVFP4_1D_SCALING, - "Amax shape override is only supported for NVFP4 test tensors."); - - auto old_amax = tensor_.get_amax(); - if (old_amax.data_ptr != nullptr) { - NVTE_CHECK_CUDA(cudaFree(old_amax.data_ptr)); - } - - float *amax = nullptr; - NVTE_CHECK_CUDA(cudaMalloc(&amax, numel * sizeof(float))); - NVTE_CHECK_CUDA(cudaMemset(amax, 0, numel * sizeof(float))); - tensor_.set_amax(amax, DType::kFloat32, shape); +void Tensor::from_cpu() { + if (data_rowwise_) { data_rowwise_->from_cpu(); } + if (data_columnwise_) { data_columnwise_->from_cpu(); } + if (scale_inv_rowwise_) { scale_inv_rowwise_->from_cpu(); } + if (scale_inv_columnwise_) { scale_inv_columnwise_->from_cpu(); } + if (amax_rowwise_) { amax_rowwise_->from_cpu(); } + if (amax_columnwise_) { amax_columnwise_->from_cpu(); } + if (scale_) { scale_->from_cpu(); } } -std::vector Tensor::tensor_amax_values() const { - const auto amax = tensor_.get_amax(); - NVTE_CHECK(static_cast(amax.dtype) == DType::kFloat32, "Tensor amax must be FP32."); - - const size_t numel = product(amax.shape); - if (numel == 0) { - return {}; - } - NVTE_CHECK(amax.data_ptr != nullptr, "Tensor amax is not allocated."); - - std::vector values(numel); - NVTE_CHECK_CUDA( - cudaMemcpy(values.data(), amax.data_ptr, numel * sizeof(float), cudaMemcpyDeviceToHost)); - return values; +void Tensor::set_amax(float amax) { + NVTE_CHECK(amax_rowwise_); + NVTE_CHECK(amax_rowwise_->size() == 1); + NVTE_CHECK(amax_rowwise_->dtype() == DType::kFloat32); + *amax_rowwise_->cpu_buffer() = amax; + amax_rowwise_->from_cpu(); } -void Tensor::copy_tensor_amax_from(const Tensor &other) { - const auto other_amax = other.tensor_.get_amax(); - NVTE_CHECK(static_cast(other_amax.dtype) == DType::kFloat32, - "Source tensor amax must be FP32."); - - auto my_amax = tensor_.get_amax(); - NVTE_CHECK(static_cast(my_amax.dtype) == DType::kFloat32, - "Destination tensor amax must be FP32."); - NVTE_CHECK(areShapesEqual(my_amax.shape, other_amax.shape), "Amax shape mismatch."); +void Tensor::set_scale(float scale) { + NVTE_CHECK(scale_); + NVTE_CHECK(scale_->size() == 1); + NVTE_CHECK(scale_->dtype() == DType::kFloat32); + *scale_->cpu_buffer() = scale; + scale_->from_cpu(); +} - const size_t numel = product(other_amax.shape); - if (numel == 0) { - return; - } +void Tensor::set_scale_inv(float scale_inv) { + NVTE_CHECK(scale_inv_rowwise_); + NVTE_CHECK(scale_inv_rowwise_->size() == 1); + NVTE_CHECK(scale_inv_rowwise_->dtype() == DType::kFloat32); + *scale_inv_rowwise_->cpu_buffer() = scale_inv; + scale_inv_rowwise_->from_cpu(); +} - NVTE_CHECK(other_amax.data_ptr != nullptr, "Source tensor amax is not allocated."); - NVTE_CHECK(my_amax.data_ptr != nullptr, "Destination tensor amax is not allocated."); - NVTE_CHECK_CUDA(cudaMemcpy(my_amax.data_ptr, other_amax.data_ptr, numel * sizeof(float), - cudaMemcpyDeviceToDevice)); +void Tensor::set_tensor_amax_columnwise(float amax) { + NVTE_CHECK(amax_columnwise_); + NVTE_CHECK(amax_columnwise_->size() == 1); + NVTE_CHECK(amax_columnwise_->dtype() == DType::kFloat32); + *amax_columnwise_->cpu_buffer() = amax; + amax_columnwise_->from_cpu(); } -void Tensor::set_scale_inv(float scale_inv) { - if (isFp8Type(dtype()) || isFp4Type(dtype())) { - if (rowwise_) { - NVTE_CHECK(rowwise_scale_inv_cpu_data_); - } - if (columnwise_) { - NVTE_CHECK(columnwise_scale_inv_cpu_data_); - } +namespace { - auto [rowwise_scale_meta, colwise_scale_meta] = get_scales(tensor_.shape(), tensor_.scaling_mode()); - if (rowwise_) { - auto num_scales = product(rowwise_scale_meta.shape); - if (num_scales == 1) { - rowwise_cpu_scale_inv_ptr()[0] = scale_inv; - } else { - std::uniform_int_distribution dis(0, 127); - auto *scale_inv_ptr = rowwise_cpu_scale_inv_ptr(); - for (size_t i = 0; i < num_scales; i++) { - scale_inv_ptr[i] = dis(gen_); - } +void fill_uniform_buffer(void *cpu_data, size_t numel, DType dtype, std::mt19937 &gen) { + switch (dtype) { + case DType::kFloat32: + { + auto *data = static_cast(cpu_data); + std::uniform_real_distribution dis(-2.0, 1.0); + for (size_t i = 0; i < numel; ++i) { + data[i] = dis(gen); } } - if (columnwise_) { - auto num_scales = product(colwise_scale_meta.shape); - if (num_scales == 1) { - columnwise_cpu_scale_inv_ptr()[0] = scale_inv; - } else { - std::uniform_int_distribution dis(0, 127); - auto *scale_inv_ptr = columnwise_cpu_scale_inv_ptr(); - for (size_t i = 0; i < num_scales; i++) { - scale_inv_ptr[i] = dis(gen_); - } + break; + case DType::kFloat8E4M3: + case DType::kFloat8E8M0: + case DType::kByte: + { + auto *data = static_cast(cpu_data); + std::uniform_int_distribution dis(0, 127); + for (size_t i = 0; i < numel; ++i) { + data[i] = dis(gen); } } - from_cpu(); + break; + default: + NVTE_ERROR("Unsupported dtype (", static_cast(dtype), ")."); + } +} + +} // namespace + +void Tensor::fill_uniform_rowwise_scale_inv() { + if (!scale_inv_rowwise_ || scale_inv_rowwise_->size() == 0) { + return; + } + fill_uniform_buffer(scale_inv_rowwise_->cpu_buffer(), scale_inv_rowwise_->size(), + scale_inv_rowwise_->dtype(), gen_); + scale_inv_rowwise_->from_cpu(); +} + +void Tensor::fill_uniform_columnwise_scale_inv() { + if (!scale_inv_columnwise_ || scale_inv_columnwise_->size() == 0) { + return; } + fill_uniform_buffer(scale_inv_columnwise_->cpu_buffer(), scale_inv_columnwise_->size(), + scale_inv_columnwise_->dtype(), gen_); + scale_inv_columnwise_->from_cpu(); } -void Tensor::shareFP8Meta(const Tensor &other) { - if ((isFp8Type(dtype()) && isFp8Type(other.dtype())) - || isFp4Type(dtype()) && isFp4Type(other.dtype())) { - auto new_tensor = TensorWrapper(other.tensor_.scaling_mode()); - auto my_rowwise_data = tensor_.get_rowwise_data(); - new_tensor.set_rowwise_data(my_rowwise_data.data_ptr, static_cast(my_rowwise_data.dtype), - my_rowwise_data.shape); - auto my_columnwise_data = tensor_.get_columnwise_data(); - new_tensor.set_columnwise_data(my_columnwise_data.data_ptr, - static_cast(my_columnwise_data.dtype), - my_columnwise_data.shape); - auto other_amax = other.tensor_.get_amax(); - new_tensor.set_amax(other_amax.data_ptr, static_cast(other_amax.dtype), - other_amax.shape); - auto other_scale = other.tensor_.get_scale(); - new_tensor.set_scale(other_scale.data_ptr, static_cast(other_scale.dtype), - other_scale.shape); - auto other_row_scale_inv = other.tensor_.get_rowwise_scale_inv(); - new_tensor.set_rowwise_scale_inv(other_row_scale_inv.data_ptr, - static_cast(other_row_scale_inv.dtype), - other_row_scale_inv.shape); - auto other_col_scale_inv = other.tensor_.get_columnwise_scale_inv(); - new_tensor.set_columnwise_scale_inv(other_col_scale_inv.data_ptr, - static_cast(other_col_scale_inv.dtype), - other_col_scale_inv.shape); - tensor_ = std::move(new_tensor); - to_cpu(); +void Tensor::fill_uniform_scale() { + if (!scale_ || scale_->size() == 0) { + return; + } + + // Generate random scales on CPU + auto *cpu_data = scale_->cpu_buffer(); + const auto numel = scale_->size(); + NVTE_CHECK(scale_->dtype() == DType::kFloat32); + std::uniform_real_distribution dis(-2.0, 1.0); + for (size_t i = 0; i < numel; ++i) { + cpu_data[i] = dis(gen_); } + + // Update GPU tensor + scale_->from_cpu(); } using std::to_string; @@ -689,7 +584,7 @@ std::vector unravel(const size_t i, const NVTEShape &shape) { return ret; } -void compareResults_sequential(const std::string &name, const Tensor &test, +void compareResults_sequential(const std::string &name, Tensor &test, const void *ref, const bool rowwise, double atol, double rtol, bool if_on_gpus, const size_t tolerable_mismatches_limit) { @@ -779,7 +674,7 @@ static size_t getFirstMismatchIdx(const DType data_type, const T* test_data, con return first_mismatch_idx; } -void compareResults_parallel(const std::string &name, const Tensor &test, const void *ref, +void compareResults_parallel(const std::string &name, Tensor &test, const void *ref, const bool rowwise, double atol, double rtol, bool if_on_gpus, const size_t tolerable_mismatches_limit) { if (if_on_gpus) test.to_cpu(); @@ -806,7 +701,7 @@ void compareResults_parallel(const std::string &name, const Tensor &test, const ); } -void compareResults(const std::string &name, const Tensor &test, const void *ref, +void compareResults(const std::string &name, Tensor &test, const void *ref, const bool rowwise, double atol, double rtol, bool if_on_gpus, const size_t tolerable_mismatches_limit) { constexpr bool sequential = false; @@ -992,6 +887,7 @@ void generate_data_uniformly(T* data, const size_t size, std::mt19937* gen) { } void fillUniform(Tensor *t) { + // Generate random row-wise data and column-wise data if (t->rowwise()) { const size_t size = product(t->rowwise_shape()); TRANSFORMER_ENGINE_TYPE_SWITCH_ALL(t->dtype(), T, @@ -1009,8 +905,12 @@ void fillUniform(Tensor *t) { } ); } - std::uniform_real_distribution<> dis(-2.0, 1.0); - t->set_scale_inv(dis(t->gen())); + + // Generate random scales + t->fill_uniform_rowwise_scale_inv(); + t->fill_uniform_columnwise_scale_inv(); + + // Update data on GPU t->from_cpu(); } @@ -1046,7 +946,20 @@ void fillCase_special(Tensor *t) { } }); } - t->set_scale_inv(1.0); + + // Fill scales + if (t->scaling_mode() == NVTE_DELAYED_TENSOR_SCALING) { + if (isFp8Type(t->dtype())) { + // FP8 tensor scale is set to 1 + t->set_scale_inv(1.0); + } + } else { + // Block scales are filled randomly + t->fill_uniform_rowwise_scale_inv(); + t->fill_uniform_columnwise_scale_inv(); + } + + // Update GPU tensor data t->from_cpu(); } @@ -1080,15 +993,12 @@ template void fillCase(Tensor *t, const InputsFillCase fill_case); #endif void setRandomScale(Tensor *t) { - std::uniform_real_distribution<> dis(-2.0, 1.0); - const float scale = dis(t->gen()); - t->set_scale(scale); + t->fill_uniform_scale(); } void setRandomScaleInv(Tensor *t) { - std::uniform_real_distribution<> dis(-2.0, 1.0); - const float scale_inv = dis(t->gen()); - t->set_scale_inv(scale_inv); + t->fill_uniform_rowwise_scale_inv(); + t->fill_uniform_columnwise_scale_inv(); } bool isFp8Type(DType type) { diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index b2a7da89c..17f36a99d 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -6,10 +6,12 @@ #pragma once -#include -#include #include +#include +#include #include +#include + #include #define FP4_TYPE_SUPPORTED (CUDA_VERSION >= 12080) @@ -27,6 +29,11 @@ namespace test { using namespace transformer_engine; +size_t typeToNumBits(DType type); +size_t product(const NVTEShape &shape); +size_t product(const std::vector &shape); +size_t bytes(const NVTEShape& shape, const DType type); + template struct BytesToType {}; @@ -114,9 +121,30 @@ struct TypeInfo { } constexpr static DType dtype = getType(); - constexpr static size_t size = BitsNumber::num_bits;; + constexpr static size_t size = BitsNumber::num_bits; +}; + +// Deleter for CUDA buffer RAII class +struct CudaDeleter { + void operator()(void* ptr) const { if (ptr != nullptr) cudaFree(ptr); } }; +// CUDA buffer RAII class +template +using CudaPtr = std::unique_ptr; + +// Construct CUDA memory +template +CudaPtr cuda_alloc(size_t bytes) { + void* ptr = nullptr; + NVTE_CHECK_CUDA(cudaMalloc(&ptr, bytes)); + return CudaPtr(static_cast(ptr)); +} + +/* Wrapper for Transformer Engine tensor + * + * Maintains matching GPU and CPU buffers. + */ class Tensor { public: Tensor(const std::string& name, @@ -133,7 +161,7 @@ class Tensor { const NVTEScalingMode &mode = NVTE_DELAYED_TENSOR_SCALING) : Tensor(name, nvte_make_shape(shape.data(), shape.size()), type, rowwise, columnwise, mode) {} - Tensor() {} + Tensor() = default; Tensor& operator=(const Tensor &other) = delete; Tensor(const Tensor &other) = delete; @@ -141,42 +169,7 @@ class Tensor { Tensor(Tensor &&other) = default; Tensor& operator=(Tensor &&other) = default; - ~Tensor() { - void *data_ptr = tensor_.dptr(); - void *scale_inv = tensor_.scale_inv(); - void *columnwise_data_ptr = tensor_.get_columnwise_data().data_ptr; - void *columnwise_scale_inv = tensor_.get_columnwise_scale_inv().data_ptr; - void *amax = tensor_.amax(); - void *columnwise_amax_ptr = tensor_.get_columnwise_amax().data_ptr; - void *scale = tensor_.scale(); - if (columnwise_data_ptr == data_ptr) { - columnwise_data_ptr = nullptr; - } - if (columnwise_scale_inv == scale_inv) { - columnwise_scale_inv = nullptr; - } - if (data_ptr != nullptr) { - cudaFree(data_ptr); - } - if (scale_inv != nullptr) { - cudaFree(scale_inv); - } - if (columnwise_data_ptr != nullptr) { - cudaFree(columnwise_data_ptr); - } - if (columnwise_scale_inv != nullptr) { - cudaFree(columnwise_scale_inv); - } - if (amax != nullptr) { - cudaFree(amax); - } - if (columnwise_amax_ptr != nullptr) { - cudaFree(columnwise_amax_ptr); - } - if (scale != nullptr) { - cudaFree(scale); - } - } + ~Tensor() = default; NVTETensor data() const noexcept { return tensor_.data(); } @@ -213,141 +206,176 @@ class Tensor { } template - T *rowwise_cpu_dptr() const { - NVTE_CHECK(TypeInfo::dtype == tensor_.dtype(), "Invalid type!"); + T *rowwise_cpu_dptr() { + NVTE_CHECK(data_rowwise_, "Tensor does not have rowwise data!"); + NVTE_CHECK(TypeInfo::dtype == data_rowwise_->dtype(), "Invalid type!"); NVTE_CHECK(rowwise_, "Tensor does not have rowwise data!"); - return reinterpret_cast(cpu_data_rowwise_.get()); + return data_rowwise_->cpu_buffer(); } template - T *columnwise_cpu_dptr() const { - NVTE_CHECK(TypeInfo::dtype == tensor_.dtype(), "Invalid type!"); + T *columnwise_cpu_dptr() { + NVTE_CHECK(data_columnwise_, "Tensor does not have columnwise data!"); + NVTE_CHECK(TypeInfo::dtype == data_columnwise_->dtype(), "Invalid type!"); NVTE_CHECK(columnwise_, "Tensor does not have columnwise data!"); - return reinterpret_cast(cpu_data_columnwise_.get()); + return data_columnwise_->cpu_buffer(); } - float amax() const { - if(amax_cpu_data_) { - to_cpu(); - return *amax_cpu_data_; - } else { - return 0; - } + float amax() { + NVTE_CHECK(amax_rowwise_); + NVTE_CHECK(amax_rowwise_->size() == 1); + NVTE_CHECK(amax_rowwise_->dtype() == DType::kFloat32); + amax_rowwise_->to_cpu(); + return *amax_rowwise_->cpu_buffer(); } - float amax_columnwise() const { - if(amax_cpu_data_columnwise_) { - to_cpu(); - return *amax_cpu_data_columnwise_; - } else { - return 0; - } + float amax_columnwise() { + NVTE_CHECK(amax_columnwise_); + NVTE_CHECK(amax_columnwise_->size() == 1); + NVTE_CHECK(amax_columnwise_->dtype() == DType::kFloat32); + amax_columnwise_->to_cpu(); + return *amax_columnwise_->cpu_buffer(); } - float scale() const { - if(scale_cpu_data_) { - NVTE_CHECK(tensor_.scaling_mode() == NVTE_DELAYED_TENSOR_SCALING, "Invalid scaling_mode!"); - to_cpu(); - return *scale_cpu_data_; - } else { - return 1; - } + float scale() { + NVTE_CHECK(scale_); + NVTE_CHECK(scale_->size() == 1); + NVTE_CHECK(scale_->dtype() == DType::kFloat32); + scale_->to_cpu(); + return *scale_->cpu_buffer(); + } + + float rowwise_scale_inv(){ + NVTE_CHECK(scale_inv_rowwise_); + NVTE_CHECK(scale_inv_rowwise_->size() == 1); + NVTE_CHECK(scale_inv_rowwise_->dtype() == DType::kFloat32); + scale_inv_rowwise_->to_cpu(); + return *scale_inv_rowwise_->cpu_buffer(); } template T *rowwise_cpu_scale_inv_ptr(){ - if (tensor_.scaling_mode() == NVTE_DELAYED_TENSOR_SCALING){ - NVTE_CHECK(TypeInfo::dtype == DType::kFloat32, "Invalid type!"); - } else if (tensor_.scaling_mode() == NVTE_BLOCK_SCALING_1D || tensor_.scaling_mode() == NVTE_BLOCK_SCALING_2D) { - NVTE_CHECK(TypeInfo::dtype == DType::kFloat32, "Invalid type!"); - } else if (tensor_.scaling_mode() == NVTE_NVFP4_1D_SCALING) { - NVTE_CHECK(TypeInfo::dtype == DType::kFloat8E4M3, "Invalid type!"); - } else { - NVTE_CHECK(TypeInfo::dtype == DType::kByte, "Invalid type!"); - } - to_cpu(); - return reinterpret_cast(rowwise_scale_inv_cpu_data_.get()); + NVTE_CHECK(scale_inv_rowwise_); + scale_inv_rowwise_->to_cpu(); + return scale_inv_rowwise_->cpu_buffer(); } template T *columnwise_cpu_scale_inv_ptr(){ - if (tensor_.scaling_mode() == NVTE_DELAYED_TENSOR_SCALING){ - NVTE_CHECK(TypeInfo::dtype == DType::kFloat32, "Invalid type!"); - } else if (tensor_.scaling_mode() == NVTE_BLOCK_SCALING_1D || tensor_.scaling_mode() == NVTE_BLOCK_SCALING_2D) { - NVTE_CHECK(TypeInfo::dtype == DType::kFloat32, "Invalid type!"); - } else if (tensor_.scaling_mode() == NVTE_NVFP4_1D_SCALING) { - NVTE_CHECK(TypeInfo::dtype == DType::kFloat8E4M3, "Invalid type!"); - } else { - NVTE_CHECK(TypeInfo::dtype == DType::kByte, "Invalid type!"); - } - to_cpu(); - return reinterpret_cast(columnwise_scale_inv_cpu_data_.get()); + NVTE_CHECK(scale_inv_columnwise_); + scale_inv_columnwise_->to_cpu(); + return scale_inv_columnwise_->cpu_buffer(); } - float rowwise_scale_inv(){ - if(rowwise_scale_inv_cpu_data_) { - float scale_inv = rowwise_cpu_scale_inv_ptr()[0]; - return scale_inv; - } else { - return 1; - } - } - - bool rowwise() const { - return rowwise_; + template + T *cpu_rowwise_amax_ptr() { + NVTE_CHECK(amax_rowwise_); + amax_rowwise_->to_cpu(); + return amax_rowwise_->cpu_buffer(); } - bool columnwise() const { - return columnwise_; + template + T *cpu_columnwise_amax_ptr() { + NVTE_CHECK(amax_columnwise_); + amax_columnwise_->to_cpu(); + return amax_columnwise_->cpu_buffer(); } - void set_tensor_amax(const float amax) { - if (amax_cpu_data_) { - *amax_cpu_data_ = amax; - from_cpu(); - } + size_t rowwise_amax_size() const noexcept { + return amax_rowwise_ ? amax_rowwise_->size() : 0; } - void set_tensor_amax_columnwise(const float amax) { - if (amax_cpu_data_columnwise_) { - *amax_cpu_data_columnwise_ = amax; - from_cpu(); - } + bool rowwise() const { + return rowwise_; } - void set_tensor_amax_nullptr(){ - tensor_.set_amax(nullptr, DType::kFloat32, tensor_.defaultShape); + bool columnwise() const { + return columnwise_; } - void set_tensor_amax_shape(const std::vector &shape); - std::vector tensor_amax_values() const; - void copy_tensor_amax_from(const Tensor &other); + void set_tensor_amax_nullptr(); - void set_with_gemm_swizzled_scales(bool with_gemm_swizzled_scales){ - tensor_.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); - } + void set_with_gemm_swizzled_scales(bool with_gemm_swizzled_scales); + void set_row_scaled_nvfp4(bool row_scaled_nvfp4); - void set_row_scaled_nvfp4(bool row_scaled_nvfp4) { - tensor_.set_row_scaled_nvfp4(row_scaled_nvfp4); - } + void to_cpu(); + void from_cpu(); - void to_cpu() const; - void from_cpu() const; + void set_amax(float amax); void set_scale(float scale); void set_scale_inv(float scale_inv); - void shareFP8Meta(const Tensor &other); + void set_tensor_amax_columnwise(float amax); + + void fill_uniform_rowwise_scale_inv(); + void fill_uniform_columnwise_scale_inv(); + void fill_uniform_scale(); std::mt19937& gen() { return gen_; } private: + + /* Manages matching GPU and CPU buffers. */ + class Buffer { + public: + + Buffer(size_t size = 0, DType dtype = DType::kByte); + ~Buffer() = default; + Buffer(const Buffer&) = delete; + Buffer& operator=(const Buffer&) = delete; + Buffer(Buffer&&) = default; + Buffer& operator=(Buffer&&) = default; + + size_t size() const noexcept { return size_; } + DType dtype() const noexcept { return dtype_; } + + // Void pointer accessors + void *cpu_buffer() { return cpu_buffer_.get(); } + const void *cpu_buffer() const { return cpu_buffer_.get(); } + void *gpu_buffer() { return gpu_buffer_.get(); } + const void *gpu_buffer() const { return gpu_buffer_.get(); } + + // Templated pointer accessors + template + T *cpu_buffer() { + return reinterpret_cast(cpu_buffer()); + } + template + const T *cpu_buffer() const { + return const_cast(this)->cpu_buffer(); + } + template + T *gpu_buffer() { + return reinterpret_cast(gpu_buffer()); + } + template + const T *gpu_buffer() const { + return const_cast(this)->gpu_buffer(); + } + + // Memory transfers between CPU and GPU + void to_cpu(); + void from_cpu(); + + private: + std::unique_ptr cpu_buffer_; + CudaPtr gpu_buffer_; + size_t size_; + DType dtype_; + size_t bytes_; + }; + + // Transformer Engine tensor TensorWrapper tensor_; - std::unique_ptr cpu_data_rowwise_; - std::unique_ptr cpu_data_columnwise_; - std::shared_ptr amax_cpu_data_; - std::shared_ptr amax_cpu_data_columnwise_; - std::shared_ptr scale_cpu_data_; - std::unique_ptr rowwise_scale_inv_cpu_data_; - std::unique_ptr columnwise_scale_inv_cpu_data_; + + // Data buffers + std::optional data_rowwise_; + std::optional data_columnwise_; + std::shared_ptr scale_inv_rowwise_; + std::shared_ptr scale_inv_columnwise_; + std::optional amax_rowwise_; + std::optional amax_columnwise_; + std::optional scale_; + bool rowwise_; bool columnwise_; std::string name_; @@ -497,17 +525,12 @@ inline float dsilu(const float x) { return x * dsigmoid(x) + sigmoid(x); } inline float srelu(const float x) { return x > 0 ? x * x : 0; } inline float dsrelu(const float x) { return fmaxf(0, 2 * x); } -size_t typeToNumBits(DType type); -size_t product(const NVTEShape &shape); -size_t product(const std::vector &shape); -size_t bytes(const NVTEShape& shape, const DType type); - size_t first_dimension(const std::vector &shape); size_t last_dimension(const std::vector &shape); bool areShapesEqual(const NVTEShape &s1, const NVTEShape &s2); -void compareResults(const std::string &name, const Tensor &test, const void *ref, +void compareResults(const std::string &name, Tensor &test, const void *ref, bool rowwise, double atol = 1e-5, double rtol = 1e-8, bool if_on_gpus = true, const size_t tolerable_mismatches_limit = 0); void compareResults(const std::string &name, const float test, const float ref, @@ -550,26 +573,14 @@ int32_t getDeviceComputeCapability(); constexpr int32_t hopperComputeCapability = 90; constexpr int32_t blackwellComputeCapability = 100; -// Custom deleters for RAII -struct CudaDeleter { - void operator()(void* p) const { if (p) cudaFree(p); } -}; +// Custom deleter for RAII struct GroupedTensorDeleter { void operator()(NVTEGroupedTensor h) const { if (h) nvte_destroy_grouped_tensor(h); } }; -template -using CudaPtr = std::unique_ptr; +// Grouped tensor RAII class using GroupedTensorHandle = std::unique_ptr, GroupedTensorDeleter>; -// Helper to allocate CUDA memory into a CudaPtr -template -CudaPtr cuda_alloc(size_t bytes) { - void* ptr = nullptr; - NVTE_CHECK_CUDA(cudaMalloc(&ptr, bytes)); - return CudaPtr(static_cast(ptr)); -} - // Helper owning GPU buffers that back NVTEGroupedTensor. // NVTEGroupedTensor does not own memory; data/offsets/scales // must be allocated and freed by the test. diff --git a/tests/cpp_distributed/test_comm_gemm.cu b/tests/cpp_distributed/test_comm_gemm.cu index cc0d760a3..45f666456 100644 --- a/tests/cpp_distributed/test_comm_gemm.cu +++ b/tests/cpp_distributed/test_comm_gemm.cu @@ -107,8 +107,10 @@ std::vector CopyMatrix(const std::vector& data, size_t mstart, size_t nsta template test::Tensor Make(size_t m, size_t n, float scale) { test::Tensor ret("", std::vector{n, m}, TypeInfo::dtype); - ret.set_scale(scale); - ret.set_scale_inv(1.0 / scale); + if (test::isFp8Type(TypeInfo::dtype)) { + ret.set_scale(scale); + ret.set_scale_inv(1.0 / scale); + } return ret; } @@ -116,8 +118,10 @@ template test::Tensor MakeFromData(const std::vector& data, size_t mstart, size_t nstart, size_t msize, size_t nsize, size_t ld, float scale) { test::Tensor ret("", std::vector{nsize, msize}, TypeInfo::dtype); - ret.set_scale(scale); - ret.set_scale_inv(1.0 / scale); + if (test::isFp8Type(TypeInfo::dtype)) { + ret.set_scale(scale); + ret.set_scale_inv(1.0 / scale); + } auto local = CopyMatrix(data, mstart, nstart, msize, nsize, ld); NVTE_CHECK_CUDA(cudaMemcpy(ret.rowwise_dptr(), local.data(), local.size() * sizeof local[0], cudaMemcpyDefault)); diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index e9a6f4f73..488f25915 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -72,7 +72,17 @@ enum NVTETensorParam { kNVTEColumnwiseScaleInv = 5, /*!< Scale inverse tensor for decoding Columnwise Data */ kNVTEColumnwiseAmax = 6, /*!< Columnwise Amax tensor */ kNVTEWithGEMMSwizzledScales = 7, /*!< Whether scaling factors are in format expected by GEMM */ - kNVTERowScaledNVFP4 = 8, /*!< Whether an NVFP4 tensor uses row scaling */ + /*! Whether an NVFP4 tensor uses row scaling instead of tensor scaling. + * + * Column-wise data is not supported with row scaling. + * + * Row scaling affects the interpretation of the amax tensor. With + * tensor scaling, the amax tensor is a single FP32 that must be + * computed prior to quantization. With row scaling, the amax + * tensor size is the number of tensor rows (flattened to 2D), and + * its values are populated during quantization. + */ + kNVTERowScaledNVFP4 = 8, kNVTENumTensorParams }; From d73bfa14fb26d186ed8e5eeaad8ef5c1d728cc86 Mon Sep 17 00:00:00 2001 From: Evgeny Tsykunov Date: Mon, 11 May 2026 19:26:39 +0200 Subject: [PATCH 047/180] [PyTorch] Introduce QuantizerRole (#2620) * Enable semantic roles emitted by module/op and comsumed by custom recipe state Signed-off-by: Evgeny * Update quantization factories Signed-off-by: Evgeny * Fix tests Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Swap tensor:module Signed-off-by: Evgeny * Better naming Signed-off-by: Evgeny * Introduce QuantizerRole frozen data class instead of a string Signed-off-by: Evgeny * Shrink module_type vocabulary Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix numerics exact test Signed-off-by: Evgeny * Set defaults, make custom recipe forward compatible Signed-off-by: Evgeny * remove position from QuantizerRole Signed-off-by: Evgeny * Set good defaults Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Resolve naming: make every module/op distinguishable via name Signed-off-by: Evgeny * Configure output/grad_input roles, defaults to None Signed-off-by: Evgeny * Remove is_gemm() Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Enable base recipes via CustomRecipe and quantization factories Signed-off-by: Evgeny * Add factory example - NVFP4 for Linear, MXFP8 for GroupedLinear Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix custom recipe test Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Test fine-grained quantization targets Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Add quantizer roles for attention (attn is wip) Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Enable statful recipes in the Custom recipe - Delayed Scaling support Signed-off-by: Evgeny * Fix save_original_input for custom delayed scaling Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Enable custom recipe for attn Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make boundary role setting more explicit in MHA Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make dpa role setting more intuitive Signed-off-by: Evgeny * Docstring for get_quantizer_roles() in base module Signed-off-by: Evgeny * Fix lint Signed-off-by: Evgeny * Restrict None roles Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Linter Signed-off-by: Evgeny * Minor fixes Signed-off-by: Evgeny * Test debug tools compat Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix pylint Signed-off-by: Evgeny * fix test Signed-off-by: Evgeny * Fix lint Signed-off-by: Evgeny * Constructor takes roles kwarg + test fix Signed-off-by: Evgeny * Constructor takes roles kwarg + test fix (quantization.py) Signed-off-by: Evgeny * Fix attention: MXFP8, w/o CP Signed-off-by: Evgeny * Add test custom recipe Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make Float8BlockScalingRecipeState and NVFP4BlockScalingRecipeState aware about QuantizerRole, dispatch on that + positional fallback if get_quantizer_roles() is not defined by the module/op Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix linter Signed-off-by: Evgeny * Fix CI Signed-off-by: Evgeny * Preserve delayed scaling state (buffers) when rebuild is triggered Signed-off-by: Evgeny * Fix test, minor Signed-off-by: Evgeny * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix distributed tests Signed-off-by: Evgeny --------- Signed-off-by: Evgeny Signed-off-by: Evgeny Tsykunov Signed-off-by: Evgeny Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Evgeny --- qa/L0_pytorch_unittest/test.sh | 1 + .../pytorch/distributed/run_numerics_exact.py | 55 +- tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 2 +- .../nvfp4/test_nvfp4_group_quantize.py | 2 +- .../test_nvfp4_group_quantize_graph_safe.py | 2 +- .../pytorch/nvfp4/test_nvfp4_module_exact.py | 51 +- .../nvfp4/test_nvfp4_quantize_exact.py | 2 +- .../nvfp4/test_nvfp4_rht_quantize_exact.py | 2 +- tests/pytorch/test_custom_recipe.py | 1588 ++++++++++++++++- .../test_float8_current_scaling_exact.py | 2 +- transformer_engine/common/recipe/__init__.py | 37 +- transformer_engine/pytorch/__init__.py | 3 + .../dot_product_attention/backends.py | 54 +- .../dot_product_attention/context_parallel.py | 32 +- .../dot_product_attention.py | 130 +- .../attention/dot_product_attention/utils.py | 31 +- .../pytorch/attention/multi_head_attention.py | 83 +- .../quantization_factory_examples.py | 271 +++ .../quantization_recipes_base.py | 179 ++ ...py => quantization_ref_current_scaling.py} | 15 +- ...ion_nvfp4.py => quantization_ref_nvfp4.py} | 35 +- transformer_engine/pytorch/module/base.py | 197 +- .../pytorch/module/grouped_linear.py | 42 +- .../pytorch/module/layernorm_linear.py | 31 +- .../pytorch/module/layernorm_mlp.py | 52 +- transformer_engine/pytorch/module/linear.py | 31 +- .../pytorch/ops/basic/basic_linear.py | 17 +- transformer_engine/pytorch/ops/op.py | 21 +- transformer_engine/pytorch/quantization.py | 712 ++++++-- 29 files changed, 3364 insertions(+), 316 deletions(-) create mode 100644 transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py create mode 100644 transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py rename transformer_engine/pytorch/custom_recipes/{quantization_current_scaling.py => quantization_ref_current_scaling.py} (98%) rename transformer_engine/pytorch/custom_recipes/{quantization_nvfp4.py => quantization_ref_nvfp4.py} (98%) diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index 22636828f..c35dc4c06 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -26,6 +26,7 @@ pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/test_sanity.py || test_fail "test_sanity.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_recipe.xml $TE_PATH/tests/pytorch/test_recipe.py || test_fail "test_recipe.py" +python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_custom_recipe.xml $TE_PATH/tests/pytorch/test_custom_recipe.py || test_fail "test_custom_recipe.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_deferred_init.xml $TE_PATH/tests/pytorch/test_deferred_init.py || test_fail "test_deferred_init.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/test_numerics.py || test_fail "test_numerics.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cuda_graphs.xml $TE_PATH/tests/pytorch/test_cuda_graphs.py || test_fail "test_cuda_graphs.py" diff --git a/tests/pytorch/distributed/run_numerics_exact.py b/tests/pytorch/distributed/run_numerics_exact.py index 0f3d2cbbf..15ae2dae6 100644 --- a/tests/pytorch/distributed/run_numerics_exact.py +++ b/tests/pytorch/distributed/run_numerics_exact.py @@ -22,7 +22,7 @@ ) from transformer_engine.pytorch import NVFP4Quantizer from transformer_engine.pytorch.constants import NVFP4_BLOCK_SCALING_SIZE -from transformer_engine.pytorch.custom_recipes import quantization_nvfp4 +from transformer_engine.pytorch.custom_recipes import quantization_ref_nvfp4 from transformer_engine.pytorch.custom_recipes import utils from run_layer_with_overlap import _compare_tensors @@ -52,44 +52,39 @@ def get_nvfp4_quantizer_factory(): """ Create a quantizer factory for NVFP4 reference implementation. - This factory returns NVFP4QuantizerRef instances with RHT and 2D quantization - enabled. + Linear/grouped-linear weight slots get 2D (16x16) quantization without RHT; + every other slot (input, gradient, boundary slots with ``role is None``, + and any unknown tensor type) gets 1D (1x16) quantization with RHT. + + Mirrors the canonical "branch on what we care about, default fall-through" + pattern from + ``transformer_engine.pytorch.custom_recipes.quantization_recipes_base``; + every slot gets a real :class:`NVFP4QuantizerRef` (``CustomRecipeState`` + rejects ``None`` returns). Returns: - A factory function that takes a role string and returns a quantizer instance + A factory function that takes a QuantizerRole and returns a quantizer instance """ def factory(role): - if role == "linear_input": - return quantization_nvfp4.NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, - quant_tile_shape=(1, 16), - pow_2_scales=False, - with_rht=True, # RHT enabled for input - ) - elif role == "linear_weight": - return quantization_nvfp4.NVFP4QuantizerRef( + is_weight = ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type == "weight" + ) + if is_weight: + return quantization_ref_nvfp4.NVFP4QuantizerRef( dtype=utils.Fp4Formats.E2M1, - quant_tile_shape=(16, 16), # 2D quantization for weight + quant_tile_shape=(16, 16), pow_2_scales=False, with_rht=False, ) - elif role == "linear_output": - # Output quantization not used - return None - elif role == "linear_grad_output": - return quantization_nvfp4.NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, - quant_tile_shape=(1, 16), - pow_2_scales=False, - with_rht=True, # RHT enabled for grad_output - ) - elif role == "linear_grad_input": - # Grad input quantization not used - return None - else: - # For any other roles, return None - return None + return quantization_ref_nvfp4.NVFP4QuantizerRef( + dtype=utils.Fp4Formats.E2M1, + quant_tile_shape=(1, 16), + pow_2_scales=False, + with_rht=True, + ) return factory diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index b93933627..a7ea4f089 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -9,7 +9,7 @@ from transformer_engine.pytorch.constants import TE_DType from transformer_engine.pytorch import NVFP4Quantizer from transformer_engine.pytorch.cpp_extensions import general_gemm, general_grouped_gemm -from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef from transformer_engine.pytorch.custom_recipes import utils diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py index 7bf288fff..20a91bf6f 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize.py @@ -13,7 +13,7 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef from transformer_engine.pytorch.custom_recipes import utils from transformer_engine.pytorch.constants import TE_DType from transformer_engine.common.recipe import NVFP4BlockScaling diff --git a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py index cf2ae50ee..d46a87469 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py +++ b/tests/pytorch/nvfp4/test_nvfp4_group_quantize_graph_safe.py @@ -6,7 +6,7 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef from transformer_engine.pytorch.custom_recipes import utils from transformer_engine.pytorch.constants import TE_DType from transformer_engine.common.recipe import NVFP4BlockScaling diff --git a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py index a96fea3af..b57b78eb1 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_module_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_module_exact.py @@ -6,7 +6,7 @@ import torch import transformer_engine.pytorch as te from transformer_engine.common import recipe -from transformer_engine.pytorch.custom_recipes import quantization_nvfp4 +from transformer_engine.pytorch.custom_recipes import quantization_ref_nvfp4 from transformer_engine.pytorch.custom_recipes import utils @@ -76,40 +76,37 @@ def get_nvfp4_quantizer_factory(with_rht: bool = False, with_2d_quantization: bo with_2d_quantization: Whether to use 2D quantization (16x16 tiles for weights) Returns: - A factory function that takes a role string and returns a quantizer instance + A factory function that takes a QuantizerRole (or None for boundary slots) + and returns a quantizer instance. """ + # Boundary slots (output, grad_input) get role=None from Linear.get_quantizer_roles + # when no consumer is configured. CustomRecipeState rejects None returns from + # qfactory, so we return a valid quantizer for those slots; it is harmless because + # the GEMM outputs in the high-precision activation dtype, not in NVFP4. + def _default_quantizer(): + return quantization_ref_nvfp4.NVFP4QuantizerRef( + dtype=utils.Fp4Formats.E2M1, + quant_tile_shape=(1, 16), + pow_2_scales=False, + with_rht=with_rht, + ) + def factory(role): - if role == "linear_input": - return quantization_nvfp4.NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, - quant_tile_shape=(1, 16), - pow_2_scales=False, - with_rht=with_rht, - ) - elif role == "linear_weight": - return quantization_nvfp4.NVFP4QuantizerRef( + if role is None: + return _default_quantizer() + if role.tensor_type == "input": + return _default_quantizer() + if role.tensor_type == "weight": + return quantization_ref_nvfp4.NVFP4QuantizerRef( dtype=utils.Fp4Formats.E2M1, quant_tile_shape=(16, 16) if with_2d_quantization else (1, 16), pow_2_scales=False, with_rht=False, ) - elif role == "linear_output": - # Output quantization not used - return None - elif role == "linear_grad_output": - return quantization_nvfp4.NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, - quant_tile_shape=(1, 16), - pow_2_scales=False, - with_rht=with_rht, - ) - elif role == "linear_grad_input": - # Grad input quantization not used - return None - else: - # For any other roles, return None - return None + if role.tensor_type == "grad_output": + return _default_quantizer() + return _default_quantizer() return factory diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 0824a5e7b..53569d90d 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -7,7 +7,7 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef from transformer_engine.pytorch.custom_recipes import utils from transformer_engine.common.recipe import NVFP4BlockScaling from transformer_engine.pytorch.constants import TE_DType diff --git a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py index 795721df0..2d159dbf6 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_rht_quantize_exact.py @@ -12,7 +12,7 @@ import transformer_engine.pytorch as te import transformer_engine_torch as tex from transformer_engine.pytorch import NVFP4Quantizer -from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import NVFP4QuantizerRef +from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import NVFP4QuantizerRef from transformer_engine.pytorch.custom_recipes import utils from transformer_engine.pytorch.constants import TE_DType from transformer_engine.common.recipe import NVFP4BlockScaling diff --git a/tests/pytorch/test_custom_recipe.py b/tests/pytorch/test_custom_recipe.py index 536d43adc..62a629179 100644 --- a/tests/pytorch/test_custom_recipe.py +++ b/tests/pytorch/test_custom_recipe.py @@ -17,8 +17,16 @@ GroupedLinear, Float8CurrentScalingQuantizer, ) +from transformer_engine.pytorch.quantization import QuantizerRole import transformer_engine.pytorch.ops as te_ops -from transformer_engine.pytorch.custom_recipes.quantization_nvfp4 import ( +from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( + current_scaling_quantizer_factory, + mxfp8_quantizer_factory, + float8_block_scaling_quantizer_factory, + nvfp4_quantizer_factory, + delayed_scaling_quantizer_factory, +) +from transformer_engine.pytorch.custom_recipes.quantization_ref_nvfp4 import ( nvfp4_ref_rht_2d_quantizer_factory, ) @@ -91,9 +99,9 @@ def test_custom_recipe_sanity(module_type): # Single factory: map roles to quantizers def quantizer_factory(role): - if role in ("linear_input", "linear_weight", "linear_output"): + if role is None: return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") - if role in ("linear_grad_output", "linear_grad_input"): + if role.tensor_type == "grad_output": return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") @@ -119,18 +127,18 @@ def test_custom_recipe_grouped_linear_sanity(): num_gemms = 3 in_features = 64 out_features = 64 - batch = 32 - base = batch // num_gemms - rem = batch % num_gemms - m_splits = [base + (1 if i < rem else 0) for i in range(num_gemms)] + # Each per-GEMM M dim must be a multiple of 16 to satisfy cuBLAS FP8 GEMM's + # leading-dimension alignment requirement on Hopper (sm_90). + m_splits = [16] * num_gemms + batch = sum(m_splits) model = GroupedLinear(num_gemms, in_features, out_features, params_dtype=torch.bfloat16).cuda() inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True) def quantizer_factory(role): - if role in ("linear_input", "linear_weight", "linear_output"): + if role is None: return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") - if role in ("linear_grad_output", "linear_grad_input"): + if role.tensor_type == "grad_output": return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") @@ -190,9 +198,9 @@ def test_custom_recipe_matches_current_scaling(): # Custom: single factory returning quantizers per role to match Float8CurrentScaling def quantizer_factory(role): - if role in ("linear_input", "linear_weight", "linear_output"): + if role is None: return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") - if role in ("linear_grad_output", "linear_grad_input"): + if role.tensor_type == "grad_output": return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") @@ -210,7 +218,7 @@ def quantizer_factory(role): assert cus_fwd_w.dtype == tex.DType.kFloat8E4M3 assert cus_fwd_out.dtype == tex.DType.kFloat8E4M3 assert cus_bwd_go.dtype == tex.DType.kFloat8E5M2 - assert cus_bwd_gi.dtype == tex.DType.kFloat8E5M2 + assert cus_bwd_gi.dtype == tex.DType.kFloat8E4M3 # role=None fallback loss_custom = (out_custom.float() * scale.view(1, -1)).sum() loss_custom.backward() @@ -247,9 +255,9 @@ def test_custom_recipe_ops_linear_2_1_layout(): inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True) def quantizer_factory(role): - if role in ("linear_input", "linear_weight", "linear_output"): + if role is None: return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") - if role in ("linear_grad_output", "linear_grad_input"): + if role.tensor_type == "grad_output": return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device="cuda") return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") @@ -272,44 +280,47 @@ def test_custom_recipe_factory_invocation_counts_and_cycling(): in_features = 64 out_features = 64 - batch = 8 + # batch must be a multiple of 16 to satisfy cuBLAS FP8 GEMM's leading-dim + # alignment requirement on Hopper (sm_90). + batch = 16 op = Linear(in_features, out_features, params_dtype=torch.bfloat16) inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True) - # Counters per role + # Counters per tensor_type. The output (fwd) and grad_input (bwd) + # slots have role=None by default (unknown consumer), so we count + # those separately. counts = { - "linear_input": 0, - "linear_weight": 0, - "linear_output": 0, - "linear_grad_output": 0, - "linear_grad_input": 0, + "input": 0, + "weight": 0, + "grad_output": 0, + None: 0, } def quantizer_factory(role): - if role in counts: - counts[role] += 1 - if role in ("linear_input", "linear_weight", "linear_output"): + if role is None: + counts[None] += 1 return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device=torch.device("cuda")) - if role in ("linear_grad_output", "linear_grad_input"): + assert isinstance(role, QuantizerRole), f"Expected QuantizerRole, got {type(role)}" + assert role.module_type == "linear" + if role.tensor_type in counts: + counts[role.tensor_type] += 1 + if role.tensor_type == "grad_output": return Float8CurrentScalingQuantizer(tex.DType.kFloat8E5M2, device=torch.device("cuda")) return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device=torch.device("cuda")) custom = recipe.CustomRecipe(qfactory=quantizer_factory) - # Run fwd+bwd once; for a single GEMM, expect forward to build 3 quantizers (cycled from 1 factory), - # and backward to build 2 quantizers (cycled from 1 factory). with autocast(enabled=True, recipe=custom): out = op(inp) loss = out.float().sum() loss.backward() - # Single GEMM: forward should request input, weight, output; backward grad_output, grad_input - assert counts["linear_input"] == 1 - assert counts["linear_weight"] == 1 - assert counts["linear_output"] == 1 - assert counts["linear_grad_output"] == 1 - assert counts["linear_grad_input"] == 1 + # Forward: input, weight, output(None); backward: grad_output, grad_input(None) + assert counts["input"] == 1 + assert counts["weight"] == 1 + assert counts["grad_output"] == 1 + assert counts[None] == 2, f"Expected 2 None roles (output + grad_input), got {counts[None]}" def test_factories_return_distinct_instances_and_buffers(): @@ -317,9 +328,15 @@ def test_factories_return_distinct_instances_and_buffers(): if not torch.cuda.is_available() or not available: pytest.skip(f"FP8 unsupported on this device: {reason}") - # Two calls should produce distinct quantizer objects and distinct tensor buffers + from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer + + # Two calls should produce distinct quantizer objects with distinct + # scale/amax buffers (Float8Quantizer / delayed-scaling is the class + # that owns persistent per-quantizer state; current scaling has none). def factory(): - return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device=torch.device("cuda")) + scale = torch.ones(1, dtype=torch.float32, device="cuda") + amax = torch.zeros(1, dtype=torch.float32, device="cuda") + return Float8Quantizer(scale=scale, amax=amax, fp8_dtype=tex.DType.kFloat8E4M3) q1 = factory() q2 = factory() @@ -331,3 +348,1504 @@ def factory(): # Mutating one should not affect the other q1.scale.fill_(123.0) assert not torch.equal(q1.scale, q2.scale) + + +def _run_linear_fwd_bwd(model, inp, recipe): + """Run forward + backward with a given recipe and return (output, inp.grad, param grads).""" + with autocast(enabled=True, recipe=recipe): + out = model(inp) + loss = out.float().sum() + loss.backward() + param_grads = {n: p.grad.clone() for n, p in model.named_parameters() if p.grad is not None} + return out.clone(), inp.grad.clone(), param_grads + + +def _make_pair(in_features=128, out_features=128, batch=32, seed=42): + """Create a pair of identical Linear models and matching inputs.""" + torch.manual_seed(seed) + model_ref = Linear(in_features, out_features, params_dtype=torch.bfloat16, bias=False).cuda() + model_cus = Linear(in_features, out_features, params_dtype=torch.bfloat16, bias=False).cuda() + model_cus.load_state_dict(model_ref.state_dict()) + + base_inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16) + inp_ref = base_inp.clone().detach().requires_grad_(True) + inp_cus = base_inp.clone().detach().requires_grad_(True) + return model_ref, model_cus, inp_ref, inp_cus + + +def _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus): + """Assert exact match of outputs and all gradients.""" + assert torch.allclose( + out_ref, out_cus, rtol=0.0, atol=0.0 + ), f"Forward mismatch: max diff = {(out_ref - out_cus).abs().max()}" + assert torch.allclose( + grad_ref, grad_cus, rtol=0.0, atol=0.0 + ), f"Input grad mismatch: max diff = {(grad_ref - grad_cus).abs().max()}" + for name in pgrads_ref: + assert torch.allclose(pgrads_ref[name], pgrads_cus[name], rtol=0.0, atol=0.0), ( + f"Param grad '{name}' mismatch: max diff = " + f"{(pgrads_ref[name] - pgrads_cus[name]).abs().max()}" + ) + + +def test_factory_matches_delayed_scaling(): + """delayed_scaling_quantizer_factory should produce bit-identical results + to the built-in DelayedScaling recipe.""" + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported: {reason}") + + model_ref, model_cus, inp_ref, inp_cus = _make_pair() + + out_ref, grad_ref, pgrads_ref = _run_linear_fwd_bwd(model_ref, inp_ref, recipe.DelayedScaling()) + out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( + model_cus, inp_cus, recipe.CustomRecipe(qfactory=delayed_scaling_quantizer_factory) + ) + _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) + + +def test_factory_matches_current_scaling(): + """current_scaling_quantizer_factory should produce bit-identical results + to the built-in Float8CurrentScaling recipe.""" + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported: {reason}") + + model_ref, model_cus, inp_ref, inp_cus = _make_pair() + + out_ref, grad_ref, pgrads_ref = _run_linear_fwd_bwd( + model_ref, inp_ref, recipe.Float8CurrentScaling() + ) + out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( + model_cus, inp_cus, recipe.CustomRecipe(qfactory=current_scaling_quantizer_factory) + ) + _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) + + +def test_factory_matches_mxfp8(): + """mxfp8_quantizer_factory should produce bit-identical results + to the built-in MXFP8BlockScaling recipe.""" + available, reason = te.is_mxfp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"MXFP8 unsupported: {reason}") + + model_ref, model_cus, inp_ref, inp_cus = _make_pair() + + out_ref, grad_ref, pgrads_ref = _run_linear_fwd_bwd( + model_ref, inp_ref, recipe.MXFP8BlockScaling() + ) + out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( + model_cus, inp_cus, recipe.CustomRecipe(qfactory=mxfp8_quantizer_factory) + ) + _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) + + +def test_factory_matches_block_scaling(): + """float8_block_scaling_quantizer_factory should produce bit-identical results + to the built-in Float8BlockScaling recipe.""" + available = te.is_fp8_block_scaling_available() + if not torch.cuda.is_available() or not available: + pytest.skip("Float8 block scaling unsupported on this device") + + model_ref, model_cus, inp_ref, inp_cus = _make_pair() + + out_ref, grad_ref, pgrads_ref = _run_linear_fwd_bwd( + model_ref, inp_ref, recipe.Float8BlockScaling() + ) + out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( + model_cus, inp_cus, recipe.CustomRecipe(qfactory=float8_block_scaling_quantizer_factory) + ) + _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) + + +def test_factory_matches_nvfp4(): + """nvfp4_quantizer_factory should produce bit-identical results + to the built-in NVFP4BlockScaling recipe.""" + available = te.is_nvfp4_available() + if not torch.cuda.is_available() or not available: + pytest.skip("NVFP4 unsupported on this device") + + model_ref, model_cus, inp_ref, inp_cus = _make_pair() + + out_ref, grad_ref, pgrads_ref = _run_linear_fwd_bwd( + model_ref, inp_ref, recipe.NVFP4BlockScaling() + ) + out_cus, grad_cus, pgrads_cus = _run_linear_fwd_bwd( + model_cus, inp_cus, recipe.CustomRecipe(qfactory=nvfp4_quantizer_factory) + ) + + _assert_match(out_ref, out_cus, grad_ref, grad_cus, pgrads_ref, pgrads_cus) + + +def test_custom_recipe_quantization_targets(): + """Validate fine-grained per-module quantization targeting via QuantizerRole. + + Four transformer layers, each assembled at a different abstraction level. + The default recipe is NVFP4; specific modules are overridden: + + Layer 0 - ``TransformerLayer`` (name="tl0") -> all MXFP8 + Layer 1 - ``TransformerLayer`` (name="tl1") -> NVFP4 (default), + except fc2 overridden to MXFP8 + Layer 2 - ``MultiheadAttention`` + ``LayerNormMLP`` + (name prefix "tl2") -> NVFP4 (default), + except qkv and fc1 overridden to Float8 block-scaling + Layer 3 - Individual blocks (name prefix "tl3") -> NVFP4 (default), + except proj overridden to Float8 current-scaling + + The test validates that: + * The factory receives QuantizerRole objects with correct names + * Different quantizer types are dispatched per module + * Forward + backward complete successfully through all four layers + """ + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported on this device: {reason}") + if not te.is_mxfp8_available(): + pytest.skip("MXFP8 unsupported on this device") + if not te.is_nvfp4_available(): + pytest.skip("NVFP4 unsupported on this device") + if not te.is_fp8_block_scaling_available(): + pytest.skip("Float8 block scaling unsupported on this device") + + torch.manual_seed(42) + + H = 64 # hidden_size + FFN = 64 # ffn_hidden_size + NH = 4 # num_heads + KV = H // NH # kv_channels + B = 4 # batch + S = 8 # seq_len + common = dict(params_dtype=torch.bfloat16, bias=False) + + # Layer 0: TransformerLayer -> MXFP8 + tl0 = te.TransformerLayer( + H, + FFN, + NH, + hidden_dropout=0.0, + attention_dropout=0.0, + name="tl0", + **common, + ).cuda() + + # Layer 1: TransformerLayer -> NVFP4 default, fc2 overridden to MXFP8 + tl1 = te.TransformerLayer( + H, + FFN, + NH, + hidden_dropout=0.0, + attention_dropout=0.0, + name="tl1", + **common, + ).cuda() + + # Layer 2: MHA + LayerNormMLP -> NVFP4 default, qkv and fc1 to block-scaling + tl2_mha = te.MultiheadAttention( + H, + NH, + KV, + attention_dropout=0.0, + input_layernorm=True, + return_bias=True, + name="tl2.self_attention", + **common, + ).cuda() + tl2_mlp = LayerNormMLP(H, FFN, name="tl2.layernorm_mlp", **common).cuda() + + # Layer 3: Individual blocks with DPA -> NVFP4 default, proj to current-scaling + tl3_qkv = LayerNormLinear(H, 3 * H, name="tl3.qkv", **common).cuda() + tl3_dpa = te.DotProductAttention(NH, KV, attention_dropout=0.0, name="tl3.core_attention") + tl3_proj = Linear(H, H, name="tl3.proj", **common).cuda() + tl3_fc1 = LayerNormLinear(H, FFN, name="tl3.fc1", **common).cuda() + tl3_fc2 = Linear(FFN, H, name="tl3.fc2", **common).cuda() + + # ------------------------------------------------------------------ + # Recording + dispatching factory + # ------------------------------------------------------------------ + recorded_roles = [] + + def targeting_factory(role): + recorded_roles.append(role) + + if role is None: + return nvfp4_quantizer_factory(role) + + assert isinstance(role, QuantizerRole), f"Expected QuantizerRole, got {type(role)}" + + # Layer 0 (tl0.*): all MXFP8 + if role.name.startswith("tl0"): + return mxfp8_quantizer_factory(role) + + # Layer 1 (tl1.*): NVFP4 default, but fc2 overridden to MXFP8 + if role.name == "tl1.layernorm_mlp.fc2": + return mxfp8_quantizer_factory(role) + + # Layer 2: block scaling for qkv and fc1, rest falls through to default + if role.name == "tl2.self_attention.layernorm_linear_qkv": + return float8_block_scaling_quantizer_factory(role) + if role.name == "tl2.layernorm_mlp.fc1": + return float8_block_scaling_quantizer_factory(role) + + # Layer 3: current-scaling for proj, rest falls through to default + if role.name == "tl3.proj": + return current_scaling_quantizer_factory(role) + + # Default: NVFP4 + return nvfp4_quantizer_factory(role) + + custom_recipe = recipe.CustomRecipe(qfactory=targeting_factory) + + # ------------------------------------------------------------------ + # Forward + backward + # ------------------------------------------------------------------ + inp = torch.randn(S, B, H, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with autocast(enabled=True, recipe=custom_recipe): + # Layer 0 & 1: TransformerLayer + h = tl1(tl0(inp)) + + # Layer 2: MHA + residual + LayerNormMLP + residual + attn_out, _ = tl2_mha(h) + h = h + attn_out + h = h + tl2_mlp(h) + + # Layer 3: individual blocks with DPA + residual = h + qkv = tl3_qkv(h).view(S, B, 3, NH, KV) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + attn = tl3_dpa(q, k, v).view(S, B, H) + h = residual + tl3_proj(attn) + residual = h + h = residual + tl3_fc2(torch.nn.functional.gelu(tl3_fc1(h))) + + loss = h.float().sum() + loss.backward() + + # ------------------------------------------------------------------ + # Assertions + # ------------------------------------------------------------------ + + assert inp.grad is not None, "Input gradient is None" + + # -- Name propagation check -- + # The factory dispatches on role.name, so if a TE module fails to propagate + # names (e.g. TransformerLayer -> MHA -> LayerNormLinear) the factory would + # silently fall through to the default recipe. The quantizer-type assertions + # below would catch that too, but checking names explicitly gives a clearer + # error message pointing at the broken name rather than a wrong quantizer type. + role_names = {r.name for r in recorded_roles if r is not None} + + def _tl_names(prefix): + """Expected role names for a standard TransformerLayer with given prefix.""" + return { + f"{prefix}.self_attention.layernorm_linear_qkv", + f"{prefix}.self_attention.proj", + f"{prefix}.layernorm_mlp.fc1", + f"{prefix}.layernorm_mlp.fc2", + } + + all_expected = ( + _tl_names("tl0") + | _tl_names("tl1") + | _tl_names("tl2") + | {"tl3.qkv", "tl3.proj", "tl3.fc1", "tl3.fc2"} + ) + missing = all_expected - role_names + assert not missing, ( + f"Expected module names not seen in QuantizerRole.name: {missing}\n" + f"Recorded names: {sorted(role_names)}" + ) + + for r in recorded_roles: + if r is not None and r.module_type: + assert r.module_type in ( + "linear", + "dpa", + ), f"Unexpected module_type={r.module_type} for role {r}" + + # -- Quantizer-type checks -- + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + from transformer_engine.pytorch.tensor.float8_blockwise_tensor import Float8BlockQuantizer + + def _check_q(mod, expected_cls, label=""): + q = mod.quantizers["scaling_fwd"][tex.FP8FwdTensors.GEMM1_INPUT] + assert isinstance(q, expected_cls), ( + f"{mod.name}{' (' + label + ')' if label else ''}: " + f"expected {expected_cls.__name__}, got {type(q).__name__}" + ) + + # Layer 0: all MXFP8 + _check_q(tl0.self_attention.layernorm_qkv, MXFP8Quantizer) + _check_q(tl0.self_attention.proj, MXFP8Quantizer) + + # Layer 1: NVFP4 default, fc2 overridden to MXFP8 + _check_q(tl1.self_attention.layernorm_qkv, NVFP4Quantizer, "default") + _check_q(tl1.self_attention.proj, NVFP4Quantizer, "default") + assert any( + r is not None and r.name == "tl1.layernorm_mlp.fc2" and r.tensor_type == "input" + for r in recorded_roles + ), "tl1.layernorm_mlp.fc2 input role not recorded" + + # Layer 2: block-scaling on qkv and fc1, NVFP4 on proj and fc2 + _check_q(tl2_mha.layernorm_qkv, Float8BlockQuantizer) + _check_q(tl2_mha.proj, NVFP4Quantizer, "default") + + # Layer 3: current-scaling on proj, NVFP4 on everything else + _check_q(tl3_proj, Float8CurrentScalingQuantizer) + for mod in [tl3_qkv, tl3_fc1, tl3_fc2]: + _check_q(mod, NVFP4Quantizer, "default") + + +def test_grouped_linear_module_type_dispatch(): + """Verify GroupedLinear emits module_type='grouped_linear' so factories can + distinguish it from regular Linear (critical for MoE mixed-recipe dispatch).""" + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported on this device: {reason}") + + torch.manual_seed(0) + + num_gemms = 2 + in_features = 64 + out_features = 64 + # Each per-GEMM M dim must be a multiple of 16 to satisfy cuBLAS FP8 GEMM's + # leading-dimension alignment requirement on Hopper (sm_90). + batch = 32 + m_splits = [batch // num_gemms] * num_gemms + + model = GroupedLinear( + num_gemms, in_features, out_features, params_dtype=torch.bfloat16, name="experts" + ).cuda() + inp = torch.randn(batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + recorded_roles = [] + + def recording_factory(role): + recorded_roles.append(role) + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + + custom_recipe = recipe.CustomRecipe(qfactory=recording_factory) + + with autocast(enabled=True, recipe=custom_recipe): + out = model(inp, m_splits) + loss = out.float().sum() + loss.backward() + + non_none = [r for r in recorded_roles if r is not None] + assert len(non_none) > 0, "No QuantizerRole objects recorded" + for r in non_none: + assert isinstance(r, QuantizerRole) + assert ( + r.module_type == "grouped_linear" + ), f"Expected module_type='grouped_linear', got '{r.module_type}'" + assert r.name == "experts", f"Expected name='experts', got '{r.name}'" + + fwd_types = {r.tensor_type for r in non_none if r.tensor_type in ("input", "weight")} + bwd_types = {r.tensor_type for r in non_none if r.tensor_type == "grad_output"} + assert "input" in fwd_types, "Missing 'input' tensor_type in forward roles" + assert "weight" in fwd_types, "Missing 'weight' tensor_type in forward roles" + assert "grad_output" in bwd_types, "Missing 'grad_output' tensor_type in backward roles" + + +def test_delayed_scaling_request_wiring(): + """Shared buffers, correct views, Float8Quantizer instances.""" + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported: {reason}") + + from transformer_engine.pytorch.quantization import ( + DelayedScalingRequest, + CustomRecipeState, + ) + from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer + from transformer_engine.common.recipe import Format + + def ds_factory(role): + return DelayedScalingRequest(fp8_format=Format.HYBRID, amax_history_len=16) + + custom_recipe = recipe.CustomRecipe(qfactory=ds_factory) + + # 3 quantizers (input, weight, output) like a Linear fwd + state = CustomRecipeState( + custom_recipe, + mode="forward", + num_quantizers=3, + roles=[ + QuantizerRole(module_type="linear", tensor_type="input"), + QuantizerRole(module_type="linear", tensor_type="weight"), + QuantizerRole(module_type="linear", tensor_type="output"), + ], + ) + quantizers = state.make_quantizers() + + # All quantizers should be Float8Quantizer + assert len(quantizers) == 3 + for q in quantizers: + assert isinstance(q, Float8Quantizer), f"Expected Float8Quantizer, got {type(q).__name__}" + + # Managed state should exist + assert state._has_delayed_scaling + assert state.scale is not None + assert state.amax_history is not None + + # Shared buffers: scale shape = (3,), amax_history shape = (16, 3) + assert state.scale.shape == (3,) + assert state.amax_history.shape == (16, 3) + + # Each quantizer's scale should be a view into the shared buffer + for i, q in enumerate(quantizers): + assert q.scale.data_ptr() == state.scale[i].data_ptr() + + # Each quantizer's amax should be a view into amax_history[0] + for i, q in enumerate(quantizers): + assert q.amax.data_ptr() == state.amax_history[0][i].reshape((1,)).data_ptr() + + # Inner recipe should be a DelayedScaling + inner = state._inner_delayed_scaling_recipe + assert isinstance(inner, recipe.DelayedScaling) + assert inner.amax_history_len == 16 + assert inner.fp8_format == Format.HYBRID + + +def test_custom_recipe_mixed_ds_and_stateless(): + """Mix DelayedScalingRequest + stateless quantizers in same CustomRecipeState.""" + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported: {reason}") + + from transformer_engine.pytorch.quantization import ( + DelayedScalingRequest, + CustomRecipeState, + ) + from transformer_engine.pytorch.tensor.float8_tensor import Float8Quantizer + from transformer_engine.common.recipe import Format + + def mixed_factory(role): + # Only weight gets delayed scaling, rest get current scaling + if role is not None and role.tensor_type == "weight": + return DelayedScalingRequest(fp8_format=Format.HYBRID) + return Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device="cuda") + + custom_recipe = recipe.CustomRecipe(qfactory=mixed_factory) + + # 3 quantizers: input(current), weight(DS), output(current) + state = CustomRecipeState( + custom_recipe, + mode="forward", + num_quantizers=3, + roles=[ + QuantizerRole(module_type="linear", tensor_type="input"), + QuantizerRole(module_type="linear", tensor_type="weight"), + QuantizerRole(module_type="linear", tensor_type="output"), + ], + ) + quantizers = state.make_quantizers() + assert len(quantizers) == 3 + + # Slot 0 (input): current scaling + assert isinstance(quantizers[0], Float8CurrentScalingQuantizer) + # Slot 1 (weight): delayed scaling + assert isinstance(quantizers[1], Float8Quantizer) + # Slot 2 (output): current scaling + assert isinstance(quantizers[2], Float8CurrentScalingQuantizer) + + # Only 1 DS request => shared buffers have size 1 + assert state._has_delayed_scaling + assert state.scale.shape == (1,) + assert state.amax_history.shape == (1024, 1) + + +def test_custom_recipe_ds_multi_step(): + """amax_history updates across multiple forward steps.""" + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported: {reason}") + + from transformer_engine.pytorch.quantization import DelayedScalingRequest + from transformer_engine.common.recipe import Format + + def ds_factory(role): + return DelayedScalingRequest(fp8_format=Format.HYBRID) + + in_features = 128 + out_features = 128 + batch = 32 + num_steps = 3 + + torch.manual_seed(99) + model = Linear(in_features, out_features, params_dtype=torch.bfloat16, bias=False).cuda() + custom = recipe.CustomRecipe(qfactory=ds_factory) + + amax_snapshots = [] + for step in range(num_steps): + inp = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + with autocast(enabled=True, recipe=custom): + out = model(inp) + loss = out.float().sum() + loss.backward() + + # Capture amax_history snapshot + fwd_state = model.fp8_meta["scaling_fwd"] + amax_snapshots.append(fwd_state.amax_history.clone()) + + # After 3 steps, amax_history should have been updated at least once + # The first row (amax_history[0]) should differ from the initial zeros + # after the first step + assert not torch.all(amax_snapshots[0] == 0), "amax_history should be updated after first step" + + +# ---------------------------------------------------------------------- +# State preservation across role-driven rebuilds +# ---------------------------------------------------------------------- +# +# Setting ``output_quantizer_role`` / ``grad_input_quantizer_role`` to a +# different value flips ``fp8_meta_tensors_initialized = False`` so the +# next ``set_meta_tensor`` call rebuilds the recipe state and quantizers +# with up-to-date roles. That rebuild MUST preserve persistent training +# buffers (delayed scaling's ``scale`` / ``amax_history``); otherwise +# checkpointed amax history is silently destroyed on the first forward +# pass after ``load_state_dict`` (when MHA wires boundary roles for the +# first time on the freshly-loaded module). The buffers must also be +# preserved by tensor-object identity, not just by value: the +# ``FP8GlobalStateManager`` reduction buffer holds a direct reference to +# the tensor created at first init, so any rebuild that allocates fresh +# tensors would break amax all-reduce. + + +def test_role_change_preserves_delayed_scaling_state(): + """Built-in DelayedScaling: role-driven rebuild preserves scale / amax_history. + + Stashes sentinel values into the buffers, forces a rebuild via the role + setter, and verifies values + tensor-object identity survive. + """ + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported: {reason}") + + torch.manual_seed(0) + model = Linear(64, 64, params_dtype=torch.bfloat16, bias=False).cuda() + inp = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) + fp8_recipe = recipe.DelayedScaling(amax_history_len=8) + + # Initialize state via a forward pass. + with autocast(enabled=True, recipe=fp8_recipe): + model(inp).float().sum().backward() + assert model.fp8_meta_tensors_initialized + + state_before = model.fp8_meta["scaling_fwd"] + state_before.scale.fill_(3.14) + state_before.amax_history.fill_(2.71) + scale_obj_id = id(state_before.scale) + amax_obj_id = id(state_before.amax_history) + scale_data_ptr = state_before.scale.data_ptr() + amax_data_ptr = state_before.amax_history.data_ptr() + + # Trigger role-driven invalidation. Setting a non-None role flips + # ``fp8_meta_tensors_initialized = False`` so the next ``set_meta_tensor`` + # falls through and creates a fresh ``RecipeState``. + model.output_quantizer_role = QuantizerRole( + module_type="dpa", tensor_type="qkv", name="downstream" + ) + assert not model.fp8_meta_tensors_initialized + + # Trigger the rebuild directly (no forward, so we can compare buffers exactly). + model.init_fp8_meta_tensors(fp8_recipe) + assert model.fp8_meta_tensors_initialized + + state_after = model.fp8_meta["scaling_fwd"] + assert state_after is not state_before, "state should have been rebuilt" + # Tensor objects must be inherited (not freshly allocated) so the + # FP8GlobalStateManager reduction buffer's reference stays valid. + assert ( + id(state_after.scale) == scale_obj_id + ), "scale tensor object replaced by rebuild; global reduction buffer would dangle" + assert id(state_after.amax_history) == amax_obj_id + assert state_after.scale.data_ptr() == scale_data_ptr + assert state_after.amax_history.data_ptr() == amax_data_ptr + # Sentinel values must be preserved. + assert state_after.scale.eq(3.14).all(), "scale was wiped by role-driven rebuild" + assert state_after.amax_history.eq(2.71).all(), "amax_history was wiped" + + +def test_role_change_preserves_custom_delayed_scaling_state(): + """CustomRecipe + DelayedScalingRequest: role-driven rebuild preserves inner DSRS. + + Same property as the built-in case, but for the + ``CustomRecipeState`` -> composed ``DelayedScalingRecipeState`` path. + The inner DS state must be re-used across the rebuild so its + accumulated buffers (and any external references to them) survive. + """ + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported: {reason}") + + from transformer_engine.pytorch.quantization import ( + CustomRecipeState, + DelayedScalingRequest, + ) + from transformer_engine.common.recipe import Format + + def ds_factory(role): + return DelayedScalingRequest(fp8_format=Format.HYBRID, amax_history_len=8) + + torch.manual_seed(0) + model = Linear(64, 64, params_dtype=torch.bfloat16, bias=False).cuda() + inp = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) + custom_recipe = recipe.CustomRecipe(qfactory=ds_factory) + + # Initialize state via a forward pass. + with autocast(enabled=True, recipe=custom_recipe): + model(inp).float().sum().backward() + assert model.fp8_meta_tensors_initialized + + state_before = model.fp8_meta["scaling_fwd"] + assert isinstance(state_before, CustomRecipeState) + assert state_before._has_delayed_scaling + inner_before = state_before._ds_state + inner_before.scale.fill_(3.14) + inner_before.amax_history.fill_(2.71) + scale_obj_id = id(inner_before.scale) + amax_obj_id = id(inner_before.amax_history) + + # Trigger role-driven invalidation. + model.output_quantizer_role = QuantizerRole( + module_type="dpa", tensor_type="qkv", name="downstream" + ) + assert not model.fp8_meta_tensors_initialized + + # Rebuild. + model.init_fp8_meta_tensors(custom_recipe) + assert model.fp8_meta_tensors_initialized + + state_after = model.fp8_meta["scaling_fwd"] + assert isinstance(state_after, CustomRecipeState) + assert state_after is not state_before, "outer CustomRecipeState should have been rebuilt" + assert state_after._has_delayed_scaling, "rebuild lost the inner DS state" + inner_after = state_after._ds_state + # Inner DSRS object identity is preserved (we reuse the existing inner state), + # which means its buffers' tensor objects are also preserved. + assert ( + inner_after is inner_before + ), "inner DSRS replaced; FP8GlobalStateManager reduction buffer would dangle" + assert id(inner_after.scale) == scale_obj_id + assert id(inner_after.amax_history) == amax_obj_id + # Sentinel values preserved. + assert inner_after.scale.eq(3.14).all() + assert inner_after.amax_history.eq(2.71).all() + + +def test_role_change_does_not_invalidate_when_role_unchanged(): + """Setting the role to its current value is a no-op (no rebuild).""" + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported: {reason}") + + torch.manual_seed(0) + model = Linear(64, 64, params_dtype=torch.bfloat16, bias=False).cuda() + inp = torch.randn(32, 64, device="cuda", dtype=torch.bfloat16, requires_grad=True) + fp8_recipe = recipe.DelayedScaling(amax_history_len=8) + + role = QuantizerRole(module_type="dpa", tensor_type="qkv", name="x") + model.output_quantizer_role = role # initial set: state not yet built, no-op + + with autocast(enabled=True, recipe=fp8_recipe): + model(inp).float().sum().backward() + assert model.fp8_meta_tensors_initialized + + # Re-setting the same role value must not invalidate. + model.output_quantizer_role = QuantizerRole(module_type="dpa", tensor_type="qkv", name="x") + assert ( + model.fp8_meta_tensors_initialized + ), "Setting role to an equal value should be a no-op (frozen-dataclass __eq__)" + + +def test_custom_recipe_dpa_fp8(): + """DotProductAttention forward+backward with CustomRecipe and role-based mixed quantizers. + + Uses the nvfp4_linear_fp8_dpa_factory which dispatches: + * DPA S/dP slots -> DelayedScalingRequest (stateful) + * DPA QKV/O/dO/dQKV slots -> Float8CurrentScalingQuantizer + * Linear slots -> NVFP4Quantizer + """ + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported on this device: {reason}") + if not te.is_nvfp4_available(): + pytest.skip("NVFP4 unsupported on this device") + + from transformer_engine.pytorch.utils import get_device_compute_capability + + cc = get_device_compute_capability() + if cc < (9, 0) or cc >= (12, 0): + pytest.skip(f"FP8 attention not supported on sm{cc[0]*10+cc[1]}") + + from transformer_engine.pytorch.quantization import ( + DelayedScalingRequest, + CustomRecipeState, + ) + from transformer_engine.pytorch.tensor.float8_tensor import ( + Float8Quantizer, + Float8CurrentScalingQuantizer, + ) + from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( + nvfp4_linear_fp8_dpa_factory, + ) + + torch.manual_seed(42) + + H = 64 + NH = 4 + KV = H // NH + B = 2 + S = 32 + + # Build a small model: Linear -> DPA -> Linear + qkv_proj = Linear(H, 3 * H, params_dtype=torch.bfloat16, bias=False, name="qkv").cuda() + dpa = te.DotProductAttention( + NH, KV, attention_dropout=0.0, qkv_format="bshd", name="core_attention" + ) + out_proj = Linear(H, H, params_dtype=torch.bfloat16, bias=False, name="proj").cuda() + + custom_recipe = recipe.CustomRecipe( + qfactory=nvfp4_linear_fp8_dpa_factory, + fp8_dpa=True, + ) + + inp = torch.randn(B, S, H, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with autocast(enabled=True, recipe=custom_recipe): + qkv = qkv_proj(inp).view(B, S, 3, NH, KV) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + attn_out = dpa(q, k, v, qkv_format="bshd").reshape(B, S, H) + out = out_proj(attn_out) + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None, "Input gradient should exist" + + # Verify DPA recipe state is CustomRecipeState + fwd_state = dpa.fp8_meta["scaling_fwd"] + assert isinstance( + fwd_state, CustomRecipeState + ), f"Expected CustomRecipeState for DPA fwd, got {type(fwd_state).__name__}" + + # Verify DPA quantizers: 9 forward slots (3 GEMMs x 3) + fwd_quantizers = dpa.quantizers["scaling_fwd"] + assert len(fwd_quantizers) == 9, f"Expected 9 fwd quantizers, got {len(fwd_quantizers)}" + + # Slots 0-2: QKV (GEMM1) -> current scaling (role: module_type="dpa") + # Slots 3-5: O (GEMM2) -> current scaling (role: name hint "dpa_output") + # Slots 6-8: S (GEMM3) -> delayed scaling (Float8Quantizer from DelayedScalingRequest) + for i in range(6): + assert isinstance(fwd_quantizers[i], Float8CurrentScalingQuantizer), ( + f"Slot {i} (QKV/O): expected Float8CurrentScalingQuantizer, " + f"got {type(fwd_quantizers[i]).__name__}" + ) + for i in range(6, 9): + assert isinstance(fwd_quantizers[i], Float8Quantizer), ( + f"Slot {i} (S): expected Float8Quantizer (delayed scaling), " + f"got {type(fwd_quantizers[i]).__name__}" + ) + + # Verify DS state exists for the S/dP delayed scaling requests + assert fwd_state._has_delayed_scaling, "DPA fwd state should have delayed scaling for S slots" + + # Verify backward quantizers exist too + bwd_quantizers = dpa.quantizers["scaling_bwd"] + assert len(bwd_quantizers) == 6, f"Expected 6 bwd quantizers, got {len(bwd_quantizers)}" + + # Slots 0-1: dQKV (GEMM1) -> current scaling (role: name hint "dpa_grad_input") + # Slots 2-3: dO (GEMM2) -> current scaling (role: module_type="dpa") + # Slots 4-5: dP (GEMM3) -> delayed scaling + for i in range(4): + assert isinstance(bwd_quantizers[i], Float8CurrentScalingQuantizer), ( + f"Bwd slot {i} (dQKV/dO): expected Float8CurrentScalingQuantizer, " + f"got {type(bwd_quantizers[i]).__name__}" + ) + for i in range(4, 6): + assert isinstance(bwd_quantizers[i], Float8Quantizer), ( + f"Bwd slot {i} (dP): expected Float8Quantizer (delayed scaling), " + f"got {type(bwd_quantizers[i]).__name__}" + ) + + # Linear modules should have CustomRecipeState with NVFP4 quantizers + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + + qkv_fwd = qkv_proj.fp8_meta["scaling_fwd"] + assert isinstance( + qkv_fwd, CustomRecipeState + ), f"Expected CustomRecipeState for qkv_proj, got {type(qkv_fwd).__name__}" + qkv_fwd_quantizers = qkv_proj.quantizers["scaling_fwd"] + for i, q in enumerate(qkv_fwd_quantizers): + if q is not None: + assert isinstance( + q, NVFP4Quantizer + ), f"qkv_proj fwd slot {i}: expected NVFP4Quantizer, got {type(q).__name__}" + + +def test_custom_recipe_dpa_mxfp8(): + """DotProductAttention forward+backward with CustomRecipe and MXFP8 attention. + + Uses the nvfp4_linear_mxfp8_dpa_factory which dispatches: + * DPA roles (QKV/O/S/dO/dP/dQKV) -> MXFP8Quantizer (S/dP later nulled + out by ``get_attention_quantizers`` since the MXFP8 fused-attention + kernel handles those slots internally) + * DPA boundary hints -> MXFP8Quantizer + * Linear slots -> NVFP4Quantizer + + Mirrors the documented "NVFP4 linear + MXFP8 attention" combo from + ``dot_product_attention.py``'s recipe-combination table. + """ + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported on this device: {reason}") + if not te.is_mxfp8_available(): + pytest.skip("MXFP8 unsupported on this device") + if not te.is_nvfp4_available(): + pytest.skip("NVFP4 unsupported on this device") + + from transformer_engine.pytorch.utils import get_device_compute_capability + + cc = get_device_compute_capability() + if cc < (9, 0) or cc >= (12, 0): + pytest.skip(f"FP8 attention not supported on sm{cc[0]*10+cc[1]}") + + from transformer_engine.pytorch.quantization import CustomRecipeState + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( + nvfp4_linear_mxfp8_dpa_factory, + ) + + torch.manual_seed(42) + + # MXFP8 fused attention requires s_q % 128 == 0, s_kv % 128 == 0, + # d_qk % 32 == 0, d_v % 32 == 0. + H = 128 + NH = 4 + KV = H // NH # 32 + B = 2 + S = 128 + + # Build a small model: Linear -> DPA -> Linear + qkv_proj = Linear(H, 3 * H, params_dtype=torch.bfloat16, bias=False, name="qkv").cuda() + dpa = te.DotProductAttention( + NH, KV, attention_dropout=0.0, qkv_format="bshd", name="core_attention" + ) + out_proj = Linear(H, H, params_dtype=torch.bfloat16, bias=False, name="proj").cuda() + + custom_recipe = recipe.CustomRecipe( + qfactory=nvfp4_linear_mxfp8_dpa_factory, + fp8_dpa=True, + ) + + inp = torch.randn(B, S, H, device="cuda", dtype=torch.bfloat16, requires_grad=True) + + with autocast(enabled=True, recipe=custom_recipe): + qkv = qkv_proj(inp).view(B, S, 3, NH, KV) + q, k, v = qkv[:, :, 0], qkv[:, :, 1], qkv[:, :, 2] + # MXFP8 fused attention requires s_q % 128 == 0, s_kv % 128 == 0, + # d_qk % 32 == 0, d_v % 32 == 0. The B/S/H values above are picked + # to satisfy all four constraints (S=128, KV=32). + attn_out = dpa(q, k, v, qkv_format="bshd").reshape(B, S, H) + out = out_proj(attn_out) + + loss = out.float().sum() + loss.backward() + + assert inp.grad is not None, "Input gradient should exist" + + # DPA recipe state should be CustomRecipeState + fwd_state = dpa.fp8_meta["scaling_fwd"] + assert isinstance( + fwd_state, CustomRecipeState + ), f"Expected CustomRecipeState for DPA fwd, got {type(fwd_state).__name__}" + + # All DPA slots should resolve to MXFP8Quantizer (the factory returns MXFP8 + # uniformly for DPA roles; S/dP nulling happens inside get_attention_quantizers + # at fused-attn dispatch time, not here). + fwd_quantizers = dpa.quantizers["scaling_fwd"] + assert len(fwd_quantizers) == 9, f"Expected 9 fwd quantizers, got {len(fwd_quantizers)}" + for i, q in enumerate(fwd_quantizers): + assert isinstance( + q, MXFP8Quantizer + ), f"DPA fwd slot {i}: expected MXFP8Quantizer, got {type(q).__name__}" + + bwd_quantizers = dpa.quantizers["scaling_bwd"] + assert len(bwd_quantizers) == 6, f"Expected 6 bwd quantizers, got {len(bwd_quantizers)}" + for i, q in enumerate(bwd_quantizers): + assert isinstance( + q, MXFP8Quantizer + ), f"DPA bwd slot {i}: expected MXFP8Quantizer, got {type(q).__name__}" + + # MXFP8 attention has no delayed-scaling state (no S/dP DS-request slots). + assert ( + not fwd_state._has_delayed_scaling + ), "DPA fwd state should NOT have delayed scaling for the all-MXFP8 factory" + + # Linear modules should still be NVFP4 + qkv_fwd = qkv_proj.fp8_meta["scaling_fwd"] + assert isinstance( + qkv_fwd, CustomRecipeState + ), f"Expected CustomRecipeState for qkv_proj, got {type(qkv_fwd).__name__}" + qkv_fwd_quantizers = qkv_proj.quantizers["scaling_fwd"] + for i, q in enumerate(qkv_fwd_quantizers): + if q is not None: + assert isinstance( + q, NVFP4Quantizer + ), f"qkv_proj fwd slot {i}: expected NVFP4Quantizer, got {type(q).__name__}" + + +def test_custom_recipe_debug_tool_compat(): + """Custom recipe quantizers should work when wrapped by DebugQuantizer. + + Verifies that the debug tool (nvdlfw_inspect) can wrap custom-recipe + quantizers produced via QuantizerRole dispatch without errors. + """ + try: + import nvdlfw_inspect.api as debug_api + except ImportError: + pytest.skip("nvdlfw_inspect not installed") + + available, reason = te.is_fp8_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 unsupported: {reason}") + + import pathlib + import tempfile + + from transformer_engine.debug.pytorch.debug_state import TEDebugState + + te_debug_features = str( + pathlib.Path(__file__).resolve().parent.parent.parent + / "transformer_engine" + / "debug" + / "features" + ) + + # Log config that keeps DebugQuantizer active (not bypassed by no_debug_features_active) + log_config = """log: + layers: + layer_types: [linear] + enabled: True + transformer_engine: + LogTensorStats: + enabled: True + tensors: [activation, weight] + stats: [max] + start_step: 0 + end_step: 3 +""" + + torch.manual_seed(0) + + in_features = 64 + out_features = 64 + batch = 16 + + with tempfile.NamedTemporaryFile(mode="w+", suffix=".yaml", delete=False) as cfg: + cfg.write(log_config) + cfg.flush() + config_path = cfg.name + + try: + with tempfile.TemporaryDirectory() as log_dir: + debug_api.initialize( + config_file=config_path, + feature_dirs=te_debug_features, + log_dir=log_dir, + ) + + model = Linear( + in_features, out_features, params_dtype=torch.bfloat16, name="layer" + ).cuda() + + custom_recipe = recipe.CustomRecipe(qfactory=current_scaling_quantizer_factory) + + assert TEDebugState.debug_enabled, "Debug mode should be active" + + for _ in range(3): + inp_step = torch.randn( + batch, in_features, device="cuda", dtype=torch.bfloat16, requires_grad=True + ) + with autocast(enabled=True, recipe=custom_recipe): + out = model(inp_step) + out.float().sum().backward() + debug_api.step() + + assert inp_step.grad is not None, "Input gradient should exist" + + log_files = list(pathlib.Path(log_dir).rglob("*.log")) + assert ( + len(log_files) > 0 + ), f"Debug log output expected in {log_dir} but no .log files found" + finally: + debug_api.end_debug() + TEDebugState._reset() + import os + + os.unlink(config_path) + + +# ---------------------------------------------------------------------- +# Role-aware dispatch in built-in block-scaling recipe states +# ---------------------------------------------------------------------- +# +# These tests exercise ``Float8BlockScalingRecipeState.make_quantizers`` and +# ``NVFP4BlockScalingRecipeState.make_quantizers`` directly to verify that +# per-slot dispatch is driven by ``QuantizerRole.tensor_type`` with a +# positional fallback that matches the legacy behavior. They construct the +# recipe state objects directly (no autocast / no fwd pass) so they don't +# depend on any module's ``get_quantizer_roles`` implementation. + + +def _fp8block_role(tensor_type): + """QuantizerRole helper for FP8-block tests.""" + return QuantizerRole(module_type="linear", tensor_type=tensor_type, name="t") + + +def test_fp8block_recipe_state_role_dispatch_forward(): + """Forward dispatch: input/output -> x cfg, weight -> w cfg.""" + available, reason = te.is_fp8_block_scaling_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 block scaling unsupported: {reason}") + + from transformer_engine.pytorch.quantization import Float8BlockScalingRecipeState + + fp8_recipe = recipe.Float8BlockScaling() + state = Float8BlockScalingRecipeState( + fp8_recipe, + mode="forward", + num_quantizers=3, + roles=[ + _fp8block_role("input"), + _fp8block_role("weight"), + _fp8block_role("output"), + ], + ) + quantizers = state.make_quantizers() + assert len(quantizers) == 3 + # input slot uses x cfg + assert quantizers[0].block_scaling_dim == fp8_recipe.x_block_scaling_dim + # weight slot uses w cfg + assert quantizers[1].block_scaling_dim == fp8_recipe.w_block_scaling_dim + # output slot mirrors input cfg (legacy behavior preserved) + assert quantizers[2].block_scaling_dim == fp8_recipe.x_block_scaling_dim + # Sanity: the recipe defaults distinguish x and w block scaling dims so + # the test would fail if dispatch were uniform. + assert fp8_recipe.x_block_scaling_dim != fp8_recipe.w_block_scaling_dim + + +def test_fp8block_recipe_state_role_dispatch_backward(): + """Backward dispatch: grad_output / grad_input both -> grad cfg.""" + available, reason = te.is_fp8_block_scaling_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 block scaling unsupported: {reason}") + + from transformer_engine.pytorch.quantization import Float8BlockScalingRecipeState + + fp8_recipe = recipe.Float8BlockScaling() + state = Float8BlockScalingRecipeState( + fp8_recipe, + mode="backward", + num_quantizers=2, + roles=[ + _fp8block_role("grad_output"), + _fp8block_role("grad_input"), + ], + ) + quantizers = state.make_quantizers() + assert len(quantizers) == 2 + for q in quantizers: + assert q.block_scaling_dim == fp8_recipe.grad_block_scaling_dim + + +def test_fp8block_recipe_state_positional_fallback_matches_explicit_roles(): + """``roles=None`` produces the same per-slot configs as explicit ``[input, weight, output]``.""" + available, reason = te.is_fp8_block_scaling_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 block scaling unsupported: {reason}") + + from transformer_engine.pytorch.quantization import Float8BlockScalingRecipeState + + fp8_recipe = recipe.Float8BlockScaling() + + explicit = Float8BlockScalingRecipeState( + fp8_recipe, + mode="forward", + num_quantizers=3, + roles=[ + _fp8block_role("input"), + _fp8block_role("weight"), + _fp8block_role("output"), + ], + ).make_quantizers() + + fallback = Float8BlockScalingRecipeState( + fp8_recipe, + mode="forward", + num_quantizers=3, + roles=None, + ).make_quantizers() + + assert len(explicit) == len(fallback) == 3 + for a, b in zip(explicit, fallback): + assert a.block_scaling_dim == b.block_scaling_dim + assert a.dtype == b.dtype + assert a.amax_epsilon == b.amax_epsilon + assert a.force_pow_2_scales == b.force_pow_2_scales + + +def test_fp8block_recipe_state_supports_non_multiple_of_three(): + """Two-slot forward (fusible-Linear shape) used to fail ``% 3 == 0`` assert.""" + available, reason = te.is_fp8_block_scaling_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 block scaling unsupported: {reason}") + + from transformer_engine.pytorch.quantization import Float8BlockScalingRecipeState + + fp8_recipe = recipe.Float8BlockScaling() + state = Float8BlockScalingRecipeState( + fp8_recipe, + mode="forward", + num_quantizers=2, + roles=[ + _fp8block_role("input"), + _fp8block_role("weight"), + ], + ) + quantizers = state.make_quantizers() + assert len(quantizers) == 2 + assert quantizers[0].block_scaling_dim == fp8_recipe.x_block_scaling_dim + assert quantizers[1].block_scaling_dim == fp8_recipe.w_block_scaling_dim + + +def test_fp8block_recipe_state_unknown_or_none_role_falls_back_positionally(): + """Per-slot ``None`` and unknown ``tensor_type`` use the positional pattern.""" + available, reason = te.is_fp8_block_scaling_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 block scaling unsupported: {reason}") + + from transformer_engine.pytorch.quantization import Float8BlockScalingRecipeState + + fp8_recipe = recipe.Float8BlockScaling() + # Slot 0: bare role (empty tensor_type) -> positional "input" -> x cfg + # Slot 1: unknown tensor_type "qkv" (DPA-style) -> positional "weight" -> w cfg + # Slot 2: None role -> positional "output" -> x cfg + state = Float8BlockScalingRecipeState( + fp8_recipe, + mode="forward", + num_quantizers=3, + roles=[ + QuantizerRole(), + QuantizerRole(module_type="dpa", tensor_type="qkv"), + None, + ], + ) + quantizers = state.make_quantizers() + assert len(quantizers) == 3 + assert quantizers[0].block_scaling_dim == fp8_recipe.x_block_scaling_dim + assert quantizers[1].block_scaling_dim == fp8_recipe.w_block_scaling_dim + assert quantizers[2].block_scaling_dim == fp8_recipe.x_block_scaling_dim + + +def _nvfp4_role(tensor_type): + return QuantizerRole(module_type="linear", tensor_type=tensor_type, name="t") + + +def test_nvfp4_recipe_state_role_dispatch_forward(): + """Forward dispatch: input/output -> inp cfg (RHT, 1D), weight -> weight cfg (no RHT, 2D).""" + if not torch.cuda.is_available() or not te.is_nvfp4_available(): + pytest.skip("NVFP4 unsupported on this device") + + from transformer_engine.pytorch.quantization import NVFP4BlockScalingRecipeState + + nvfp4_recipe = recipe.NVFP4BlockScaling() + state = NVFP4BlockScalingRecipeState( + nvfp4_recipe, + mode="forward", + num_quantizers=3, + roles=[ + _nvfp4_role("input"), + _nvfp4_role("weight"), + _nvfp4_role("output"), + ], + ) + quantizers = state.make_quantizers() + assert len(quantizers) == 3 + # input slot + assert quantizers[0].with_rht == nvfp4_recipe.fp4_quant_fwd_inp.random_hadamard_transform + assert quantizers[0].with_2d_quantization == nvfp4_recipe.fp4_quant_fwd_inp.fp4_2d_quantization + # weight slot + assert quantizers[1].with_rht == nvfp4_recipe.fp4_quant_fwd_weight.random_hadamard_transform + assert ( + quantizers[1].with_2d_quantization == nvfp4_recipe.fp4_quant_fwd_weight.fp4_2d_quantization + ) + # output slot mirrors input cfg + assert quantizers[2].with_rht == nvfp4_recipe.fp4_quant_fwd_inp.random_hadamard_transform + assert quantizers[2].with_2d_quantization == nvfp4_recipe.fp4_quant_fwd_inp.fp4_2d_quantization + # Sanity: defaults distinguish input vs weight (RHT and 2D toggles differ). + assert ( + nvfp4_recipe.fp4_quant_fwd_inp.random_hadamard_transform + != nvfp4_recipe.fp4_quant_fwd_weight.random_hadamard_transform + ) or ( + nvfp4_recipe.fp4_quant_fwd_inp.fp4_2d_quantization + != nvfp4_recipe.fp4_quant_fwd_weight.fp4_2d_quantization + ) + + +def test_nvfp4_recipe_state_role_dispatch_backward(): + """Backward dispatch: any slot -> grad cfg (uniform).""" + if not torch.cuda.is_available() or not te.is_nvfp4_available(): + pytest.skip("NVFP4 unsupported on this device") + + from transformer_engine.pytorch.quantization import NVFP4BlockScalingRecipeState + + nvfp4_recipe = recipe.NVFP4BlockScaling() + state = NVFP4BlockScalingRecipeState( + nvfp4_recipe, + mode="backward", + num_quantizers=2, + roles=[ + _nvfp4_role("grad_output"), + _nvfp4_role("grad_input"), + ], + ) + quantizers = state.make_quantizers() + assert len(quantizers) == 2 + for q in quantizers: + assert q.with_rht == nvfp4_recipe.fp4_quant_bwd_grad.random_hadamard_transform + assert q.with_2d_quantization == nvfp4_recipe.fp4_quant_bwd_grad.fp4_2d_quantization + assert q.stochastic_rounding == nvfp4_recipe.fp4_quant_bwd_grad.stochastic_rounding + + +def test_nvfp4_recipe_state_positional_fallback_matches_explicit_roles(): + """``roles=None`` matches explicit ``[input, weight, output]`` slot-for-slot.""" + if not torch.cuda.is_available() or not te.is_nvfp4_available(): + pytest.skip("NVFP4 unsupported on this device") + + from transformer_engine.pytorch.quantization import NVFP4BlockScalingRecipeState + + nvfp4_recipe = recipe.NVFP4BlockScaling() + + explicit = NVFP4BlockScalingRecipeState( + nvfp4_recipe, + mode="forward", + num_quantizers=3, + roles=[ + _nvfp4_role("input"), + _nvfp4_role("weight"), + _nvfp4_role("output"), + ], + ).make_quantizers() + + fallback = NVFP4BlockScalingRecipeState( + nvfp4_recipe, + mode="forward", + num_quantizers=3, + roles=None, + ).make_quantizers() + + assert len(explicit) == len(fallback) == 3 + for a, b in zip(explicit, fallback): + assert a.with_rht == b.with_rht + assert a.with_post_rht_amax == b.with_post_rht_amax + assert a.with_2d_quantization == b.with_2d_quantization + assert a.stochastic_rounding == b.stochastic_rounding + assert a.dtype == b.dtype + + +def test_nvfp4_recipe_state_supports_non_multiple_of_three(): + """Two-slot forward (fusible-Linear shape) succeeds with role-driven dispatch.""" + if not torch.cuda.is_available() or not te.is_nvfp4_available(): + pytest.skip("NVFP4 unsupported on this device") + + from transformer_engine.pytorch.quantization import NVFP4BlockScalingRecipeState + + nvfp4_recipe = recipe.NVFP4BlockScaling() + state = NVFP4BlockScalingRecipeState( + nvfp4_recipe, + mode="forward", + num_quantizers=2, + roles=[ + _nvfp4_role("input"), + _nvfp4_role("weight"), + ], + ) + quantizers = state.make_quantizers() + assert len(quantizers) == 2 + assert quantizers[0].with_rht == nvfp4_recipe.fp4_quant_fwd_inp.random_hadamard_transform + assert quantizers[1].with_rht == nvfp4_recipe.fp4_quant_fwd_weight.random_hadamard_transform + + +def test_nvfp4_recipe_state_unknown_or_none_role_falls_back_positionally(): + """Per-slot ``None`` and unknown ``tensor_type`` use the positional pattern.""" + if not torch.cuda.is_available() or not te.is_nvfp4_available(): + pytest.skip("NVFP4 unsupported on this device") + + from transformer_engine.pytorch.quantization import NVFP4BlockScalingRecipeState + + nvfp4_recipe = recipe.NVFP4BlockScaling() + # Slot 0: bare role (empty tensor_type) -> positional "input" -> inp cfg + # Slot 1: DPA-style unknown tensor_type "qkv" -> positional "weight" -> weight cfg + # Slot 2: None role -> positional "output" -> inp cfg + state = NVFP4BlockScalingRecipeState( + nvfp4_recipe, + mode="forward", + num_quantizers=3, + roles=[ + QuantizerRole(), + QuantizerRole(module_type="dpa", tensor_type="qkv"), + None, + ], + ) + quantizers = state.make_quantizers() + assert len(quantizers) == 3 + assert quantizers[0].with_rht == nvfp4_recipe.fp4_quant_fwd_inp.random_hadamard_transform + assert quantizers[1].with_rht == nvfp4_recipe.fp4_quant_fwd_weight.random_hadamard_transform + assert quantizers[2].with_rht == nvfp4_recipe.fp4_quant_fwd_inp.random_hadamard_transform + + +# ---------------------------------------------------------------------- +# RecipeState._slot_role primitive +# ---------------------------------------------------------------------- +# +# `_slot_role` is the primitive that role-driven recipe states use to +# resolve per-slot dispatch info. It returns the real role when one was +# provided and synthesizes one with the positional ``tensor_type`` fallback +# (and empty ``module_type``/``name``) otherwise. Future recipes that +# dispatch on ``module_type`` / ``name`` rely on this contract. +# +# We exercise these via a concrete ``Float8BlockScalingRecipeState`` since +# ``RecipeState`` is abstract; the helper itself is mode-aware but +# recipe-agnostic. + + +def _make_fp8block_state(*, mode, num_quantizers, roles): + from transformer_engine.pytorch.quantization import Float8BlockScalingRecipeState + + return Float8BlockScalingRecipeState( + recipe.Float8BlockScaling(), + mode=mode, + num_quantizers=num_quantizers, + roles=roles, + ) + + +def test_slot_role_passes_real_role_through_unchanged(): + """A real ``QuantizerRole`` from the producer is returned as-is.""" + available, reason = te.is_fp8_block_scaling_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 block scaling unsupported: {reason}") + + real = QuantizerRole(module_type="linear", tensor_type="weight", name="layer37.fc1") + state = _make_fp8block_state(mode="forward", num_quantizers=1, roles=[real]) + resolved = state._slot_role(0) + # Identity: no copying, the real instance is returned. + assert resolved is real + assert resolved.module_type == "linear" + assert resolved.tensor_type == "weight" + assert resolved.name == "layer37.fc1" + + +def test_slot_role_passes_unknown_tensor_type_through_unchanged(): + """A real role with non-canonical ``tensor_type`` is NOT remapped by ``_slot_role``. + + ``_slot_tensor_type`` would fall back to positional, but ``_slot_role`` + must preserve the original so module-type / name dispatch still works. + """ + available, reason = te.is_fp8_block_scaling_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 block scaling unsupported: {reason}") + + dpa_role = QuantizerRole(module_type="dpa", tensor_type="qkv", name="self_attention.dpa") + state = _make_fp8block_state(mode="forward", num_quantizers=1, roles=[dpa_role]) + resolved = state._slot_role(0) + assert resolved is dpa_role + assert resolved.tensor_type == "qkv" # unchanged, NOT folded into known set + # ``_slot_tensor_type`` still falls back to positional pattern[0] = "input". + assert state._slot_tensor_type(0) == "input" + + +def test_slot_role_returns_bare_role_when_per_slot_role_is_none(): + """Boundary slot (``roles[i] is None``) returns a bare ``QuantizerRole()``. + + The primitive does NOT synthesize a positional ``tensor_type`` — that's + a tensor-type-dispatch policy owned by ``_slot_tensor_type``. + """ + available, reason = te.is_fp8_block_scaling_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 block scaling unsupported: {reason}") + + real_input = QuantizerRole(module_type="linear", tensor_type="input", name="t") + real_weight = QuantizerRole(module_type="linear", tensor_type="weight", name="t") + # Slot 2 (output) is None: typical for Linear without parent setting + # ``_output_quantizer_role``. + state = _make_fp8block_state( + mode="forward", num_quantizers=3, roles=[real_input, real_weight, None] + ) + # Real slots pass through. + assert state._slot_role(0) is real_input + assert state._slot_role(1) is real_weight + # None slot returns a bare QuantizerRole(): all fields empty, no + # tensor-type-specific synthesis. + bare = state._slot_role(2) + assert bare.tensor_type == "" + assert bare.module_type == "" + assert bare.name == "" + # Consumers get positional fallback through _slot_tensor_type, not _slot_role. + assert state._slot_tensor_type(2) == "output" + + +def test_slot_role_returns_bare_role_when_roles_list_is_none(): + """``roles=None`` yields bare ``QuantizerRole()`` for every slot, fwd and bwd. + + Positional fallback for tensor types lives in ``_slot_tensor_type``, not here. + """ + available, reason = te.is_fp8_block_scaling_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 block scaling unsupported: {reason}") + + fwd = _make_fp8block_state(mode="forward", num_quantizers=4, roles=None) + # _slot_role is field-agnostic: every slot is a bare QuantizerRole(). + for i in range(4): + role = fwd._slot_role(i) + assert role.tensor_type == "" + assert role.module_type == "" + assert role.name == "" + # _slot_tensor_type applies the positional fallback (with wrap). + fwd_types = [fwd._slot_tensor_type(i) for i in range(4)] + assert fwd_types == ["input", "weight", "output", "input"] + + bwd = _make_fp8block_state(mode="backward", num_quantizers=3, roles=None) + for i in range(3): + assert bwd._slot_role(i).tensor_type == "" + bwd_types = [bwd._slot_tensor_type(i) for i in range(3)] + assert bwd_types == ["grad_output", "grad_input", "grad_output"] + + +def test_slot_role_supports_module_type_only_role(): + """A role that fills ONLY ``module_type`` is preserved as-is. + + This is the producer convention for future module-type-driven recipes: + fill only the field(s) you have signal for. ``_slot_role`` must not + invent a ``tensor_type`` to mask the empty one (otherwise the module-type + branch in a mixed recipe would never see a clean signal). + """ + available, reason = te.is_fp8_block_scaling_available(return_reason=True) + if not torch.cuda.is_available() or not available: + pytest.skip(f"FP8 block scaling unsupported: {reason}") + + moe = QuantizerRole(module_type="moe_expert") + state = _make_fp8block_state(mode="forward", num_quantizers=1, roles=[moe]) + resolved = state._slot_role(0) + assert resolved is moe + assert resolved.module_type == "moe_expert" + assert resolved.tensor_type == "" # NOT auto-filled + assert resolved.name == "" + # Tensor-type-only recipes fall back to positional for this slot. + assert state._slot_tensor_type(0) == "input" diff --git a/tests/pytorch/test_float8_current_scaling_exact.py b/tests/pytorch/test_float8_current_scaling_exact.py index 99ab9c498..3b964a5af 100644 --- a/tests/pytorch/test_float8_current_scaling_exact.py +++ b/tests/pytorch/test_float8_current_scaling_exact.py @@ -14,7 +14,7 @@ from transformer_engine.pytorch.quantization import autocast, get_fp8_torch_dtype from transformer_engine.pytorch.constants import TE_DType from transformer_engine.pytorch.custom_recipes.quantization import MMParams -from transformer_engine.pytorch.custom_recipes.quantization_current_scaling import ( +from transformer_engine.pytorch.custom_recipes.quantization_ref_current_scaling import ( CurrentScalingQuantizerRef, ) diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 0d0b2fd37..959966369 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -558,19 +558,33 @@ class CustomRecipe(Recipe): Parameters ---------- qfactory : Callable - Factory callable that returns a quantizer instance for a - given semantic tensor role. - The callable is typically invoked as:: + Factory callable that returns a quantizer instance *or* a + ``QuantizerRequest`` subclass for a given ``QuantizerRole``. + The callable is invoked as:: qfactory( - role: str, - ) + role: QuantizerRole, + ) -> Union[Quantizer, QuantizerRequest] - Where `role` is one of the following strings for e.g. te.Linear - (stable public contract): + ``QuantizerRole`` is a frozen dataclass with the following fields: + + - ``module_type`` (str): module type (empty string when not set), e.g. + ``"linear"``, ``"grouped_linear"``, ``"dpa"``. + - ``tensor_type`` (str): what tensor is being quantized (empty + string when not set), e.g. ``"input"``, ``"weight"``, ``"grad_output"``. + - ``name`` (str): caller-provided module instance name (empty + string when not set), e.g. ``"qkv"``, ``"proj"``, ``"fc1"``, ``"fc2"``. + + For stateful quantizers (delayed scaling), return a + ``DelayedScalingRequest`` dataclass instead of a quantizer. + TE will allocate shared scale/amax_history buffers and create + ``Float8Quantizer`` instances integrated with the existing + delayed-scaling reduction infrastructure. + + See ``transformer_engine.pytorch.quantization.QuantizerRole`` + and ``transformer_engine.pytorch.quantization.DelayedScalingRequest`` + for full documentation. - - forward: "linear_input", "linear_weight", "linear_output" - - backward: "linear_grad_output", "linear_grad_input" backward_override : {None, 'high_precision', 'dequantized'}, default = None Backward precision mode. None does not modify backward behavior, `high_precision` keeps original high-precision operands for backward, @@ -580,6 +594,11 @@ class CustomRecipe(Recipe): qfactory: Callable[..., Any] + # fp8_format does not affect quantization (quantization factory controls that), + # but TE internals (e.g. get_fp8_te_dtype, backend selection) read it + # from the recipe. HYBRID (E4M3 fwd, E5M2 bwd) is a safe default. + fp8_format: Format = Format.HYBRID + fp8_dpa: bool = False fp8_mha: bool = False backward_override: Optional[str] = os.getenv("NVTE_BACKWARD_OVERRIDE", None) diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index d145cf0a2..3ff0d75ee 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -48,6 +48,9 @@ from transformer_engine.pytorch.quantization import is_fp8_block_scaling_available from transformer_engine.pytorch.quantization import is_nvfp4_available from transformer_engine.pytorch.quantization import get_default_recipe +from transformer_engine.pytorch.quantization import QuantizerRole +from transformer_engine.pytorch.quantization import QuantizerRequest +from transformer_engine.pytorch.quantization import DelayedScalingRequest from transformer_engine.pytorch.utils import get_cudnn_version from transformer_engine.pytorch.utils import get_device_compute_capability from transformer_engine.pytorch.utils import is_bf16_available diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 79ebbd4af..6e097265f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -175,6 +175,25 @@ _dpa_fp8_cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1" +def _qkv_quantizer_type(qkv_quantizer): + """Map a DPA QKV quantizer instance to its kernel-facing FP8 sub-recipe label. + + Returns one of ``'delayed'`` / ``'current'`` / ``'mxfp8'``. Used by FP8 + attention forward/backward to dispatch save-for-backward and + re-quantization decisions from the *quantizer instance* rather than + the top-level ``Recipe`` type, so that ``CustomRecipe`` + is handled correctly. Built-in recipes already + produce the matching quantizer instances, so behavior is preserved. + """ + if isinstance(qkv_quantizer, Float8Quantizer): + return "delayed" + if isinstance(qkv_quantizer, Float8CurrentScalingQuantizer): + return "current" + if isinstance(qkv_quantizer, MXFP8Quantizer): + return "mxfp8" + raise TypeError(f"Unsupported FP8 attention QKV quantizer: {type(qkv_quantizer).__name__}") + + class FP8EmulationFunc(torch.autograd.Function): """ Emulate the effects of FP8 quantization on tensors. Used in UnfusedDotProductAttention as follows: @@ -491,7 +510,7 @@ def forward( fp8_recipe = fp8_meta["local_recipes"][0] # get quantizers from DPA; all Nones if not fp8 QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer = ( - dpa_utils.get_attention_quantizers(fp8, fp8_recipe, quantizers) + dpa_utils.get_attention_quantizers(fp8, quantizers) ) # S/dP are forced to use DS quantizers in DPA.init_fp8_metadata; revert them here for true CS emulation if fp8_recipe.float8_current_scaling(): @@ -1318,9 +1337,15 @@ def forward( # get quantizers from DPA; all Nones if not fp8 QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer = ( - dpa_utils.get_attention_quantizers(fp8, fp8_recipe, quantizers) + dpa_utils.get_attention_quantizers(fp8, quantizers) ) + # Effective FP8 sub-recipe label inferred from the QKV quantizer + # instance. Drives save-for-backward and re-quantization dispatch + # below so that CustomRecipe (and built-in recipes alike) work + # without depending on `fp8_recipe.()`. + qkv_type = _qkv_quantizer_type(QKV_quantizer) if fp8 else None + # get nominal data type for out # FP16/BF16 attention: torch.float16 or torch.bfloat16 # FP8 attention: torch.float16 or torch.bfloat16 @@ -1402,19 +1427,13 @@ def forward( not is_bwd_fp8 or ( is_bwd_fp8 - and ( - (fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16) - or fp8_recipe.mxfp8() - ) + and ((qkv_type == "current" and _dpa_fp8_cs_o_in_f16) or qkv_type == "mxfp8") ) ) bwd_requires_o_fp8 = ( is_training and is_bwd_fp8 - and ( - fp8_recipe.delayed() - or (fp8_recipe.float8_current_scaling() and not _dpa_fp8_cs_o_in_f16) - ) + and (qkv_type == "delayed" or (qkv_type == "current" and not _dpa_fp8_cs_o_in_f16)) ) if isinstance(out_, QuantizedTensorStorage): if not is_output_fp8 or bwd_requires_o_f16: @@ -1442,14 +1461,10 @@ def forward( fp8_tensors = (None, None, None, None) f16_tensors = (None, None, None, None) if is_bwd_fp8: - if ( - fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16 - ) or fp8_recipe.mxfp8(): + if (qkv_type == "current" and _dpa_fp8_cs_o_in_f16) or qkv_type == "mxfp8": fp8_tensors = (q_fp8, k_fp8, v_fp8, None) f16_tensors = (None, None, None, out_f16) - elif fp8_recipe.delayed() or ( - fp8_recipe.float8_current_scaling() and not _dpa_fp8_cs_o_in_f16 - ): + elif qkv_type == "delayed" or (qkv_type == "current" and not _dpa_fp8_cs_o_in_f16): fp8_tensors = (q_fp8, k_fp8, v_fp8, out_fp8) else: if is_input_fp8: @@ -1536,6 +1551,7 @@ def forward( ctx.dO_quantizer = dO_quantizer ctx.dP_quantizer = dP_quantizer ctx.S_quantizer = S_quantizer + ctx.qkv_type = qkv_type if ctx.fp8 and isinstance(ctx.S_quantizer, Float8Quantizer): ctx.S_quantizer = S_quantizer.copy() ctx.S_quantizer.scale = S_quantizer.scale.clone() @@ -1706,9 +1722,9 @@ def backward(ctx, d_out, *_args): # MXFP8BlockScaling: # out_, dq_, dk_, dv_, d_out: torch.Tensor; dtype = torch.float16 or torch.bfloat16 out_ = out_fp8 - if ctx.fp8_recipe.float8_current_scaling() and _dpa_fp8_cs_o_in_f16: + if ctx.qkv_type == "current" and _dpa_fp8_cs_o_in_f16: out_ = out - if ctx.fp8_recipe.mxfp8(): + if ctx.qkv_type == "mxfp8": out_ = out aux_ctx_tensors.append(d_out) dq_, dk_, dv_, *rest = fused_attn_bwd( @@ -2059,7 +2075,7 @@ def forward( " with FP8!" ) if fp8_recipe.float8_current_scaling() and context_parallel: - all_quantizers = dpa_utils.get_attention_quantizers(fp8, fp8_recipe, quantizers) + all_quantizers = dpa_utils.get_attention_quantizers(fp8, quantizers) for q in all_quantizers: if isinstance(q, Float8CurrentScalingQuantizer): q.with_amax_reduction = True diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 32eb1b597..995ecf31b 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -58,6 +58,29 @@ _dpa_fp8_cs_o_in_f16 = os.getenv("NVTE_DPA_FP8CS_O_in_F16", "1") == "1" +def _reject_custom_recipe_under_cp(fp8, fp8_recipe): + """Fail fast when CustomRecipe meets context-parallel FP8 attention. + + Single-device FP8 attention dispatch was migrated to read quantizer + instance types (see ``backends._qkv_quantizer_type`` and + ``utils.get_attention_quantizers``), which makes CustomRecipe work end + to end. The CP code path in this module still dispatches on + ``fp8_recipe.()`` at ~90 sites; under CustomRecipe those + predicates all return False and the dispatch silently falls through to + incorrect tensor save / amax-reduction shapes. Until that migration + lands, surface the limitation here with a clear error rather than + failing later in C++ assertions or with silently-wrong gradients. + """ + if fp8 and fp8_recipe is not None and fp8_recipe.custom(): + raise NotImplementedError( + "CustomRecipe + Context Parallelism is not yet supported for FP8 " + "DotProductAttention. Either disable context parallelism, or use a " + "built-in FP8 recipe (DelayedScaling, Float8CurrentScaling, " + "MXFP8BlockScaling) for CP. The single-device CustomRecipe + DPA " + "path is supported." + ) + + def get_bsh_dims(tensor_format): """Get batch dimension and sequence dimension from tensor format""" if tensor_format in ["bshd", "sbhd", "bhsd"]: @@ -1453,6 +1476,7 @@ def forward( fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None: fp8_recipe = fp8_meta["local_recipes"][0] + _reject_custom_recipe_under_cp(fp8, fp8_recipe) ( QKV_quantizer, O_quantizer, @@ -1460,7 +1484,7 @@ def forward( dQKV_quantizer, dO_quantizer, dP_quantizer, - ) = dpa_utils.get_attention_quantizers(fp8, fp8_recipe, quantizers) + ) = dpa_utils.get_attention_quantizers(fp8, quantizers) # q, k, v a2a: gather s and split h # FP8DS/CS: Float8Tensor -> torch.uint8 -> Float8Tensor @@ -3043,6 +3067,7 @@ def forward( fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None: fp8_recipe = fp8_meta["local_recipes"][0] + _reject_custom_recipe_under_cp(fp8, fp8_recipe) ( QKV_quantizer, O_quantizer, @@ -3050,7 +3075,7 @@ def forward( dQKV_quantizer, dO_quantizer, dP_quantizer, - ) = dpa_utils.get_attention_quantizers(fp8, fp8_recipe, quantizers) + ) = dpa_utils.get_attention_quantizers(fp8, quantizers) fwd_nominal_dtype = q.dtype q_fp8, k_fp8, v_fp8 = (q, k, v) if is_input_fp8 else (None, None, None) q_f16, k_f16, v_f16 = (None, None, None) if is_input_fp8 else (q, k, v) @@ -3904,13 +3929,14 @@ def forward( fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_meta is not None and fp8_meta.get("local_recipes", None) is not None: fp8_recipe = fp8_meta["local_recipes"][0] + _reject_custom_recipe_under_cp(fp8, fp8_recipe) fwd_nominal_dtype = q.dtype fused_attn_backend = None max_logit = None QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer = ( - dpa_utils.get_attention_quantizers(fp8, fp8_recipe, quantizers) + dpa_utils.get_attention_quantizers(fp8, quantizers) ) q_fp8, k_fp8, v_fp8 = (None, None, None) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index 17e9a337a..b38b66c3e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -23,6 +23,7 @@ ) from transformer_engine.pytorch.utils import get_cudnn_version from transformer_engine.pytorch.quantization import ( + QuantizerRole, get_fp8_te_dtype, FP8GlobalStateManager, RecipeState, @@ -312,6 +313,8 @@ class DotProductAttention(TransformerEngineBaseModule): `_). :math:`\text{max_logit} = \max(S)`, where :math:`S = \text{mask}(Q \cdot K^T \cdot \text{softmax_scale} + \text{bias})` of shape ``[b, h, s_q, s_kv]``, and :math:`\text{max_logit}` is of shape ``[h]``. + name : Optional[str], default = None + module instance name. Parallelism parameters ---------------------- @@ -371,8 +374,9 @@ def __init__( softmax_scale: Optional[float] = None, softmax_type: str = "vanilla", return_max_logit: Optional[bool] = False, + name: Optional[str] = None, ) -> None: - super().__init__() + super().__init__(name=name) self.logger = logging.getLogger("DotProductAttention") self.logger.setLevel(attn_log._log_level) @@ -612,6 +616,7 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: # global recipe set in autocast() fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8_recipe.custom(): + super().init_fp8_metadata(num_gemms=num_gemms) return # switch/append recipe: fp8_recipe stays unchanged, but DPA.fp8_meta["recipe"] may be set to @@ -820,6 +825,9 @@ def init_fp8_metadata(self, num_gemms: int = 1) -> None: def set_meta_tensor(self, fwd: bool, recipe: Union[Recipe, List[Recipe]]) -> None: """Override to allow multiple recipes. Init scales and amaxes for fwd | bwd.""" + if isinstance(recipe, Recipe) and recipe.custom(): + TransformerEngineBaseModule.set_meta_tensor(self, fwd, recipe) + return if isinstance(recipe, Recipe): recipe = [recipe] fp8_recipe_dpa = recipe[-1] @@ -859,13 +867,127 @@ def set_meta_tensor(self, fwd: bool, recipe: Union[Recipe, List[Recipe]]) -> Non for i in range(len(recipe)) ] - self.fp8_meta[fp8_meta_tensor_key] = ( - recipe_states[-1] if len(recipe) == 2 else recipe_states[0] - ) + # Reached the rebuild path because ``fp8_meta_tensors_initialized`` + # was flipped to False after first init — most commonly because the + # base-class ``output_quantizer_role`` / ``grad_input_quantizer_role`` + # setter invalidated state when MHA wired boundary roles. That + # setter is recipe-agnostic, so this code fires for built-in + # recipes too even though they don't consume role information here + # (e.g. ``test_dpa_fp8_extra_state`` reaches this path with pure + # DelayedScaling). + # + # Rebuilding the recipe state must preserve persistent training + # buffers (delayed-scaling ``scale`` / ``amax_history``) so the new + # quantizer instances and the ``FP8GlobalStateManager`` reduction + # buffers end up viewing the SAME tensor objects, and so any + # checkpoint-loaded state isn't silently destroyed on the first + # forward after ``load_state_dict``. + # + # Inheritance targets the "primary" state stored under + # ``fp8_meta[fp8_meta_tensor_key]`` — the one tracked across + # ``set_meta_tensor`` calls. Auxiliary states in a multi-recipe + # splice (e.g. the CS half of ``[CS, DS]``) are stateless and have + # nothing to inherit. + old_state = self.fp8_meta.get(fp8_meta_tensor_key) + primary_idx = -1 if len(recipe) == 2 else 0 + if old_state is not None: + recipe_states[primary_idx].inherit_state_from(old_state) + + self.fp8_meta[fp8_meta_tensor_key] = recipe_states[primary_idx] self.quantizers[fp8_meta_tensor_key] = [] for recipe_state in recipe_states: self.quantizers[fp8_meta_tensor_key].extend(recipe_state.make_quantizers()) + def get_quantizer_roles( + self, + *, + fwd: bool, + num_quantizers: int, + ) -> Optional[List[QuantizerRole]]: + """QuantizerRole list for quantizers used by ``DotProductAttention``. + + DPA internally performs two matmuls:: + + S = softmax(Q · K^T) (GEMM1) + O = S · V (GEMM2) + + cuDNN's fused-attention API exposes FP8 scale/amax descriptors as + a flat array of **slot groups** numbered 1-3. The numbering is a + cuDNN convention — it does *not* correspond to operation order + inside DPA: + + Forward (3 slot groups × 3 positions = 9 slots): + + =========== =========================================== =========== + Slot group Primary tensor cuDNN enum + =========== =========================================== =========== + Group 1 QKV — inputs to GEMM1 (Q·K^T) GEMM1_OUTPUT + Group 2 O — output of GEMM2 (S·V) GEMM2_INPUT + Group 3 S — post-softmax, input to GEMM2 (S·V) GEMM3_OUTPUT + =========== =========================================== =========== + + Backward (3 slot groups × 2 positions = 6 slots): + + =========== =========================================== =========== + Slot group Primary tensor cuDNN enum + =========== =========================================== =========== + Group 1 dQKV — gradients flowing back to Q, K, V GRAD_OUTPUT1 + Group 2 dO — gradient of the attention output GRAD_INPUT2 + Group 3 dP — gradient of the softmax output GRAD_INPUT3 + =========== =========================================== =========== + + Unused positions within a group share the role of the group's + primary tensor. + + **Boundary slots** — O (fwd) and dQKV (bwd) leave DPA and enter + the next module (e.g. proj linear). DPA does not know that + consumer, so these default to ``None``. The parent module + (e.g. ``MultiheadAttention``) can set + :attr:`output_quantizer_role` / :attr:`grad_input_quantizer_role` + to fill in the consumer identity. + + When not set, a hint-only ``QuantizerRole`` with empty + ``module_type`` / ``tensor_type`` is emitted, with ``name`` + containing ``"dpa_output"`` or ``"dpa_grad_input"``. This lets + the factory return a DPA-compatible quantizer (required by the + fused kernel) even when the downstream consumer is unknown. + """ + name = self.name or "" + if fwd: + qkv_role = QuantizerRole(module_type="dpa", tensor_type="qkv", name=name) + o_role = self._output_quantizer_role + if o_role is None: + o_role = QuantizerRole(name=f"{name}.dpa_output" if name else "dpa_output") + s_role = QuantizerRole(module_type="dpa", tensor_type="s", name=name) + base = [ + qkv_role, + qkv_role, + qkv_role, # Group 1: QKV (inputs to Q·K^T) + o_role, + o_role, + o_role, # Group 2: O (output of S·V) — boundary + s_role, + s_role, + s_role, # Group 3: S (post-softmax, input to S·V) + ] + else: + dqkv_role = self._grad_input_quantizer_role + if dqkv_role is None: + dqkv_role = QuantizerRole( + name=f"{name}.dpa_grad_input" if name else "dpa_grad_input" + ) + do_role = QuantizerRole(module_type="dpa", tensor_type="do", name=name) + dp_role = QuantizerRole(module_type="dpa", tensor_type="dp", name=name) + base = [ + dqkv_role, + dqkv_role, # Group 1: dQKV (grads to Q,K,V) — boundary + do_role, + do_role, # Group 2: dO (grad of attention output) + dp_role, + dp_role, # Group 3: dP (grad of softmax output) + ] + return base[:num_quantizers] + @no_torch_dynamo(recursive=False) def forward( self, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 7df5daabe..1f1637cec 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -2334,7 +2334,7 @@ def check_set_window_size( return window_size -def get_attention_quantizers(fp8, fp8_recipe, quantizers): +def get_attention_quantizers(fp8, quantizers): """Get the list of quantizers used in attention from the quantizers list.""" if not fp8: return [None] * 6 @@ -2363,7 +2363,11 @@ def get_attention_quantizers(fp8, fp8_recipe, quantizers): dQKV_quantizer.internal = False dQKV_quantizer.set_usage(rowwise=True, columnwise=False) - if fp8_recipe.mxfp8(): + # MXFP8 attention: detect from the QKV quantizer instance rather than the + # recipe predicate so that CustomRecipe (whose `mxfp8()` predicate returns + # False) gets the same treatment as the built-in MXFP8 recipe. The kernel + # handles S/dP internally for MXFP8, hence S/dP are nulled out. + if isinstance(QKV_quantizer, MXFP8Quantizer): QKV_quantizer.columnwise_usage = True QKV_quantizer.optimize_for_gemm = True S_quantizer = None @@ -2374,6 +2378,29 @@ def get_attention_quantizers(fp8, fp8_recipe, quantizers): dP_quantizer = None dQKV_quantizer.columnwise_usage = True + _fp8_types = (Float8Quantizer, Float8CurrentScalingQuantizer, MXFP8Quantizer) + # S/dP are intentionally None under MXFP8 attention; skip the type check + # for those slots in that case. + _allow_none = {"S", "dP"} if isinstance(QKV_quantizer, MXFP8Quantizer) else set() + for _name, _q in [ + ("QKV", QKV_quantizer), + ("O", O_quantizer), + ("S", S_quantizer), + ("dQKV", dQKV_quantizer), + ("dO", dO_quantizer), + ("dP", dP_quantizer), + ]: + if _q is None and _name in _allow_none: + continue + assert isinstance(_q, _fp8_types), ( + "FP8 attention requires FP8-compatible quantizers for all DPA tensor slots, " + f"but {_name} quantizer is {type(_q).__name__}. " + "When using CustomRecipe with fp8_dpa=True, ensure the factory returns an " + "FP8 quantizer (Float8Quantizer, Float8CurrentScalingQuantizer, or " + "MXFP8Quantizer) for all DPA roles (module_type='dpa') and for None roles " + "(boundary slots like O output and dQKV grad-input)." + ) + return QKV_quantizer, O_quantizer, S_quantizer, dQKV_quantizer, dO_quantizer, dP_quantizer diff --git a/transformer_engine/pytorch/attention/multi_head_attention.py b/transformer_engine/pytorch/attention/multi_head_attention.py index afc4622b2..70ae9dfc2 100644 --- a/transformer_engine/pytorch/attention/multi_head_attention.py +++ b/transformer_engine/pytorch/attention/multi_head_attention.py @@ -8,7 +8,7 @@ from typing import Any, Callable, List, Optional, Tuple, Union import torch -from transformer_engine.pytorch.quantization import FP8GlobalStateManager +from transformer_engine.pytorch.quantization import FP8GlobalStateManager, QuantizerRole from transformer_engine.pytorch.tensor.float8_tensor import Float8Tensor from transformer_engine.pytorch.module.base import TransformerEngineBaseModule from transformer_engine.pytorch.module import LayerNormLinear, Linear, RMSNorm, LayerNorm @@ -461,6 +461,7 @@ def __init__( layer_number=self.layer_number, attention_type=self.attention_type, softmax_type=self.softmax_type, + name=name + ".core_attention" if name is not None else None, ) # Linear @@ -478,6 +479,84 @@ def __init__( **common_gemm_kwargs, ) + def _update_output_quantizer_roles( + self, + qkv_fp8_output: bool, + proj_fp8_grad: bool, + dpa_fp8_output: bool, + ) -> None: + """Set quantizer roles at the boundaries between QKV, DPA, and proj. + + MHA contains three submodules connected as follows:: + + Forward: QKV linear ──(QKV tensor)──> DPA ──(O tensor)──> Proj linear + Backward: QKV linear <──(dQKV tensor)── DPA <──(dO tensor)── Proj linear + + Each submodule owns quantizers for its internal tensors, but the + *boundary* tensors (the arrows above) need to know which module + will *consume* them so the quantizer factory can pick the right + format. This method sets those boundary roles on all four edges: + + 1. ``qkv_fp8_output`` — **QKV linear → DPA (fwd)**: the QKV + linear's ``output_quantizer_role`` is told its consumer is DPA. + 2. ``proj_fp8_grad`` — **Proj linear ← DPA (bwd)**: proj's + ``grad_input_quantizer_role`` is told its producer is DPA. + 3. ``dpa_fp8_output`` — **DPA → Proj linear (fwd)**: DPA's + ``output_quantizer_role`` is told its consumer is the proj linear. + 4. ``dpa_fp8_output`` — **DPA ← QKV linear (bwd)**: DPA's + ``grad_input_quantizer_role`` is told its consumer is QKV linear. + + When a flag is ``False`` the corresponding role is reset to ``None`` + so the module falls back to its own default. + """ + dpa_name = self.core_attention.name or "" + + # ── Boundary 1 (fwd): QKV linear output → consumed by DPA ──────── + qkv_output_role = ( + QuantizerRole(module_type="dpa", tensor_type="qkv", name=dpa_name) + if qkv_fp8_output + else None + ) + if self.attention_type == "self": + if self.input_layernorm: + self.layernorm_qkv.output_quantizer_role = qkv_output_role + else: + self.qkv.output_quantizer_role = qkv_output_role + elif self.attention_type == "cross": + if self.input_layernorm: + self.layernorm_query.output_quantizer_role = qkv_output_role + else: + self.query_layer.output_quantizer_role = qkv_output_role + self.key_value.output_quantizer_role = qkv_output_role + + # ── Boundary 2 (bwd): Proj grad-input ← produced by DPA ────────── + proj_grad_input_role = ( + QuantizerRole(module_type="dpa", tensor_type="do", name=dpa_name) + if proj_fp8_grad + else None + ) + self.proj.grad_input_quantizer_role = proj_grad_input_role + + # ── Boundary 3 (fwd): DPA output (O) → consumed by Proj linear ─── + proj_name = self.proj.name or "" + self.core_attention.output_quantizer_role = ( + QuantizerRole(module_type="linear", tensor_type="input", name=proj_name) + if dpa_fp8_output + else None + ) + + # ── Boundary 4 (bwd): DPA grad-input (dQKV) → consumed by QKV linear + if self.attention_type == "self": + qkv_linear = self.layernorm_qkv if self.input_layernorm else self.qkv + else: + qkv_linear = self.layernorm_query if self.input_layernorm else self.query_layer + qkv_name = qkv_linear.name or "" + self.core_attention.grad_input_quantizer_role = ( + QuantizerRole(module_type="linear", tensor_type="grad_output", name=qkv_name) + if dpa_fp8_output + else None + ) + def fast_setattr(self, name: str, value: Any) -> None: """Fast attribute set for non-parameter fields.""" self.__dict__[name] = value @@ -822,6 +901,8 @@ def forward( # 1. FP8CS recipe: produce F16 grads; again, due to cuBLAS limitation proj_fp8_grad = dpa_fp8_output and not float8_current_scaling + self._update_output_quantizer_roles(qkv_fp8_output, proj_fp8_grad, dpa_fp8_output) + layernorm_output = None if self.attention_type == "self": # Attention heads [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn] diff --git a/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py b/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py new file mode 100644 index 000000000..d660e5a53 --- /dev/null +++ b/transformer_engine/pytorch/custom_recipes/quantization_factory_examples.py @@ -0,0 +1,271 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +""" +Quantizer factory examples. + +Demonstrates how to use the ``CustomRecipe`` + ``qfactory`` interface to apply +*different* quantization recipes to different module/tensor types/instances within the same model. + +Usage:: + + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.quantization import autocast + from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( + nvfp4_linear_mxfp8_grouped_linear_factory, + nvfp4_linear_fp8_dpa_factory, + nvfp4_linear_mxfp8_dpa_factory, + ) + + # Mixed module types: NVFP4 for Linear, MXFP8 for GroupedLinear + recipe = CustomRecipe(qfactory=nvfp4_linear_mxfp8_grouped_linear_factory) + with autocast(recipe=recipe): + output = model(input) + + # NVFP4 for Linear, FP8 current-scaling + delayed-scaling for DPA + recipe = CustomRecipe(qfactory=nvfp4_linear_fp8_dpa_factory, fp8_dpa=True) + with autocast(recipe=recipe): + output = model(input) + + # NVFP4 for Linear, MXFP8 for DPA + recipe = CustomRecipe(qfactory=nvfp4_linear_mxfp8_dpa_factory, fp8_dpa=True) + with autocast(recipe=recipe): + output = model(input) +""" + +from __future__ import annotations + +from typing import Optional + +import transformer_engine_torch as tex + +from transformer_engine.pytorch.quantization import QuantizerRole + + +def nvfp4_linear_mxfp8_grouped_linear_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: NVFP4 for ``Linear``, MXFP8 for ``GroupedLinear``. + + Dispatch logic: + * ``role.module_type == "grouped_linear"`` -> MXFP8 (E4M3, block-32) + * everything else (``"linear"`` or unknown) -> NVFP4 (E2M1) + + NVFP4 settings follow the built-in ``NVFP4BlockScaling`` defaults: + * Weights: 2D quantization (16x16), no RHT, no stochastic rounding + * Inputs: 1D quantization, RHT enabled, no stochastic rounding + * Grads: 1D quantization, RHT enabled, stochastic rounding enabled + """ + is_grouped_linear = role is not None and role.module_type == "grouped_linear" + + if is_grouped_linear: + return _make_mxfp8_quantizer() + + return _make_nvfp4_quantizer(role) + + +def _make_mxfp8_quantizer(): + """Return an MXFP8 quantizer with default settings (E4M3, block-32, E8M0 scales).""" + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + return MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + ) + + +def _make_nvfp4_quantizer(role: Optional[QuantizerRole]): + """Return an NVFP4 quantizer configured per tensor role. + + Mirrors :class:`NVFP4BlockScaling` recipe defaults. + """ + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + + is_linear = role is not None and role.module_type == "linear" + is_weight = is_linear and role.tensor_type == "weight" + is_grad = is_linear and role.tensor_type == "grad_output" + + if is_weight: + return NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=True, + stochastic_rounding=False, + with_random_sign_mask=True, + ) + + if is_grad: + return NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_rht=True, + with_post_rht_amax=True, + with_2d_quantization=False, + stochastic_rounding=True, + with_random_sign_mask=True, + ) + + return NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_rht=True, + with_post_rht_amax=True, + with_2d_quantization=False, + stochastic_rounding=False, + with_random_sign_mask=True, + ) + + +def nvfp4_linear_fp8_dpa_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: NVFP4 for ``Linear``, mixed FP8 for ``DotProductAttention``. + + This factory demonstrates how to use ``CustomRecipe`` with ``fp8_dpa=True`` + to combine NVFP4 quantization for linear layers with FP8 attention. + + DPA tensor types (``role.module_type == "dpa"``): + + =========== ============================================================ + tensor_type Description + =========== ============================================================ + ``"qkv"`` Query, Key, Value inputs to the first attention GEMM + ``"s"`` Softmax output (S = softmax(Q·K^T)), fed into the second GEMM + ``"o"`` Attention output (O = S·V) + ``"do"`` Gradient of the attention output (dO), backward input + ``"dp"`` Gradient of the softmax output (dP = dO·V^T), backward + ``"dqkv"`` Gradient flowing back to Q, K, V + =========== ============================================================ + + Dispatch logic: + * ``role.module_type == "dpa"`` with ``tensor_type in ("s", "dp")`` + -> FP8 delayed scaling (stateful amax tracking) + * ``role.module_type == "dpa"`` (QKV, dO) + -> FP8 current scaling (E4M3) + * DPA boundary hints (``"dpa_output"`` / ``"dpa_grad_input"`` in ``role.name``) + -> FP8 current scaling placeholder. The fused attention kernel requires + FP8-compatible quantizers in all DPA slots, even when the output is + produced in BF16 (``fp8_mha=False``). DPA emits these hint-only roles + (with empty ``module_type`` and ``tensor_type``) when the downstream + consumer is unknown. + * everything else (``"linear"`` / ``"grouped_linear"`` / ``None``) + -> NVFP4 (E2M1), configured per tensor role + + Usage:: + + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.quantization import autocast + from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( + nvfp4_linear_fp8_dpa_factory, + ) + + recipe = CustomRecipe( + qfactory=nvfp4_linear_fp8_dpa_factory, + fp8_dpa=True, + ) + with autocast(recipe=recipe): + output = model(input) + """ + from transformer_engine.pytorch.quantization import DelayedScalingRequest + from transformer_engine.pytorch.tensor.float8_tensor import Float8CurrentScalingQuantizer + + is_dpa = role is not None and role.module_type == "dpa" + is_softmax_or_dp = is_dpa and role.tensor_type in ("s", "dp") + + if is_softmax_or_dp: + return DelayedScalingRequest() + + if is_dpa: + return Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + ) + + # DPA boundary slots (O output / dQKV grad-input): the fused attention + # kernel only supports FP8 quantizers here, regardless of the linear recipe. + is_dpa_boundary = ( + role is not None + and not role.module_type + and ("dpa_output" in role.name or "dpa_grad_input" in role.name) + ) + if is_dpa_boundary: + return Float8CurrentScalingQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + device="cuda", + ) + + return _make_nvfp4_quantizer(role) + + +def nvfp4_linear_mxfp8_dpa_factory( + role: Optional[QuantizerRole], +): + """Quantizer factory: NVFP4 for ``Linear``, MXFP8 for ``DotProductAttention``. + + Mirrors the documented "NVFP4 linear + MXFP8 attention" combo from + :mod:`transformer_engine.pytorch.attention.dot_product_attention.dot_product_attention` + (see the recipe-combination table at the top of that module). With + ``CustomRecipe`` the per-tensor decision is made directly here, so the + ``NVTE_DPA_FP8_RECIPE="MXFP8BlockScaling"`` env override that the + built-in recipes would otherwise need is unnecessary. + + DPA tensor types (``role.module_type == "dpa"``): + + =========== ============================================================ + tensor_type Description + =========== ============================================================ + ``"qkv"`` Query, Key, Value inputs to the first attention GEMM + ``"s"`` Softmax output (S = softmax(Q·K^T)), fed into the second GEMM + ``"o"`` Attention output (O = S·V) + ``"do"`` Gradient of the attention output (dO), backward input + ``"dp"`` Gradient of the softmax output (dP = dO·V^T), backward + ``"dqkv"`` Gradient flowing back to Q, K, V + =========== ============================================================ + + Dispatch logic: + * ``role.module_type == "dpa"`` -> MXFP8 (E4M3, block-32) + The MXFP8 fused-attention kernel handles the S/dP slots + internally, so any quantizer returned for those roles is later + nulled out by ``get_attention_quantizers``. Returning MXFP8 is + the simplest valid choice. + * DPA boundary hints (``"dpa_output"`` / ``"dpa_grad_input"`` in + ``role.name``) -> MXFP8 placeholder. The fused attention kernel + requires FP8-compatible quantizers in all DPA slots. + * everything else (``"linear"`` / ``"grouped_linear"`` / ``None``) + -> NVFP4 (E2M1), configured per tensor role. + + Usage:: + + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.quantization import autocast + from transformer_engine.pytorch.custom_recipes.quantization_factory_examples import ( + nvfp4_linear_mxfp8_dpa_factory, + ) + + recipe = CustomRecipe( + qfactory=nvfp4_linear_mxfp8_dpa_factory, + fp8_dpa=True, + ) + with autocast(recipe=recipe): + output = model(input) + """ + is_dpa = role is not None and role.module_type == "dpa" + if is_dpa: + return _make_mxfp8_quantizer() + + # DPA boundary slots (O output / dQKV grad-input): emitted by DPA with + # empty `module_type` and a `name` like ".dpa_output". The fused + # attention kernel requires an FP8-compatible quantizer here even when + # the downstream consumer is unknown. + is_dpa_boundary = ( + role is not None + and not role.module_type + and ("dpa_output" in role.name or "dpa_grad_input" in role.name) + ) + if is_dpa_boundary: + return _make_mxfp8_quantizer() + + return _make_nvfp4_quantizer(role) diff --git a/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py b/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py new file mode 100644 index 000000000..22eafaa66 --- /dev/null +++ b/transformer_engine/pytorch/custom_recipes/quantization_recipes_base.py @@ -0,0 +1,179 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +""" +Quantizer factory examples using real silicon quantizers. + +Each factory below replicates the behaviour of built-in TE recipe but via the +``CustomRecipe`` + ``qfactory`` interface. This is useful when you want to +start from a known-good recipe and then selectively override quantizer settings +for specific layers / tensor types. + +Usage (any factory):: + + from transformer_engine.common.recipe import CustomRecipe + from transformer_engine.pytorch.quantization import autocast + from transformer_engine.pytorch.custom_recipes.quantization_recipes_base import ( + nvfp4_quantizer_factory, + ) + + recipe = CustomRecipe(qfactory=nvfp4_quantizer_factory) + with autocast(recipe=recipe): + output = model(input) +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import transformer_engine_torch as tex + +from transformer_engine.pytorch.quantization import QuantizerRole + + +def delayed_scaling_quantizer_factory( + role: Optional[QuantizerRole], # pylint: disable=unused-argument +) -> "DelayedScalingRequest": + """Factory that mirrors :class:`DelayedScaling` recipe defaults. + + Returns a :class:`DelayedScalingRequest` for every slot. TE allocates + shared scale/amax_history buffers and wires them into the existing + delayed-scaling reduction path. + + * HYBRID format: E4M3 forward, E5M2 backward + * amax_history_len = 1024 + * reduce_amax = True + """ + from transformer_engine.pytorch.quantization import DelayedScalingRequest + from transformer_engine.common.recipe import Format + + return DelayedScalingRequest(fp8_format=Format.HYBRID) + + +def current_scaling_quantizer_factory( + role: Optional[QuantizerRole], +) -> "Float8CurrentScalingQuantizer": + """Factory that mirrors :class:`Float8CurrentScaling` recipe defaults. + + * Forward tensors (input, weight) → E4M3 + * Backward tensors (grad_output) → E5M2 + """ + from transformer_engine.pytorch.tensor.float8_tensor import ( + Float8CurrentScalingQuantizer, + ) + + is_backward = role is not None and role.tensor_type == "grad_output" + fp8_dtype = tex.DType.kFloat8E5M2 if is_backward else tex.DType.kFloat8E4M3 + + return Float8CurrentScalingQuantizer( + fp8_dtype=fp8_dtype, + device=torch.device("cuda"), + force_pow_2_scales=False, # constrain scale to powers of 2 + amax_epsilon=0.0, # clamp amax from below to avoid div-by-zero + ) + + +def mxfp8_quantizer_factory( + role: Optional[QuantizerRole], # pylint: disable=unused-argument +) -> "MXFP8Quantizer": + """Factory that mirrors :class:`MXFP8BlockScaling` recipe defaults. + + * E4M3 by default for all tensors + * Block size 32, power-of-2 (E8M0) scales + """ + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + return MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + ) + + +def float8_block_scaling_quantizer_factory( + role: Optional[QuantizerRole], +) -> "Float8BlockQuantizer": + """Factory that mirrors :class:`Float8BlockScaling` recipe defaults. + + * E4M3 by default for all tensors + * Weights use 2D block scaling, everything else uses 1D + * Power-of-2 scales by default + """ + from transformer_engine.pytorch.tensor.float8_blockwise_tensor import ( + Float8BlockQuantizer, + ) + + is_weight = ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type == "weight" + ) + block_scaling_dim = 2 if is_weight else 1 + + return Float8BlockQuantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + amax_epsilon=0.0, # clamp amax from below to avoid div-by-zero + force_pow_2_scales=True, + block_scaling_dim=block_scaling_dim, # 1 = 1D (1×128), 2 = 2D (128×128) + ) + + +def nvfp4_quantizer_factory( + role: Optional[QuantizerRole], +) -> "NVFP4Quantizer": + """Factory that mirrors :class:`NVFP4BlockScaling` recipe defaults. + + * All tensors quantized to E2M1 (FP4) + * Weights: 2D quantization (16x16 blocks), no RHT, no stochastic rounding + * Inputs: 1D quantization, RHT enabled, no stochastic rounding + * Grads: 1D quantization, RHT enabled, stochastic rounding enabled + + Quantizer knobs: + fp4_dtype - E2M1 (only supported format) + with_rht - randomized Hadamard transform (smooths outliers) + with_post_rht_amax - recompute amax after RHT (should match with_rht) + with_2d_quantization - 16x16 2D blocks (vs 1x16 1D) + stochastic_rounding - probabilistic rounding to reduce quant bias + with_random_sign_mask - random sign flip in the Hadamard matrix + """ + from transformer_engine.pytorch.tensor.nvfp4_tensor import NVFP4Quantizer + + is_linear = role is not None and role.module_type in ("linear", "grouped_linear") + is_weight = is_linear and role.tensor_type == "weight" + is_grad = is_linear and role.tensor_type == "grad_output" + + if is_weight: + return NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=True, + stochastic_rounding=False, + with_random_sign_mask=True, + ) + + if is_grad: + return NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_rht=True, + with_post_rht_amax=True, + with_2d_quantization=False, + stochastic_rounding=True, + with_random_sign_mask=True, + ) + + # For input and unknown roles + return NVFP4Quantizer( + fp4_dtype=tex.DType.kFloat4E2M1, + rowwise=True, + columnwise=True, + with_rht=True, + with_post_rht_amax=True, + with_2d_quantization=False, + stochastic_rounding=False, + with_random_sign_mask=True, + ) diff --git a/transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py b/transformer_engine/pytorch/custom_recipes/quantization_ref_current_scaling.py similarity index 98% rename from transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py rename to transformer_engine/pytorch/custom_recipes/quantization_ref_current_scaling.py index 8580cf4a3..ecbb667ec 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_current_scaling.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_ref_current_scaling.py @@ -18,17 +18,18 @@ def current_scaling_ref_quantizer_factory(role): """Factory function for current scaling reference quantizer. - Usage with CustomRecipe and autocast: + Receives a :class:`~transformer_engine.pytorch.quantization.QuantizerRole`. + + Backward tensors use E5M2, everything else uses E4M3. + + Usage with CustomRecipe and autocast:: + custom_recipe = recipe.CustomRecipe(qfactory=current_scaling_ref_quantizer_factory) with autocast(recipe=custom_recipe): output = model(input) """ - if role in ("linear_input", "linear_weight"): - dtype = torch.float8_e4m3fn - elif role in ("linear_output", "linear_grad_output"): - dtype = torch.float8_e5m2 - else: - return None + is_backward = role is not None and role.tensor_type == "grad_output" + dtype = torch.float8_e5m2 if is_backward else torch.float8_e4m3fn return CurrentScalingQuantizerRef( dtype=dtype, rowwise=True, diff --git a/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py b/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py similarity index 98% rename from transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py rename to transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py index 12f8ef8f5..acb7abefd 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py @@ -18,33 +18,32 @@ def nvfp4_ref_rht_2d_quantizer_factory(role): """ Quantizer factory for NVFP4 recipe reference implementation (RHT and 2D quantization for weights). - Usage with CustomRecipe and autocast: + Receives a :class:`~transformer_engine.pytorch.quantization.QuantizerRole`. + + Usage with CustomRecipe and autocast:: + custom_recipe = recipe.CustomRecipe(qfactory=nvfp4_ref_rht_2d_quantizer_factory) - with autocast(fp8_recipe=custom_recipe): + with autocast(recipe=custom_recipe): output = model(input) """ - if role == "linear_input": - return NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, - quant_tile_shape=(1, 16), - pow_2_scales=False, - with_rht=True, - ) - if role == "linear_weight": + is_weight_tensor_in_gemm = ( + role is not None + and role.module_type in ("linear", "grouped_linear") + and role.tensor_type == "weight" + ) + if is_weight_tensor_in_gemm: # 2D quantization for weights in GEMM-based modules return NVFP4QuantizerRef( dtype=utils.Fp4Formats.E2M1, quant_tile_shape=(16, 16), pow_2_scales=False, with_rht=False, ) - if role == "linear_grad_output": - return NVFP4QuantizerRef( - dtype=utils.Fp4Formats.E2M1, - quant_tile_shape=(1, 16), - pow_2_scales=False, - with_rht=True, - ) - return None + return NVFP4QuantizerRef( + dtype=utils.Fp4Formats.E2M1, + quant_tile_shape=(1, 16), + pow_2_scales=False, + with_rht=True, + ) def cast_to_fp4x2(x): diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index e6bedee0c..746177ec7 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -27,8 +27,11 @@ Float8CurrentScalingRecipeState, Float8BlockScalingRecipeState, NVFP4BlockScalingRecipeState, + CustomRecipeState, FP8GlobalStateManager, + QuantizerRole, RecipeState, + _has_delayed_scaling_state, ) from ..distributed import ( gather_along_first_dim, @@ -789,6 +792,8 @@ def __init__(self, name: Optional[str] = None) -> None: self.activation_dtype: Optional[torch.dtype] = None self.wgrad_accumulation_and_reduce_hooks = [] self.wgrad_store = None + self._output_quantizer_role: Optional[QuantizerRole] = None + self._grad_input_quantizer_role: Optional[QuantizerRole] = None if not TEDebugState.debug_enabled: TEDebugState.initialize() @@ -809,6 +814,72 @@ def module_setattr(self, name: str, value: Any) -> None: """ super().__setattr__(name, value) + @property + def output_quantizer_role(self) -> Optional[QuantizerRole]: + """Caller-configurable :class:`QuantizerRole` for the forward output quantizer. + + When set, overrides the default role used by :meth:`get_quantizer_roles` + for the forward-pass output quantizer slot. Setting this after + quantizers have been created forces their recreation on the next + forward pass. + + See also :attr:`grad_input_quantizer_role` for the backward-pass + counterpart. + """ + return self._output_quantizer_role + + @output_quantizer_role.setter + def output_quantizer_role(self, role: Optional[QuantizerRole]) -> None: + if role == self._output_quantizer_role: + return + self._output_quantizer_role = role + if self.fp8_meta_tensors_initialized: + self.fp8_meta_tensors_initialized = False + + @property + def grad_input_quantizer_role(self) -> Optional[QuantizerRole]: + """Caller-configurable :class:`QuantizerRole` for the grad-input quantizer. + + Backward-pass counterpart of :attr:`output_quantizer_role`. + """ + return self._grad_input_quantizer_role + + @grad_input_quantizer_role.setter + def grad_input_quantizer_role(self, role: Optional[QuantizerRole]) -> None: + if role == self._grad_input_quantizer_role: + return + self._grad_input_quantizer_role = role + if self.fp8_meta_tensors_initialized: + self.fp8_meta_tensors_initialized = False + + def _warn_missing_output_quantizer_role( + self, + fp8_output: bool, + fp8_grad: bool, + ) -> None: + """Warn when quantized output is requested but no consumer role is set. + + Only relevant for ``CustomRecipe`` where the ``qfactory`` dispatches + on roles. Built-in recipes ignore role metadata. + """ + recipe = FP8GlobalStateManager.get_fp8_recipe() + if not recipe.custom(): + return + if fp8_output and self._output_quantizer_role is None: + warnings.warn( + f"{type(self).__name__}: fp8_output=True but " + "output_quantizer_role is not set. The CustomRecipe qfactory " + "will receive None for the output quantizer role.", + stacklevel=3, + ) + if fp8_grad and self._grad_input_quantizer_role is None: + warnings.warn( + f"{type(self).__name__}: fp8_grad=True but " + "grad_input_quantizer_role is not set. The CustomRecipe " + "qfactory will receive None for the grad-input quantizer role.", + stacklevel=3, + ) + @property def is_fsdp2(self) -> bool: """Whether this module is wrapped with FSDP2.""" @@ -901,21 +972,124 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: return if recipe.nvfp4() and isinstance(recipe_state, NVFP4BlockScalingRecipeState): return + if recipe.custom() and isinstance(recipe_state, CustomRecipeState): + return # Max. number of fp8 tensors per GEMM = 3 (input, weight, output) for fwd and # 2 (grad_output and grad_input) for bwd num_fp8_tensors = self.fp8_meta["num_gemms"] * 3 if fwd else self.fp8_meta["num_gemms"] * 2 # Initialize recipe state and quantizers - recipe_state = RecipeState.create( + roles = self.get_quantizer_roles( # pylint: disable=assignment-from-none + fwd=fwd, num_quantizers=num_fp8_tensors + ) + if roles is not None: + assert ( + len(roles) == num_fp8_tensors + ), f"Recipe roles must match number of quantizers ({len(roles)=} vs {num_fp8_tensors=})" + recipe_state = RecipeState.create( # pylint: disable=assignment-from-none recipe, mode=("forward" if fwd else "backward"), num_quantizers=num_fp8_tensors, + roles=roles, ) + # Reached the rebuild path because ``fp8_meta_tensors_initialized`` + # was flipped to False after first init — most commonly because the + # ``output_quantizer_role`` / ``grad_input_quantizer_role`` setter + # invalidated state when a parent module (e.g. ``MultiheadAttention``) + # wired boundary roles. That setter is recipe-agnostic, so this code + # fires even for built-in recipes that don't consume role information + # in ``make_quantizers``. + # + # Rebuilding the recipe state must preserve persistent training + # buffers (delayed-scaling ``scale`` / ``amax_history``) so the new + # quantizer instances and the ``FP8GlobalStateManager`` reduction + # buffers end up viewing the SAME tensor objects, and so any + # checkpoint-loaded state isn't silently destroyed on the first + # forward after ``load_state_dict``. + old_state = self.fp8_meta.get(fp8_meta_tensor_key) + if old_state is not None: + recipe_state.inherit_state_from(old_state) + self.fp8_meta[fp8_meta_tensor_key] = recipe_state self.quantizers[fp8_meta_tensor_key] = recipe_state.make_quantizers() + def get_quantizer_roles( + self, + *, + fwd: bool, # pylint: disable=unused-argument + num_quantizers: int, # pylint: disable=unused-argument + ) -> Optional[List[QuantizerRole]]: + """Return an ordered list of :class:`QuantizerRole` for quantizers. + + Overview + -------- + When using ``CustomRecipe``, the quantizer factory is called once + per quantizer slot. Each call receives a ``QuantizerRole`` that + tells the factory *what* is being quantized so it can return the + right quantizer. + + This method builds the role list. Subclasses override it to + describe their internal GEMM layout. + + Slot layout + ----------- + Return one ``QuantizerRole`` per slot, in the same order as the + module's quantizer array. For example, ``Linear`` uses 3 + forward slots ``[input, weight, output]`` and 2 backward slots + ``[grad_output, grad_input]``. Multi-GEMM modules like + ``LayerNormMLP`` repeat that pattern for each GEMM: + ``[fc1_input, fc1_weight, fc1_output, fc2_input, fc2_weight, fc2_output]``. + + What to put in each slot + ------------------------ + Create a ``QuantizerRole(module_type=..., tensor_type=..., + name=...)`` for each slot: + + * **module_type** — the kind of module: ``"linear"``, + ``"grouped_linear"``, ``"dpa"``, etc. The factory can dispatch + on this to use different quantization formats per module type. + * **tensor_type** — what tensor this slot holds, in the module's + own vocabulary. For linears: ``"input"``, ``"weight"``, + ``"grad_output"``, etc. For DPA: ``"qkv"``, ``"s"``, + ``"do"``, ``"dp"``, etc. + * **name** — the instance name (e.g. ``"encoder.layer0.fc1"``), + propagated from the ``name=`` constructor argument. The factory + can dispatch on this to target specific layers. + + Boundary slots + -------------- + The last slot of a forward GEMM group (output) and the last slot + of a backward group (grad_input) are **boundary** slots — the + tensor leaves this module and enters an unknown consumer. For + these slots, use ``self._output_quantizer_role`` (fwd) and + ``self._grad_input_quantizer_role`` (bwd), which default to + ``None``. A parent module (e.g. ``MultiheadAttention``) can set + these attributes to fill in the consumer identity; see + ``MultiheadAttention._update_output_quantizer_roles`` for an + example. + + Return value + ------------ + * A list of ``QuantizerRole`` with length ``num_quantizers``. + * ``None`` to opt out of role-based dispatch. + + Not implemented (default) + ~~~~~~~~~~~~~~~~~~~~~~~~~ + The base implementation returns ``None``. When ``None`` is + returned, ``CustomRecipeState`` emits a warning and falls back + to bare ``QuantizerRole()`` instances (all fields empty) for + every slot. The factory still gets called once per slot, but + every call receives an identical empty role — it cannot + distinguish input from weight, forward from backward, or one + module from another. What happens then depends entirely on the + factory: it may return the same quantizer for all slots (acting + as a uniform recipe), or it may raise an error if it requires + meaningful roles to dispatch on. + """ + return None + def _update_weight_quantizers(self) -> None: """Update the quantizers for the weight tensors.""" weight_tensors = self._get_weight_tensors() @@ -1024,7 +1198,7 @@ def to_cpu(src: torch.Tensor) -> torch.Tensor: # Copy tensors to CPU and store state = {} state["recipe"] = self.fp8_meta["recipe"] - if state["recipe"].delayed(): + if _has_delayed_scaling_state(self.fp8_meta): state["scale_fwd"] = to_cpu(self.fp8_meta["scaling_fwd"].scale) state["amax_history_fwd"] = to_cpu(self.fp8_meta["scaling_fwd"].amax_history) state["scale_bwd"] = to_cpu(self.fp8_meta["scaling_bwd"].scale) @@ -1096,7 +1270,7 @@ def copy_tensor(src: torch.Tensor, dst: torch.Tensor) -> None: dst.copy_(src, non_blocking=True) # Load tensors - if self.fp8_meta["recipe"].delayed(): + if _has_delayed_scaling_state(self.fp8_meta): copy_tensor(state["scale_fwd"], self.fp8_meta["scaling_fwd"].scale) copy_tensor(state["amax_history_fwd"], self.fp8_meta["scaling_fwd"].amax_history) copy_tensor(state["scale_bwd"], self.fp8_meta["scaling_bwd"].scale) @@ -1223,7 +1397,7 @@ def prepare_forward( # Activation recomputation is used and this is the second forward phase. if self.fp8 and in_fp8_activation_recompute_phase(): - delayed_scaling_recipe = self.fp8_meta["recipe"].delayed() + delayed_scaling_recipe = _has_delayed_scaling_state(self.fp8_meta) FP8GlobalStateManager.get_old_fp8_meta_tensors_for_recompute(self.fp8_meta) else: if not inp.is_cuda: @@ -1242,14 +1416,15 @@ def prepare_forward( self.init_fp8_metadata(num_gemms=num_gemms) self._check_weight_tensor_recipe_correspondence() - delayed_scaling_recipe = self.fp8 and self.fp8_meta["recipe"].delayed() + delayed_scaling_recipe = self.fp8 and _has_delayed_scaling_state(self.fp8_meta) if delayed_scaling_recipe: if self.sequence_parallel: - if not self.fp8_meta["recipe"].reduce_amax: - raise ValueError( - "Amax reduction across tensor parallel group is " - "necessary when using sequence parallelism with FP8." - ) + assert ( + self.fp8_meta["recipe"].custom() or self.fp8_meta["recipe"].reduce_amax + ), ( + "Amax reduction across tensor parallel group is " + "necessary when using sequence parallelism with FP8." + ) if not FP8GlobalStateManager.fp8_graph_capturing(): FP8GlobalStateManager.add_fp8_tensors_to_global_buffer(self.fp8_meta) @@ -1268,7 +1443,7 @@ def end_forward(self): Required to be called at the end of the forward function to properly handle DelayedScaling metadata handling and the NVTX ranges. """ - delayed_scaling_recipe = self.fp8 and self.fp8_meta["recipe"].delayed() + delayed_scaling_recipe = self.fp8 and _has_delayed_scaling_state(self.fp8_meta) if delayed_scaling_recipe and self.fp8 and in_fp8_activation_recompute_phase(): FP8GlobalStateManager.restore_fp8_meta_tensors(self.fp8_meta) nvtx_range_pop() diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 4ae7b47b9..e950f2657 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -24,7 +24,7 @@ _2X_ACC_WGRAD, ) from ._common import WeightGradStore -from ..quantization import FP8GlobalStateManager +from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( divide, cast_if_needed, @@ -116,7 +116,18 @@ def forward( # Configure quantizers if save_original_input and isinstance(input_quantizers[0], Float8Quantizer): - raise ValueError("DelayedScaling recipe is not supported with save_original_input") + if FP8GlobalStateManager.get_fp8_recipe().custom(): + # Custom recipe factory may produce DS quantizers unknown to caller. + # TODO(negvet): fix on Megatron side — guard should also exclude 'custom', or + # better: check at runtime whether quantizers are DS-based. + warnings.warn( + "save_original_input is incompatible with delayed-scaling quantizers " + "(Float8Quantizer). Disabling save_original_input for this module.", + stacklevel=2, + ) + save_original_input = False + else: + raise ValueError("DelayedScaling recipe is not supported with save_original_input") if input_quantizers[0] is not None: for input_quantizer in input_quantizers: input_quantizer.set_usage( @@ -829,6 +840,33 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: if recipe.float8_current_scaling(): self._customize_quantizers_float8_current_scaling(fwd, recipe) + def get_quantizer_roles( + self, + *, + fwd: bool, + num_quantizers: int, + ) -> Optional[List[QuantizerRole]]: + """QuantizerRole list for quantizers used by ``GroupedLinear``. + + For grouped GEMMs we repeat the same pattern for each GEMM in + order. The output (fwd) and grad-input (bwd) slots default to + ``None`` (unknown consumer). Set :attr:`output_quantizer_role` / + :attr:`grad_input_quantizer_role` to provide consumer identity. + """ + name = self.name or "" + if fwd: + base = [ + QuantizerRole(module_type="grouped_linear", tensor_type="input", name=name), + QuantizerRole(module_type="grouped_linear", tensor_type="weight", name=name), + self._output_quantizer_role, + ] + else: + base = [ + QuantizerRole(module_type="grouped_linear", tensor_type="grad_output", name=name), + self._grad_input_quantizer_role, + ] + return [base[i % len(base)] for i in range(num_quantizers)] + def make_grouped_weights(self, defer_init=False) -> None: """ Convert parameters into a GroupedTensor and re-register them as parameters. diff --git a/transformer_engine/pytorch/module/layernorm_linear.py b/transformer_engine/pytorch/module/layernorm_linear.py index abfa6af03..8c88f3ee8 100644 --- a/transformer_engine/pytorch/module/layernorm_linear.py +++ b/transformer_engine/pytorch/module/layernorm_linear.py @@ -28,7 +28,7 @@ _2X_ACC_DGRAD, _2X_ACC_WGRAD, ) -from ..quantization import FP8GlobalStateManager +from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( assert_dim_for_fp8_exec, cast_if_needed, @@ -1504,6 +1504,32 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: elif recipe.nvfp4(): self._customize_quantizers_nvfp4(fwd, recipe) + def get_quantizer_roles( + self, + *, + fwd: bool, + num_quantizers: int, + ) -> Optional[List[QuantizerRole]]: + """QuantizerRole list for quantizers used by ``LayerNormLinear``. + + The output (fwd) and grad-input (bwd) slots default to ``None`` + (unknown consumer). Set :attr:`output_quantizer_role` / + :attr:`grad_input_quantizer_role` to provide consumer identity. + """ + name = self.name or "" + if fwd: + base = [ + QuantizerRole(module_type="linear", tensor_type="input", name=name), + QuantizerRole(module_type="linear", tensor_type="weight", name=name), + self._output_quantizer_role, + ] + else: + base = [ + QuantizerRole(module_type="linear", tensor_type="grad_output", name=name), + self._grad_input_quantizer_role, + ] + return [base[i % len(base)] for i in range(num_quantizers)] + def reset_layer_norm_parameters(self) -> None: """Init LN params""" warnings.warn( @@ -1713,6 +1739,9 @@ def forward( def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): if not self.fp8: return [None] * 6 + + self._warn_missing_output_quantizer_role(fp8_output, fp8_grad) + grad_input_quantizer = None grad_weight_quantizer = None grad_output_quantizer = None diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 4fa7eb285..46918ff0f 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -29,7 +29,7 @@ _2X_ACC_DGRAD, _2X_ACC_WGRAD, ) -from ..quantization import FP8GlobalStateManager +from ..quantization import FP8GlobalStateManager, QuantizerRole from ..jit import ( bias_gelu_fused, bgrad_dgelu_fused, @@ -2104,6 +2104,53 @@ def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: elif recipe.nvfp4(): self._customize_quantizers_nvfp4(fwd, recipe) + def get_quantizer_roles( + self, + *, + fwd: bool, + num_quantizers: int, + ) -> Optional[List[QuantizerRole]]: + """QuantizerRole list for quantizers used by ``LayerNormMLP``. + + Each internal GEMM (fc1, fc2) gets a distinct name suffix so that + custom-recipe factories can target them individually. + + The module's final output (fc2 fwd) and final grad (fc1 bwd) + slots default to ``None`` (unknown consumer). Set + :attr:`output_quantizer_role` / :attr:`grad_input_quantizer_role` + to provide consumer identity. Internal boundaries use fixed + roles with known consumer identity. + """ + base_name = self.name or "" + fc1_name = f"{base_name}.fc1" if base_name else "fc1" + fc2_name = f"{base_name}.fc2" if base_name else "fc2" + # Roles use the *consumer's* identity: internal boundary tensors are + # labeled with the downstream module that will consume them. + # + # Forward: fc1_input -> fc1 GEMM -> [act] -> fc2_input -> fc2 GEMM -> output + # Backward: grad_input <- fc1 GEMM <- [act'] <- fc2 GEMM <- grad_output + if fwd: + base = [ + QuantizerRole(module_type="linear", tensor_type="input", name=fc1_name), + QuantizerRole(module_type="linear", tensor_type="weight", name=fc1_name), + # fc1 output — consumed by fc2 (via activation), so labeled as fc2 input + QuantizerRole(module_type="linear", tensor_type="input", name=fc2_name), + QuantizerRole(module_type="linear", tensor_type="input", name=fc2_name), + QuantizerRole(module_type="linear", tensor_type="weight", name=fc2_name), + # fc2 output — boundary, consumer unknown + self._output_quantizer_role, + ] + else: + base = [ + QuantizerRole(module_type="linear", tensor_type="grad_output", name=fc1_name), + # fc1 grad_input — boundary, consumer unknown + self._grad_input_quantizer_role, + QuantizerRole(module_type="linear", tensor_type="grad_output", name=fc2_name), + # fc2 grad_input — consumed by fc1 (via activation'), so labeled as fc1 grad_output + QuantizerRole(module_type="linear", tensor_type="grad_output", name=fc1_name), + ] + return [base[i % len(base)] for i in range(num_quantizers)] + def reset_layer_norm_parameters(self) -> None: """Init LN params""" warnings.warn( @@ -2336,6 +2383,9 @@ def forward( return out def _get_quantizers(self, fp8_output, is_grad_enabled): + if self.fp8: + self._warn_missing_output_quantizer_role(fp8_output, False) + ( fc1_input_quantizer, fc1_output_quantizer, diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index 2b14eaaf2..e725387e7 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -27,7 +27,7 @@ _2X_ACC_WGRAD, ) from ._common import noop_cat, WeightGradStore -from ..quantization import FP8GlobalStateManager +from ..quantization import FP8GlobalStateManager, QuantizerRole from ..utils import ( cast_if_needed, clear_tensor_data, @@ -1510,6 +1510,32 @@ def __init__( if name in self.weight_names or name in self.bias_names: param.skip_backward_post_hook = True + def get_quantizer_roles( + self, + *, + fwd: bool, + num_quantizers: int, + ) -> Optional[List[QuantizerRole]]: + """QuantizerRole list for quantizers used by ``Linear``. + + The output (fwd) and grad-input (bwd) slots default to ``None`` + (unknown consumer). Set :attr:`output_quantizer_role` / + :attr:`grad_input_quantizer_role` to provide consumer identity. + """ + name = self.name or "" + if fwd: + base = [ + QuantizerRole(module_type="linear", tensor_type="input", name=name), + QuantizerRole(module_type="linear", tensor_type="weight", name=name), + self._output_quantizer_role, + ] + else: + base = [ + QuantizerRole(module_type="linear", tensor_type="grad_output", name=name), + self._grad_input_quantizer_role, + ] + return [base[i % len(base)] for i in range(num_quantizers)] + def set_meta_tensor(self, fwd: bool, recipe: Recipe) -> None: """Init scales and amaxes for fwd | bwd.""" super().set_meta_tensor(fwd, recipe) @@ -1724,6 +1750,9 @@ def forward( def _get_quantizers(self, fp8_output, fp8_grad, is_grad_enabled): if not self.fp8: return [None] * 6 + + self._warn_missing_output_quantizer_role(fp8_output, fp8_grad) + grad_input_quantizer = None grad_weight_quantizer = None grad_output_quantizer = None diff --git a/transformer_engine/pytorch/ops/basic/basic_linear.py b/transformer_engine/pytorch/ops/basic/basic_linear.py index 41f0855f1..95e044030 100644 --- a/transformer_engine/pytorch/ops/basic/basic_linear.py +++ b/transformer_engine/pytorch/ops/basic/basic_linear.py @@ -19,7 +19,7 @@ gather_along_first_dim, reduce_scatter_along_first_dim, ) -from ...quantization import FP8GlobalStateManager, Recipe +from ...quantization import FP8GlobalStateManager, QuantizerRole, Recipe from ...module.base import ( _2X_ACC_FPROP, _2X_ACC_DGRAD, @@ -275,6 +275,21 @@ def num_quantizers(self, mode: str) -> int: return 1 return 0 + def get_quantizer_roles(self, mode: str) -> Optional[list[QuantizerRole]]: + name = getattr(self, "name", "") or "" + if mode == "forward": + # BasicLinear owns input and weight quantizers. + # Output quantizer is provided by the next op (as its input quantizer). + return [ + QuantizerRole(module_type="linear", tensor_type="input", name=name), + QuantizerRole(module_type="linear", tensor_type="weight", name=name), + ] + if mode == "backward": + # BasicLinear owns grad_output quantizer. + # Grad_input quantizer is provided by the previous op (as its grad_output quantizer). + return [QuantizerRole(module_type="linear", tensor_type="grad_output", name=name)] + return None + def reset_parameters(self) -> None: """Initialize parameter buffers and values""" diff --git a/transformer_engine/pytorch/ops/op.py b/transformer_engine/pytorch/ops/op.py index c5c8ea346..168718723 100644 --- a/transformer_engine/pytorch/ops/op.py +++ b/transformer_engine/pytorch/ops/op.py @@ -16,6 +16,7 @@ from transformer_engine.common.recipe import Recipe from ..quantization import ( FP8GlobalStateManager, + QuantizerRole, RecipeState, autocast, ) @@ -209,6 +210,17 @@ def num_quantizers( """ return 0 + def get_quantizer_roles( + self, mode: str # pylint: disable=unused-argument + ) -> Optional[list[QuantizerRole]]: + """Return an ordered list of :class:`QuantizerRole` for quantizers. + + The returned list must be aligned with the internal quantizer ordering and + must have length ``num_quantizers(mode)`` for supported modes. + Returning ``None`` means "no explicit roles". + """ + return None + def get_input_quantizer(self) -> Optional[Quantizer]: if self.num_quantizers("forward") > 0: return self.get_quantizer("forward", 0) @@ -268,10 +280,17 @@ def reset_recipe_state( ) # Construct quantization recipe state - recipe_state = RecipeState.create( + roles = self.get_quantizer_roles(mode) # pylint: disable=assignment-from-none + if roles is not None: + assert len(roles) == num_quantizers, ( + "Recipe roles must match number of quantizers " + f"({len(roles)=} vs {num_quantizers=})" + ) + recipe_state = RecipeState.create( # pylint: disable=assignment-from-none recipe, mode=mode, num_quantizers=num_quantizers, + roles=roles, ) fp8_meta_key = FP8GlobalStateManager.get_meta_tensor_key( forward=(mode == "forward"), diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index e9f009d93..82b827437 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -6,7 +6,7 @@ from __future__ import annotations import abc -import itertools +import dataclasses import warnings import os from dataclasses import dataclass, field @@ -41,6 +41,9 @@ "is_nvfp4_available", "get_default_recipe", "get_align_size_for_quantization", + "QuantizerRole", + "QuantizerRequest", + "DelayedScalingRequest", ] @@ -50,6 +53,99 @@ _FP8_BLOCK_SCALING_SUPPORT: Optional[Tuple[bool, str]] = None +@dataclasses.dataclass(frozen=True) +class QuantizerRole: + """Identity of a tensor slot requesting a quantizer. + + TE modules populate all fields they know about. + User factories inspect only the fields they care about. + + .. warning:: + **EXPERIMENTAL**: QuantizerRole is experimental, still under active development, + and the API is subject to change without notice. Use at your own risk. + + Fields + ------ + module_type : str + Module type that emits this role, e.g. `"linear"`, `"grouped_linear"`, `"dpa"`. + Empty string when not provided. + tensor_type : str + What tensor is being quantized, in the module's own vocabulary. + Linear modules: `"input"`, `"weight"`, `"grad_output"`, etc. + DPA: `"qkv"`, `"s"`, etc. + Empty string when not provided. + name : str + Caller-provided module instance name (e.g. set by the training + framework), e.g. + `"qkv"`, `"proj"`, `"fc1"`, `"fc2"`, `"linear_39"`. + Empty string when not provided. + """ + + module_type: str = "" + tensor_type: str = "" + name: str = "" + + def __str__(self) -> str: + parts = [] + if self.module_type: + parts.append(f"module_type={self.module_type}") + if self.tensor_type: + parts.append(f"tensor_type={self.tensor_type}") + if self.name: + parts.append(f"name={self.name}") + return "|".join(parts) if parts else "QuantizerRole()" + + +@dataclasses.dataclass(frozen=True) +class QuantizerRequest: + """Base class for stateful quantizer requests. + + Custom recipe factories return ``QuantizerRequest`` subclasses (instead of + quantizer instances) when the quantizer requires TE-managed shared state. + TE detects these requests, allocates the required state, and replaces them + with real quantizer instances. + + .. warning:: + **EXPERIMENTAL**: QuantizerRequest is experimental, still under active + development, and the API is subject to change without notice. + """ + + +@dataclasses.dataclass(frozen=True) +class DelayedScalingRequest(QuantizerRequest): + """Request a Float8Quantizer with TE-managed delayed scaling state. + + .. warning:: + **EXPERIMENTAL**: DelayedScalingRequest is experimental, still under active + development, and the API is subject to change without notice. + + All ``DelayedScalingRequest`` instances within the same ``CustomRecipeState`` + must share identical parameter values. + + Parameters + ---------- + fp8_format : Format, default = Format.HYBRID + Controls fwd/bwd dtype (HYBRID = E4M3 fwd, E5M2 bwd). + margin : int, default = 0 + Margin for scaling factor computation. + amax_history_len : int, default = 1024 + Length of the amax history window. + amax_compute_algo : str or Callable, default = "max" + Algorithm for choosing amax from history. + scaling_factor_compute_algo : Callable or None, default = None + Custom scaling factor computation. + reduce_amax : bool, default = True + Whether to all-reduce amax across the distributed group. + """ + + fp8_format: Format = Format.HYBRID + margin: int = 0 + amax_history_len: int = 1024 + amax_compute_algo: Union[str, Callable] = "max" + scaling_factor_compute_algo: Optional[Callable] = None + reduce_amax: bool = True + + def _compute_fp8_support() -> Tuple[bool, str]: """Return if fp8 support is available""" if get_device_compute_capability() >= (9, 0): # hopper and above @@ -383,7 +479,7 @@ def add_fp8_tensors_to_global_buffer( fp8_meta: Dict[str, Any], ) -> None: """ - Delayed scaling only. + Delayed scaling only (built-in or custom recipe with DS requests). The amax reduction process happens completely outside the FP8 modules. To participate in the reduction, the only role played by a module is @@ -398,8 +494,8 @@ def add_fp8_tensors_to_global_buffer( wrapper. For non CG case, it's called from within the module. """ - # delayed scaling only function, noop for any other recipe - if not fp8_meta["recipe"].delayed(): + # noop unless delayed scaling state is present + if not _has_delayed_scaling_state(fp8_meta): return # Every module must call this function exactly once since @@ -417,7 +513,17 @@ def add_fp8_tensors_to_global_buffer( # Handles non-parameter FP8 modules, e.g. DPA. continue - key = cls.get_key_in_buffer(forward, fp8_meta["recipe"], fp8_meta["fp8_group"]) + state = fp8_meta[fp8_meta_tensor_key] + + # Determine recipe + buffers: built-in DS or custom with DS requests + if isinstance(state, CustomRecipeState) and state._has_delayed_scaling: + inner_recipe = state._inner_delayed_scaling_recipe + key = cls.get_key_in_buffer(forward, inner_recipe, fp8_meta["fp8_group"]) + # Register inner recipe in autocast_arguments for reduction + autocast_key = cls.get_unique_autocast_key(inner_recipe, fp8_meta["fp8_group"]) + qstate.autocast_arguments[autocast_key] = (inner_recipe, fp8_meta["fp8_group"]) + else: + key = cls.get_key_in_buffer(forward, fp8_meta["recipe"], fp8_meta["fp8_group"]) if key not in qstate.global_amax_buffer: qstate.global_amax_buffer[key] = [fp8_meta[fp8_meta_tensor_key].amax_history[0]] @@ -655,7 +761,7 @@ def copy_forward_fp8_meta_tensors_for_recompute(cls, fp8_meta: Dict[str, Any]) - """ # delayed scaling only function, noop for any other recipe - if not fp8_meta["recipe"].delayed(): + if not _has_delayed_scaling_state(fp8_meta): return buffer_position_key = "global_fp8_buffer_pos_fwd_recompute" @@ -682,7 +788,7 @@ def get_old_fp8_meta_tensors_for_recompute(cls, fp8_meta: Dict[str, Any]) -> Non 1 forward for indentical numerical outputs. """ # delayed scaling only function, noop for any other recipe - if not fp8_meta["recipe"].delayed(): + if not _has_delayed_scaling_state(fp8_meta): return # Store updated amaxes and scales from phase 1 post forward. @@ -703,7 +809,7 @@ def get_old_fp8_meta_tensors_for_recompute(cls, fp8_meta: Dict[str, Any]) -> Non def restore_fp8_meta_tensors(fp8_meta: Dict[str, Any]) -> None: """Restore latest scaling factors and amaxes after recompute forward run.""" # delayed scaling only function, noop for any other recipe - if not fp8_meta["recipe"].delayed(): + if not _has_delayed_scaling_state(fp8_meta): return fp8_meta["scaling_fwd"].amax_history.copy_(fp8_meta["updated_amax_history_fwd"]) @@ -1026,8 +1132,112 @@ class RecipeState(abc.ABC): This class may pack together the state for multiple quantizers, which is helpful for applying fused kernels with less overhead. + Subclasses that own mutable training buffers (e.g. delayed scaling's + ``scale`` / ``amax_history``) MUST list them in + :attr:`_persistent_state_buffers`. These buffers are preserved across + role-driven rebuilds and post-checkpoint resume via + :meth:`inherit_state_from`. Stateless subclasses leave the attribute + empty. """ + roles: Optional[List[QuantizerRole]] + mode: str + + # Names of mutable torch.Tensor attributes that represent persistent + # training state (e.g. running scale, amax history). The default + # ``inherit_state_from`` rebinds these from a predecessor RecipeState + # so external references (e.g. ``FP8GlobalStateManager`` reduction + # buffers) keep pointing at the same backing tensor. + _persistent_state_buffers: Tuple[str, ...] = () + + # Canonical tensor types that a recipe state can dispatch on. + _KNOWN_TENSOR_TYPES = ("input", "weight", "output", "grad_output", "grad_input") + # Positional fallback used when no role information is available: the + # tensor type at slot ``i`` defaults to ``_FWD_DEFAULT_TENSOR_TYPES[i % len]`` + # (forward) or ``_BWD_DEFAULT_TENSOR_TYPES[i % len]`` (backward). Mirrors + # the ``[input, weight, output, ...]`` / ``[grad_output, grad_input, ...]`` + # convention assumed by ``module/base.py::set_meta_tensor``. + _FWD_DEFAULT_TENSOR_TYPES = ("input", "weight", "output") + _BWD_DEFAULT_TENSOR_TYPES = ("grad_output", "grad_input") + + @staticmethod + def _validate_roles( + roles: Optional[List[QuantizerRole]], + num_quantizers: int, + ) -> None: + """Validate that ``roles``, if provided, has length ``num_quantizers``.""" + if roles is not None and len(roles) != num_quantizers: + raise ValueError( + "RecipeState requires roles to match num_quantizers " + f"({len(roles)=} vs {num_quantizers=})" + ) + + def _slot_role(self, idx: int) -> QuantizerRole: + """Resolve slot ``idx`` to a non-``None`` :class:`QuantizerRole`. + + This is the field-agnostic primitive that role-driven recipe states + use to dispatch on any combination of role fields (``tensor_type``, + ``module_type``, ``name``, future fields). + + Resolution rules: + + * If a real ``QuantizerRole`` was provided for this slot, it is + returned unchanged. Producers fill only the fields they know about; + the rest carry the dataclass defaults (empty strings). Consumers + should treat an empty field as "no signal" rather than as "no role + provided". + * Otherwise (whole ``roles`` list missing, or this slot is ``None``), + a bare ``QuantizerRole()`` with all fields empty is returned. + Field-specific fallback policies belong to the individual + dispatch convenience accessors (e.g. :meth:`_slot_tensor_type`), + not to this primitive — that way a future recipe state that + dispatches on, say, ``module_type`` is free to define its own + fallback policy without impacting tensor-type dispatch. + + The "real role vs bare-default role" distinction is hidden from + dispatch logic here. Recipe states that need to *warn* on missing + roles (as :class:`CustomRecipeState` does) should consult + ``self.roles[idx]`` directly. + """ + if self.roles is not None: + role = self.roles[idx] + if role is not None: + return role + return QuantizerRole() + + def _slot_tensor_type(self, idx: int) -> str: + """Convenience accessor: tensor-type dispatch with positional fallback. + + Resolves to one of :attr:`_KNOWN_TENSOR_TYPES`. Used by recipe states + whose dispatch only depends on the tensor's role within a GEMM + (input / weight / output / grad_output / grad_input), e.g. + Float8BlockScalingRecipeState, NVFP4BlockScalingRecipeState. + + Behavior: + + * If the resolved :meth:`_slot_role` carries a ``tensor_type`` in + :attr:`_KNOWN_TENSOR_TYPES`, return it. + * Otherwise (no role provided, a role with empty / non-canonical + ``tensor_type`` like DPA's ``"qkv"``, or a role that intentionally + only sets ``module_type``/``name``), fall back to the positional + default (forward: ``[input, weight, output, ...]``; + backward: ``[grad_output, grad_input, ...]``) indexed by + ``idx % len(default_tensor_types)``. + + This fallback policy is local to tensor-type dispatch; it does not + affect :meth:`_slot_role` or any other accessor. + """ + role = self._slot_role(idx) + if role.tensor_type in self._KNOWN_TENSOR_TYPES: + return role.tensor_type + # Positional fallback: tensor_type is missing or non-canonical. + default_tensor_types = ( + self._FWD_DEFAULT_TENSOR_TYPES + if self.mode == "forward" + else self._BWD_DEFAULT_TENSOR_TYPES + ) + return default_tensor_types[idx % len(default_tensor_types)] + @staticmethod def create( recipe: Recipe, @@ -1035,6 +1245,7 @@ def create( mode: str, num_quantizers: int = 1, device: Optional[torch.device] = None, + roles: Optional[List[QuantizerRole]] = None, ) -> RecipeState: """Factory method to create the state for a quantization recipe @@ -1048,6 +1259,9 @@ def create( Number of quantizers to create state for. device: torch.device, default = default CUDA device Device for quantized tensors. + roles: list of QuantizerRole, optional + Semantic roles for each quantizer slot. When provided, must + have length ``num_quantizers``. Returns ------- @@ -1076,6 +1290,7 @@ def create( mode=mode, num_quantizers=num_quantizers, device=device, + roles=roles, ) @abc.abstractmethod @@ -1088,6 +1303,43 @@ def make_quantizers(self) -> list: """ + def inherit_state_from(self, other: "RecipeState") -> bool: + """Take over persistent training buffers from a predecessor state. + + Used when a ``RecipeState`` is being replaced (e.g. role-driven + rebuild, post-checkpoint resume) but its mutable buffers must + survive. The default implementation rebinds attributes listed in + :attr:`_persistent_state_buffers` to ``other``'s tensor objects. + Rebinding (rather than copying values) ensures any external + references — most importantly the + :class:`FP8GlobalStateManager` reduction buffers — keep pointing + at storage that is also visible to this state's quantizers, so + amax reductions and quantization stay consistent. + + Subclasses with composed sub-states (e.g. :class:`CustomRecipeState` + owning an inner :class:`DelayedScalingRecipeState`) override this + to recurse / stash for later use during ``make_quantizers``. + + Returns + ------- + bool + ``True`` if any persistent buffer was inherited; ``False`` if + the states are incompatible (different class, mismatched + shapes / dtypes) and a fresh state should be used instead. + """ + if type(self) is not type(other): + return False + if not self._persistent_state_buffers: + return False + for name in self._persistent_state_buffers: + src = getattr(other, name) + dst = getattr(self, name) + if src.shape != dst.shape or src.dtype != dst.dtype: + return False + for name in self._persistent_state_buffers: + setattr(self, name, getattr(other, name)) + return True + class DelayedScalingRecipeState(RecipeState): """State for FP8 quantization with per-tensor delayed scaling. @@ -1105,6 +1357,10 @@ class DelayedScalingRecipeState(RecipeState): scale: torch.Tensor amax_history: torch.Tensor + # Persistent training state inherited across role-driven rebuilds. + # See ``RecipeState.inherit_state_from``. + _persistent_state_buffers = ("scale", "amax_history") + def __init__( self, recipe: DelayedScaling, @@ -1112,10 +1368,13 @@ def __init__( mode: str, num_quantizers: int = 1, device: Optional[torch.device] = None, + roles: Optional[List[QuantizerRole]] = None, ) -> None: + self._validate_roles(roles, num_quantizers) self.recipe = recipe self.mode = mode self.num_quantizers = num_quantizers + self.roles = roles self.dtype = get_fp8_te_dtype(recipe, mode == "forward") # Allocate buffers @@ -1158,10 +1417,13 @@ def __init__( mode: str, num_quantizers: int = 1, device: Optional[torch.device] = None, + roles: Optional[List[QuantizerRole]] = None, ) -> None: + self._validate_roles(roles, num_quantizers) self.recipe = recipe self.mode = mode self.num_quantizers = num_quantizers + self.roles = roles self.dtype = get_fp8_te_dtype(recipe, mode == "forward") # Allocate buffers @@ -1198,10 +1460,13 @@ def __init__( mode: str, num_quantizers: int = 1, device: Optional[torch.device] = None, + roles: Optional[List[QuantizerRole]] = None, ) -> None: + self._validate_roles(roles, num_quantizers) self.recipe = recipe self.mode = mode self.num_quantizers = num_quantizers + self.roles = roles self.dtype = get_fp8_te_dtype(recipe, mode == "forward") # Allocate buffers @@ -1235,10 +1500,13 @@ def __init__( mode: str, num_quantizers: int = 1, device: Optional[torch.device] = None, + roles: Optional[List[QuantizerRole]] = None, ) -> None: + self._validate_roles(roles, num_quantizers) self.recipe = recipe self.mode = mode self.num_quantizers = num_quantizers + self.roles = roles self.qx_dtype = get_fp8_te_dtype(recipe, True) self.qw_dtype = get_fp8_te_dtype(recipe, True) self.qgrad_dtype = get_fp8_te_dtype(recipe, False) @@ -1249,75 +1517,54 @@ def __init__( self.device = device def make_quantizers(self) -> list: + """Build one ``Float8BlockQuantizer`` per slot, dispatched by tensor type. + + Per-slot behavior, resolved via :meth:`RecipeState._slot_tensor_type`: + + * ``"weight"`` uses ``recipe.fp8_quant_fwd_weight`` and + ``recipe.w_block_scaling_dim``. + * ``"input"`` / ``"output"`` (and any unknown forward slot) use + ``recipe.fp8_quant_fwd_inp`` and ``recipe.x_block_scaling_dim``. + * ``"grad_output"`` / ``"grad_input"`` (and any unknown backward slot) + use ``recipe.fp8_quant_bwd_grad`` and ``recipe.grad_block_scaling_dim``. + + When the owning module/op provides a role list via + ``get_quantizer_roles``, the per-slot ``tensor_type`` drives dispatch. + Otherwise (or for boundary slots whose role is ``None``), the + positional fallback ``[input, weight, output, ...]`` / + ``[grad_output, grad_input, ...]`` is used. This matches the legacy + index-based convention, so behavior is unchanged for + modules that haven't adopted roles yet. + """ # TODO(ksivamani); Find better design for this, adding here to avoid circular import. from .tensor.float8_blockwise_tensor import Float8BlockQuantizer - if self.mode == "forward": - # The index convention (coming from base.py set_meta_tensor) - # is somewhat awkward, and doesn't play nicely with QuantizeOp, - # which is not associated with a GEMM. - assert self.num_quantizers % 3 == 0 # x, w, output per gemm - return list( - itertools.chain.from_iterable( - [ - [ - Float8BlockQuantizer( - fp8_dtype=self.qx_dtype, - rowwise=True, - columnwise=True, - amax_epsilon=self.recipe.fp8_quant_fwd_inp.amax_epsilon, - force_pow_2_scales=self.recipe.fp8_quant_fwd_inp.power_2_scale, - block_scaling_dim=self.recipe.x_block_scaling_dim, - ), - Float8BlockQuantizer( - fp8_dtype=self.qw_dtype, - rowwise=True, - columnwise=True, - amax_epsilon=self.recipe.fp8_quant_fwd_weight.amax_epsilon, - force_pow_2_scales=self.recipe.fp8_quant_fwd_weight.power_2_scale, - block_scaling_dim=self.recipe.w_block_scaling_dim, - ), - Float8BlockQuantizer( - fp8_dtype=self.qx_dtype, - rowwise=True, - columnwise=True, - amax_epsilon=self.recipe.fp8_quant_fwd_inp.amax_epsilon, - force_pow_2_scales=self.recipe.fp8_quant_fwd_inp.power_2_scale, - block_scaling_dim=self.recipe.x_block_scaling_dim, - ), - ] - for _ in range(self.num_quantizers // 3) - ] - ) + def _make(tensor_type: str) -> Float8BlockQuantizer: + if tensor_type == "weight": + qparams = self.recipe.fp8_quant_fwd_weight + fp8_dtype = self.qw_dtype + block_scaling_dim = self.recipe.w_block_scaling_dim + elif tensor_type in ("grad_output", "grad_input"): + qparams = self.recipe.fp8_quant_bwd_grad + fp8_dtype = self.qgrad_dtype + block_scaling_dim = self.recipe.grad_block_scaling_dim + else: + # "input", "output", or any unknown forward type fall back to + # the input config, matching the legacy positional behavior. + qparams = self.recipe.fp8_quant_fwd_inp + fp8_dtype = self.qx_dtype + block_scaling_dim = self.recipe.x_block_scaling_dim + return Float8BlockQuantizer( + fp8_dtype=fp8_dtype, + rowwise=True, + columnwise=True, + amax_epsilon=qparams.amax_epsilon, + force_pow_2_scales=qparams.power_2_scale, + block_scaling_dim=block_scaling_dim, ) - assert self.mode == "backward", f"Unexpected mode {self.mode}" - assert self.num_quantizers % 2 == 0 # grad_output and grad_input per gemm - return list( - itertools.chain.from_iterable( - [ - [ - Float8BlockQuantizer( - fp8_dtype=self.qgrad_dtype, - rowwise=True, - columnwise=True, - amax_epsilon=self.recipe.fp8_quant_bwd_grad.amax_epsilon, - force_pow_2_scales=self.recipe.fp8_quant_bwd_grad.power_2_scale, - block_scaling_dim=self.recipe.grad_block_scaling_dim, - ), - Float8BlockQuantizer( - fp8_dtype=self.qgrad_dtype, - rowwise=True, - columnwise=True, - amax_epsilon=self.recipe.fp8_quant_bwd_grad.amax_epsilon, - force_pow_2_scales=self.recipe.fp8_quant_bwd_grad.power_2_scale, - block_scaling_dim=self.recipe.grad_block_scaling_dim, - ), - ] - for _ in range(self.num_quantizers // 2) - ] - ) - ) + assert self.mode in ("forward", "backward"), f"Unexpected mode {self.mode}" + return [_make(self._slot_tensor_type(idx)) for idx in range(self.num_quantizers)] class NVFP4BlockScalingRecipeState(RecipeState): @@ -1338,10 +1585,13 @@ def __init__( mode: str, num_quantizers: int = 1, device: Optional[torch.device] = None, + roles: Optional[List[QuantizerRole]] = None, ) -> None: + self._validate_roles(roles, num_quantizers) self.recipe = recipe self.mode = mode self.num_quantizers = num_quantizers + self.roles = roles self.dtype = get_fp4_te_dtype(recipe) # Allocate buffers @@ -1349,63 +1599,184 @@ def __init__( device = torch.device("cuda") def make_quantizers(self) -> list: + """Build one ``NVFP4Quantizer`` per slot, dispatched by tensor type. + + Per-slot behavior, resolved via :meth:`RecipeState._slot_tensor_type`: + + * Forward, ``"weight"`` -> ``recipe.fp4_quant_fwd_weight``. + * Forward, ``"input"`` / ``"output"`` (and any unknown forward type) -> + ``recipe.fp4_quant_fwd_inp``. + * Backward, any slot -> ``recipe.fp4_quant_bwd_grad``. + + When the owning module/op provides a role list via + ``get_quantizer_roles``, the per-slot ``tensor_type`` drives dispatch. + Otherwise (or for boundary slots whose role is ``None``), the + positional fallback ``[input, weight, output, ...]`` is used; on this + layout slot ``idx % 3 == 1`` is always weight and the rest fall into + the input config, matching the legacy index-based behavior. + """ from .tensor.nvfp4_tensor import NVFP4Quantizer - # The index convention (coming from base.py set_meta_tensor) - # is somewhat awkward. It assumes forward quantizers are - # ordered [input, weight, output, ...] and backward quantizers - # are ordered [grad_output, grad_input, ...]. This doesn't - # play nicely with fusible ops: Linear op doesn't own output - # or grad input quantizers, Quantize op only owns input and - # grad output quantizers. - - if self.mode == "forward": - - def _make_quantizer(idx: int) -> NVFP4Quantizer: - qparams = ( - self.recipe.fp4_quant_fwd_weight - if idx % 3 == 1 - else self.recipe.fp4_quant_fwd_inp - ) - return NVFP4Quantizer( - fp4_dtype=self.dtype, - rowwise=True, - columnwise=True, - with_rht=qparams.random_hadamard_transform, - with_post_rht_amax=qparams.random_hadamard_transform, - with_2d_quantization=qparams.fp4_2d_quantization, - stochastic_rounding=qparams.stochastic_rounding, - row_scaled_nvfp4=self.recipe.row_scaled_activation and idx % 3 != 1, - ) + def _qparams(tensor_type: str): + if self.mode == "backward": + return self.recipe.fp4_quant_bwd_grad + if tensor_type == "weight": + return self.recipe.fp4_quant_fwd_weight + return self.recipe.fp4_quant_fwd_inp + + def _make(tensor_type: str) -> NVFP4Quantizer: + qparams = _qparams(tensor_type) + return NVFP4Quantizer( + fp4_dtype=self.dtype, + rowwise=True, + columnwise=True, + with_rht=qparams.random_hadamard_transform, + with_post_rht_amax=qparams.random_hadamard_transform, + with_2d_quantization=qparams.fp4_2d_quantization, + stochastic_rounding=qparams.stochastic_rounding, + row_scaled_nvfp4=( + self.mode == "forward" + and tensor_type != "weight" + and self.recipe.row_scaled_activation + ), + ) - return [_make_quantizer(idx) for idx in range(self.num_quantizers)] - - if self.mode == "backward": - return [ - NVFP4Quantizer( - fp4_dtype=self.dtype, - rowwise=True, - columnwise=True, - with_rht=self.recipe.fp4_quant_bwd_grad.random_hadamard_transform, - with_post_rht_amax=self.recipe.fp4_quant_bwd_grad.random_hadamard_transform, - with_2d_quantization=self.recipe.fp4_quant_bwd_grad.fp4_2d_quantization, - stochastic_rounding=self.recipe.fp4_quant_bwd_grad.stochastic_rounding, - row_scaled_nvfp4=False, + if self.mode not in ("forward", "backward"): + raise RuntimeError(f"Unexpected recipe mode ({self.mode})") + + return [_make(self._slot_tensor_type(idx)) for idx in range(self.num_quantizers)] + + +def _handle_delayed_scaling_requests( + raw: list, + device: torch.device, + mode: str, + *, + existing_ds_state: Optional["DelayedScalingRecipeState"] = None, +) -> Optional["DelayedScalingRecipeState"]: + """Detect DelayedScalingRequest items, allocate shared state, replace with real quantizers. + + All DS requests in the same RecipeState must share identical parameters. + + When ``existing_ds_state`` is provided and compatible (same dtype, + same number of DS slots, same ``amax_history_len``), it is reused + instead of allocating fresh buffers. Reusing preserves accumulated + ``scale`` / ``amax_history`` across role-driven rebuilds — important + for post-checkpoint resume and mid-training factory swaps. The + ``Float8Quantizer`` instances built here will then view into the + SAME tensor objects already registered with + ``FP8GlobalStateManager``'s reduction buffers, keeping reduction + and quantization consistent. + + Returns a ``DelayedScalingRecipeState`` owning the shared buffers, or + ``None`` when no DS requests are present. + """ + ds_items = [(i, r) for i, r in enumerate(raw) if isinstance(r, DelayedScalingRequest)] + if not ds_items: + return None + + r0 = ds_items[0][1] + + # Validate all DS requests share same params + for idx, req in ds_items[1:]: + for field_name in ( + "fp8_format", + "margin", + "amax_history_len", + "amax_compute_algo", + "scaling_factor_compute_algo", + "reduce_amax", + ): + v0 = getattr(r0, field_name) + vi = getattr(req, field_name) + if v0 != vi: + raise ValueError( + "All DelayedScalingRequests in one CustomRecipeState must match. " + f"Slot 0 has {field_name}={v0!r}, slot {idx} has {vi!r}." ) - for _ in range(self.num_quantizers) - ] - raise RuntimeError(f"Unexpected recipe mode ({self.mode})") + # Build a real DelayedScalingRecipeState to own the shared buffers. + inner_recipe = DelayedScaling( + fp8_format=r0.fp8_format, + margin=r0.margin, + amax_history_len=r0.amax_history_len, + amax_compute_algo=r0.amax_compute_algo, + scaling_factor_compute_algo=r0.scaling_factor_compute_algo, + reduce_amax=r0.reduce_amax, + ) + n = len(ds_items) + + # Reuse a compatible existing DSRS so its scale / amax_history (and any + # external references to them) survive the rebuild. + expected_dtype = get_fp8_te_dtype(inner_recipe, mode == "forward") + dsrs = None + if existing_ds_state is not None: + if ( + existing_ds_state.num_quantizers == n + and existing_ds_state.dtype == expected_dtype + and existing_ds_state.amax_history.shape[0] == r0.amax_history_len + ): + dsrs = existing_ds_state + + if dsrs is None: + dsrs = DelayedScalingRecipeState( + inner_recipe, + mode=mode, + num_quantizers=n, + device=device, + ) + + # Splice Float8Quantizer instances (backed by dsrs buffers) into raw list. + quantizers = dsrs.make_quantizers() + for j, (idx, _req) in enumerate(ds_items): + raw[idx] = quantizers[j] + + return dsrs + + +def _has_delayed_scaling_state(fp8_meta: Dict[str, Any]) -> bool: + """Check if fp8_meta has delayed scaling state (built-in or custom).""" + if fp8_meta["recipe"].delayed(): + return True + if fp8_meta["recipe"].custom(): + for key in ("scaling_fwd", "scaling_bwd"): + state = fp8_meta.get(key) + if isinstance(state, CustomRecipeState) and state._has_delayed_scaling: + return True + return False class CustomRecipeState(RecipeState): - """State for CustomRecipe: produce quantizers per tensor.""" + """State for CustomRecipe: produce quantizers per tensor. + + Stateful quantizer support: + - Supports stateful quantizers (e.g. delayed scaling) via ``DelayedScalingRequest``. + - The factory returns request dataclasses for stateful quantizers; TE detects them, + allocates shared buffers, and replaces with real quantizer instances. + - Stateful recipe state is composed via real TE recipe state objects (e.g. + ``DelayedScalingRecipeState``), not reimplemented. + """ recipe: CustomRecipe mode: str num_quantizers: int device: Optional[torch.device] + # -- Composed sub-states for stateful sub-recipes -- + # + # When the qfactory returns request objects (e.g. ``DelayedScalingRequest``) + # for a stateful built-in recipe, ``make_quantizers`` allocates a real + # built-in ``RecipeState`` for those slots and reuses its persistent + # buffers across role-driven rebuilds via ``inherit_state_from``. One + # ``__state`` / ``__state_to_inherit`` pair per stateful recipe. + + # Delayed scaling (``DelayedScalingRequest`` -> ``DelayedScalingRecipeState``): + # ``_ds_state`` owns shared ``scale`` / ``amax_history`` for DS slots in this + # CustomRecipeState; ``_ds_state_to_inherit`` is a transient stash set by + # ``inherit_state_from`` and consumed by the next ``make_quantizers`` call. + _ds_state: Optional[DelayedScalingRecipeState] + _ds_state_to_inherit: Optional[DelayedScalingRecipeState] + def __init__( self, recipe: CustomRecipe, @@ -1413,39 +1784,106 @@ def __init__( mode: str, num_quantizers: int = 1, device: Optional[torch.device] = None, + roles: Optional[List[QuantizerRole]] = None, ) -> None: + self._validate_roles(roles, num_quantizers) self.recipe = recipe self.mode = mode self.num_quantizers = num_quantizers + self.roles = roles if device is None: device = torch.device("cuda") self.device = device + # -- Stateful sub-state slots (initialized empty) -- + # Delayed scaling + self._ds_state = None + self._ds_state_to_inherit = None + if getattr(recipe, "qfactory", None) is None: raise ValueError("CustomRecipe requires `qfactory`.") def make_quantizers(self) -> list: qfactory = self.recipe.qfactory - out = [] - - # TODO(negvet): make_quantizers() should take roles from the operation - # Hardcode linear-specific roles for now - roles: List[str] - if self.mode == "forward": - roles = [ - ("linear_input", "linear_weight", "linear_output")[i % 3] - for i in range(self.num_quantizers) - ] - elif self.mode == "backward": - roles = [ - ("linear_grad_output", "linear_grad_input")[i % 2] - for i in range(self.num_quantizers) - ] - else: - roles = ["unknown"] * self.num_quantizers - for i in range(self.num_quantizers): - # Get quantizer from the user defined factory - quantizer = qfactory(roles[i]) - out.append(quantizer) - return out + roles = self.roles + if roles is None: + warnings.warn( + "CustomRecipeState: no QuantizerRole list provided by the module/op. " + "Falling back to bare QuantizerRole() defaults. " + "Override get_quantizer_roles() to provide meaningful roles.", + stacklevel=2, + ) + roles = [QuantizerRole() for _ in range(self.num_quantizers)] + + # qfactory must return a Quantizer or QuantizerRequest for every slot. + # None is not a valid return value — it would silently disable quantization + # for that tensor, risking hard-to-detect performance regressions. + # TODO(negvet): Introduce an explicit IdentityQuantizer for intentional no-op + # quantization. Until then, None is rejected. + raw = [qfactory(roles[i]) for i in range(self.num_quantizers)] + for i, q in enumerate(raw): + if q is None: + raise ValueError( + f"CustomRecipe qfactory returned None for slot {i} " + f"(role={roles[i]}). Every slot must return a Quantizer " + "instance or a QuantizerRequest." + ) + + # -- Delayed scaling sub-state -- + # If a predecessor stashed a compatible inner DSRS via + # ``inherit_state_from``, reuse it so accumulated scale / amax_history + # survive the rebuild. Consume the stash so a subsequent + # ``make_quantizers`` doesn't reuse it again unintentionally. + existing_ds_state = self._ds_state_to_inherit + self._ds_state_to_inherit = None + self._ds_state = _handle_delayed_scaling_requests( + raw, + self.device, + self.mode, + existing_ds_state=existing_ds_state, + ) + + return raw + + def inherit_state_from(self, other: "RecipeState") -> bool: + """Stash ``other``'s composed sub-states for reuse on next ``make_quantizers``. + + ``CustomRecipeState`` cannot inherit declaratively because its + persistent state lives in composed sub-states (one per stateful + sub-recipe) that are allocated only when ``make_quantizers`` runs. + For each stateful sub-recipe we stash the predecessor's sub-state + and let the next ``make_quantizers`` decide whether the + predecessor's shape is compatible with the new factory output. + """ + if not isinstance(other, CustomRecipeState): + return False + + inherited_any = False + + # -- Delayed scaling sub-state -- + if other._ds_state is not None: + self._ds_state_to_inherit = other._ds_state + inherited_any = True + + return inherited_any + + # -- Delegation to composed DelayedScalingRecipeState -- + + @property + def _has_delayed_scaling(self) -> bool: + return self._ds_state is not None + + @property + def amax_history(self) -> Optional[torch.Tensor]: + """Amax history from the composed delayed-scaling state, if any.""" + return self._ds_state.amax_history if self._ds_state else None + + @property + def scale(self) -> Optional[torch.Tensor]: + """Current scale from the composed delayed-scaling state, if any.""" + return self._ds_state.scale if self._ds_state else None + + @property + def _inner_delayed_scaling_recipe(self) -> Optional[DelayedScaling]: + return self._ds_state.recipe if self._ds_state else None From b7323b1ecf92a94a0b15dab423f15549f19a4d39 Mon Sep 17 00:00:00 2001 From: Xi Jingyi <48742253+jing-4369@users.noreply.github.com> Date: Tue, 12 May 2026 04:57:15 +0800 Subject: [PATCH 048/180] [Common][PyTorch] Fix int32 overflow and -1 sentinel handling in moe_permute (#2907) * [Common][PyTorch] Fix int32 overflow and -1 sentinel handling in moe_permute Two independent bugs in transformer_engine/common/permutation/permutation.cu and the PyTorch extension caller reproduce on main (264da2b) and v2.13: 1. int32 overflow in moe_unpermute_kernel and moe_permute_kernel. `source_token * num_cols` and `source_row * num_cols` are computed with int, so for long-sequence MoE workloads where num_out_tokens * num_cols reaches 2**31 (e.g. 2**18 tokens x 2**13 hidden), the pointer offset wraps and the kernel either reads garbage or raises `an illegal memory access was encountered`. Widening source_token, source_row and dest_row to int64_t inside the kernels keeps the index arithmetic in 64 bits without changing any public types. 2. Incorrect handling of -1 sentinels in the routing indices. Libraries such as DeepEP (and any expert-parallel mask that sets non-local (token, slot) pairs to -1) feed a routing_map that contains -1 entries. `cub::DeviceRadixSort::SortPairs` is signed ascending, so those sentinels land at the HEAD of sorted_row_id, not the tail. moe_permute_row_map currently writes -1 only for idx >= num_out_tokens and reads the sentinel prefix as if it were a valid sorted id, producing bogus row_id_map writes (for instance `source_row / topK == 0, source_row % topK == -1`). The caller now advances sorted_row_id_ptr past the num_minus_ones prefix and pre-fills row_id_map with -1 via torch::full, so the kernel only processes the valid suffix and never dereferences a sentinel. The launcher's grid switches from num_rows*topK blocks to num_out_tokens blocks to match the new valid range. No behaviour change on happy-path routing_map (no -1, no overflow). Reproducers: - 8-token, topK=2 routing_map with -1 masking: max |TE - ref| = 4.5e0 on bf16 with current main; 0.0 with this patch. - num_tokens=2**18+1, num_cols=2**13, topK=1: current main raises CUDA illegal memory access at permutation.cu:252; with this patch it succeeds. Signed-off-by: Jingyi Xi * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Guard against invalid num_out_tokens in moe_permute_fwd Add an NVTE_CHECK that num_out_tokens <= num_tokens * topK and cast num_minus_ones to size_t before the pointer advance, so a negative num_minus_ones (from an invalid num_out_tokens) cannot silently wrap into a huge pointer offset. Signed-off-by: Jingyi Xi * Switch radix sort keys to uint32_t to fix -1 sentinel ordering The MoE permute path was correct for the existing capacity-drop convention (drops encoded as a large positive expert id, sorted to the tail by the signed cub::DeviceRadixSort), but it broke for callers that mark dropped (token, slot) pairs with -1 (expert-parallel rank masking, e.g. DeepEP). With signed sort the -1 sentinels land at the HEAD of sorted_row_id, while moe_permute_row_map's `idx >= num_out_tokens` branch assumes drops are at the tail. Reinterpret the keys as uint32_t inside nvte_device_radix_sort_pairs so -1 (= UINT_MAX) sorts to the tail, unifying the EP-mask case with the existing capacity-drop convention. The kernel and caller sides are unchanged - this is a one-place fix that makes both drop conventions land in the existing drop branch. Also widen the loop-carried indices in moe_unpermute_kernel and moe_permute_kernel to int64_t (`source_token`, `source_row`, `dest_row`) to keep `row * num_cols` in 64 bits. We hit this on DeepSeek-V3 long- context training (hidden = 7168, topK = 8): once `num_out_tokens * num_cols` reaches 2**31 the int product wraps and the kernel either silently corrupts rows or raises CUDA `illegal memory access`. Signed-off-by: Jingyi Xi * Widen num_rows * topK products in moe_permute_row_map for consistency Per reviewer feedback in NVIDIA/TransformerEngine#2907, promote the int * int multiplications in moe_permute_row_map and its launcher to int64_t. These are not the overflow path this PR was originally fixing (DeepSeek-V3 long-context hits row * num_cols, where num_cols is the hidden dim ~ 7-8k), and num_rows * topK only crosses 2**31 at unrealistic per-rank token counts (>= 268M at topK=8). The change is purely defensive but keeps the index arithmetic in this kernel consistent with the int64_t source_token / source_row / dest_row widening already applied to moe_unpermute_kernel and moe_permute_kernel. Signed-off-by: Jingyi Xi --------- Signed-off-by: Jingyi Xi Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Teddy Do --- .../common/permutation/permutation.cu | 27 ++++++++++++------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/transformer_engine/common/permutation/permutation.cu b/transformer_engine/common/permutation/permutation.cu index fbba27941..aa7cb50e8 100644 --- a/transformer_engine/common/permutation/permutation.cu +++ b/transformer_engine/common/permutation/permutation.cu @@ -19,7 +19,7 @@ static __global__ void moe_permute_row_map(const int *sorted_row_id, int *row_id const int tid = threadIdx.x; const int idx = bid * blockDim.x + tid; - if (idx >= num_rows * topK) return; + if (idx >= static_cast(num_rows) * topK) return; int source_row = sorted_row_id[idx]; int source_token_id = source_row / topK; @@ -27,10 +27,10 @@ static __global__ void moe_permute_row_map(const int *sorted_row_id, int *row_id if (idx >= num_out_tokens) { // Set the indices of dropped tokens to -1 - row_id_map[source_topK_id * num_rows + source_token_id] = -1; + row_id_map[static_cast(source_topK_id) * num_rows + source_token_id] = -1; } else { // Create a row id map for subsequent unpermute operation - row_id_map[source_topK_id * num_rows + source_token_id] = idx; + row_id_map[static_cast(source_topK_id) * num_rows + source_token_id] = idx; } } @@ -42,7 +42,7 @@ __global__ void moe_unpermute_kernel(const T *input, T *unpermuted_output, const TCompute *s_prob = reinterpret_cast(s_mem); // Each block corresponds to one dest token - const int source_token = blockIdx.x; + const int64_t source_token = blockIdx.x; const int tid = threadIdx.x; if (hasProb) { @@ -65,7 +65,7 @@ __global__ void moe_unpermute_kernel(const T *input, T *unpermuted_output, const TCompute frag_elem[kElementsPerAccess]; TCompute frag_sum[kElementsPerAccess]; - int source_row = row_id_map[source_token]; + int64_t source_row = row_id_map[source_token]; // source_row == -1 represents a dropped token if (source_row != -1) { @@ -134,7 +134,7 @@ __global__ void moe_permute_kernel(const T *input_bwd, const T *input_fwd, T *ac TCompute *s_prob = reinterpret_cast(s_mem); // Each block corresponds to one source token - const int source_token = blockIdx.x; + const int64_t source_token = blockIdx.x; const int tid = threadIdx.x; if (hasProb) { @@ -172,7 +172,7 @@ __global__ void moe_permute_kernel(const T *input_bwd, const T *input_fwd, T *ac for (int k = 0; k < topKTile; k++) { if (k == topK) break; - int dest_row = row_id_map[index]; + int64_t dest_row = row_id_map[index]; index += num_rows; if (dest_row != -1) { @@ -239,7 +239,7 @@ void nvte_permute_launcher(const T *input, T *output, const int *sorted_row_id, // moe_permute_fwd int threads = 64; - int blocks = (num_rows * topK + threads - 1) / threads; + int blocks = (static_cast(num_rows) * topK + threads - 1) / threads; moe_permute_row_map<<>>(sorted_row_id, row_id_map, num_rows, topK, num_out_tokens); @@ -371,6 +371,13 @@ void nvte_device_radix_sort_pairs(void *temp_storage, size_t *temp_storage_bytes int *keys_out, int *values_in, int *values_out, size_t num_items) { NVTE_API_CALL(nvte_device_radix_sort_pairs); - cub::DeviceRadixSort::SortPairs(temp_storage, *temp_storage_bytes, keys_in, keys_out, values_in, - values_out, num_items); + // Sort keys as uint32_t so any negative-int sentinel (e.g. `-1` placed by an + // expert-parallel rank mask) becomes a large unsigned value and lands at the + // tail of the sorted output, matching the existing capacity-drop convention + // (drops encoded as a large positive expert id) and the + // `idx >= num_out_tokens` drop branch in moe_permute_row_map. + auto *u_keys_in = reinterpret_cast(keys_in); + auto *u_keys_out = reinterpret_cast(keys_out); + cub::DeviceRadixSort::SortPairs(temp_storage, *temp_storage_bytes, u_keys_in, u_keys_out, + values_in, values_out, num_items); } From 282b4fb29b5f4658b617af37b01c97785673ecb1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Mon, 11 May 2026 23:00:24 +0200 Subject: [PATCH 049/180] [torch.compile][PyTorch] Prepare linear for torch compile (#2967) * [PyTorch] Linear: minor cleanups for compile-friendliness Three small refactors that make the module easier to reason about and pave the way for the dataclass / saved-tensor refactors: - Add a TensorOrQuantized type alias (Union[Tensor, QuantizedTensorStorage]) used pervasively in helper signatures. - Hoist the conditional bias argument into a local linear_bias_tensor variable instead of an inline expression at the linear_fn() call site. - Only forward self.wgrad_store into the autograd Function when it is actually active (delay_wgrad_compute() is True); pass None otherwise so the autograd graph does not carry an unused Python object. Pure rename / hoisting; no behavioural change. Signed-off-by: Pawel Gadzinski Co-authored-by: Cursor * [PyTorch] Linear: pack forward/backward state into dataclasses Replace the loosely typed ``non_tensor_args`` tuple and the ad-hoc ``ctx.`` plumbing with two dataclasses, ``LinearFwdArgs`` and ``LinearBwdArgs``, that act as the single argument to every helper in the forward/backward pipeline. What changes: * ``LinearFwdArgs`` carries the (positional) tensors ``weight``, ``inp`` and ``bias`` plus all quantizers, ``requires_grad`` flags, the cached ``weight_workspace`` and every former ``non_tensor_args`` knob. ``_Linear.forward`` still takes ``weight/inp/bias`` as positional Tensor inputs so autograd tracks them, then immediately re-attaches them to ``fwd_args`` so every downstream helper has a single-argument signature. * ``LinearBwdArgs`` mirrors that on the backward side: it owns the saved tensors (``inputmat``, ``weight_fp8``, ``saved_weight``, ``bias``), the per-call quantizers, every flag previously stored directly on ``ctx`` and a ``setup_saved_tensors(saved_tensors, tensor_objects)`` helper that rehydrates the saved-tensor fields. * ``ctx.backward_objects = bwd_args`` is now the single attribute the autograd context needs (besides ``saved_tensors``/``tensor_objects``). * ``weight_workspace`` is no longer a positional Tensor arg of the autograd Function; it is read from ``fwd_args.weight_workspace`` and the freshly produced workspace is returned alongside ``out`` so the module can refresh its cache without autograd tracking the cache. * ``prepare_for_saving`` now lives at the autograd boundary in ``_Linear.forward``; ``_linear_setup_ctx`` only returns the merged list of tensors that should be saved. * ``grad_output_preprocess`` is invoked with ``bwd_args`` directly (it is duck-typed on the same attribute names) so backward never reaches into ``ctx.`` for non-tensor state. Behaviour preserved (verified numerically against ``torch.nn.Linear`` and on FP8 + workspace-cache paths). Signed-off-by: Pawel Gadzinski Co-authored-by: Cursor * [PyTorch] Linear: deduplicate saved tensors that alias forward inputs When ``saved_inputmat is inp``, ``wt_save is weight`` or ``bias`` is the exact bias passed in, there is no point asking ``prepare_for_saving`` to serialize the same Python object twice. Make ``_linear_forward_impl`` emit ``None`` in those slots (and a parallel ``saved_tensor_aliases`` tuple in ``ctx_attrs`` describing which slot points where), and have ``_linear_setup_ctx`` rebuild the tuple with the original references before handing it to ``prepare_for_saving``. Saves a Python ref per alias in eager and, more importantly, keeps the forward helper from "returning" a tensor that aliases its own inputs -- a pattern ``torch.compile`` would otherwise need to reason about when the helper is wrapped in an opaque op. Numerically equivalent (validated against ``torch.nn.Linear`` and on a multi-iteration FP8 path with workspace caching). Signed-off-by: Pawel Gadzinski Co-authored-by: Cursor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Linear: tighten LinearFwdArgs/BwdArgs and trim ctx_attrs Follow-up cleanups on top of the dataclass refactor: * Sort ``LinearFwdArgs`` / ``LinearBwdArgs`` fields into labelled groups (tensors, requires_grad flags, quantizers, dtype/numerical config, parallelism, userbuffers, FSDP, wgrad scheduling, misc) and mirror that ordering in their construction sites. * Add ``slots=True`` to both dataclasses so typos in ``fwd_args.X`` / ``bwd_args.X`` raise ``AttributeError`` immediately instead of silently creating a new attribute. * Inline single-use ``args.X`` aliases in ``_linear_forward_impl`` (``weight_workspace``, ``fp8_calibration``, ``tp_size``, ``tensor_parallel``, ``cache_weight``, ``skip_fp8_weight_update``, ``custom``, ``backward_input_needs_gather``) so the prelude only keeps aliases that are actually reused. * Shrink ``ctx_attrs`` to ``{fsdp_shapes, saved_tensor_aliases}``: ``weight_quantizer`` is re-derived in ``_linear_setup_ctx`` from ``fwd_args.weight`` (matching the resolution done in forward), ``is_fsdp2`` already lives on ``fwd_args``, and ``owns_input`` is equivalent to ``saved_tensor_aliases[0] != "inp"``. * Replace ``setup_saved_tensors(saved_tensors, tensor_objects)`` with ``setup_saved_tensors(ctx)`` backed by ``restore_from_func_ctx``, matching ``layernorm_mlp`` / ``layernorm_linear`` / ``grouped_linear`` and dropping the manual ``ctx.tensor_objects = None`` cleanup. Signed-off-by: Pawel Gadzinski Co-authored-by: Cursor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] tests: snapshot backward ctx state from LinearBwdArgs After packing the Linear backward state into ``LinearBwdArgs`` the attributes the test was reading (``backward_override``, ``fp8``, ``grad_output_quantizer``, ``reduce_and_update_bwd_fp8_tensors``) no longer live directly on ``grad_fn``. Read them from ``grad_fn.backward_objects`` when present, falling back to ``grad_fn`` for the linear-like modules that have not been refactored yet (``layernorm_linear``, ``ops_linear``). Signed-off-by: Pawel Gadzinski Co-authored-by: Cursor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Linear: add docstrings to LinearFwdArgs / LinearBwdArgs Restore the one-line class docstrings dropped during the field reorganization so pylint stops warning about C0115. Signed-off-by: Pawel Gadzinski Co-authored-by: Cursor * [PyTorch] Linear: drop ctx.backward_objects after backward Saved tensors, quantizers, weakrefs and main_grad closures referenced from LinearBwdArgs survived until ctx GC, extending peak GPU memory under retain_graph=True. Null out ctx.backward_objects right after _linear_backward so they are released as soon as backward returns. Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski Co-authored-by: Cursor Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/test_backward_override.py | 17 +- transformer_engine/pytorch/module/linear.py | 962 ++++++++++++-------- 2 files changed, 576 insertions(+), 403 deletions(-) diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py index c7c5a5b99..43e9587d9 100644 --- a/tests/pytorch/test_backward_override.py +++ b/tests/pytorch/test_backward_override.py @@ -400,23 +400,26 @@ def _snapshot_backward_ctx_state( ) -> tuple[str, bool, object, bool]: if output.grad_fn is None: raise RuntimeError("Output tensor has no grad_fn; cannot inspect backward context state.") + # ``Linear`` packs backward state into ``grad_fn.backward_objects`` + # (``LinearBwdArgs``); other linear-like modules still set the attributes + # directly on the autograd ctx. + state_holder = getattr(output.grad_fn, "backward_objects", output.grad_fn) required_attrs = ( "backward_override", "fp8", "grad_output_quantizer", "reduce_and_update_bwd_fp8_tensors", ) - missing_attrs = [attr for attr in required_attrs if not hasattr(output.grad_fn, attr)] + missing_attrs = [attr for attr in required_attrs if not hasattr(state_holder, attr)] if missing_attrs: raise RuntimeError( - "grad_fn does not expose required backward context attributes: " - f"{', '.join(missing_attrs)}." + f"Backward context does not expose required attributes: {', '.join(missing_attrs)}." ) return ( - getattr(output.grad_fn, "backward_override"), - bool(getattr(output.grad_fn, "fp8")), - getattr(output.grad_fn, "grad_output_quantizer"), - bool(getattr(output.grad_fn, "reduce_and_update_bwd_fp8_tensors")), + getattr(state_holder, "backward_override"), + bool(getattr(state_holder, "fp8")), + getattr(state_holder, "grad_output_quantizer"), + bool(getattr(state_holder, "reduce_and_update_bwd_fp8_tensors")), ) diff --git a/transformer_engine/pytorch/module/linear.py b/transformer_engine/pytorch/module/linear.py index e725387e7..dcbb9eaf9 100644 --- a/transformer_engine/pytorch/module/linear.py +++ b/transformer_engine/pytorch/module/linear.py @@ -3,7 +3,8 @@ # See LICENSE for license information. """Linear API""" -from typing import Callable, Dict, Optional, Tuple, Union, List +from dataclasses import dataclass +from typing import Any, Callable, Dict, Optional, Tuple, Union, List from functools import reduce from operator import mul as multiply_op import warnings @@ -33,7 +34,6 @@ clear_tensor_data, divide, init_method_constant, - requires_grad, needs_quantized_gemm, assert_dim_for_fp8_exec, nvtx_range_pop, @@ -80,6 +80,163 @@ __all__ = ["Linear"] +TensorOrQuantized = Union[torch.Tensor, QuantizedTensorStorage] + + +@dataclass(slots=True) +class LinearFwdArgs: + """Single-argument bag for the forward path of :class:`_Linear`.""" + + # --- Differentiable tensors (also passed positionally to autograd) --- + weight: TensorOrQuantized + inp: torch.Tensor + bias: Optional[torch.Tensor] + + # --- Non-differentiable cached tensors --- + weight_workspace: Optional[torch.Tensor] + + # --- requires_grad flags (cached so backward does not re-query) --- + input_requires_grad: bool + weight_requires_grad: bool + bias_requires_grad: bool + + # --- Quantizers --- + input_quantizer: Optional[Quantizer] + weight_quantizer: Optional[Quantizer] + output_quantizer: Optional[Quantizer] + grad_input_quantizer: Optional[Quantizer] + grad_weight_quantizer: Optional[Quantizer] + grad_output_quantizer: Optional[Quantizer] + + # --- Numerical / dtype config --- + activation_dtype: torch.dtype + fp8: bool + fp8_calibration: bool + fp8_output: bool + save_original_input: bool + backward_override: Optional[str] + custom: bool + debug: bool + + # --- Weight-workspace caching --- + is_first_microbatch: Optional[bool] + cache_weight: bool + skip_fp8_weight_update: Optional[bool] + + # --- Tensor / sequence parallelism --- + parallel_mode: Optional[str] + tp_group: Optional[Any] + tp_size: int + tensor_parallel: bool + sequence_parallel: bool + symmetric_ar_type: Optional[str] + backward_input_needs_gather: bool + + # --- Userbuffers (comm + GEMM overlap) --- + ub_name: Optional[str] + ub_overlap_ag_fprop: bool + ub_overlap_rs_fprop: bool + ub_overlap_ag_dgrad: bool + ub_overlap_rs_dgrad: bool + ub_bulk_dgrad: bool + ub_bulk_wgrad: bool + + # --- FSDP --- + fsdp_group: Optional[Any] + is_fsdp2: bool + + # --- Weight-grad scheduling --- + fuse_wgrad_accumulation: bool + wgrad_store: Optional[Any] + + # --- Misc --- + cpu_offloading: bool + is_grad_enabled: bool + + +@dataclass(slots=True) +class LinearBwdArgs: + """Single-argument bag for the backward path of :class:`_Linear`.""" + + # --- Saved / restored tensors (populated at backward entry) --- + grad_output: Optional[torch.Tensor] = None + inputmat: Optional[TensorOrQuantized] = None + weight_fp8: Optional[TensorOrQuantized] = None + saved_weight: Optional[TensorOrQuantized] = None + bias: Optional[torch.Tensor] = None + + # --- Quantizers --- + input_quantizer: Optional[Quantizer] = None + weight_quantizer: Optional[Quantizer] = None + grad_input_quantizer: Optional[Quantizer] = None + grad_weight_quantizer: Optional[Quantizer] = None + grad_output_quantizer: Optional[Quantizer] = None + + # --- Differentiability summary --- + use_bias: bool = False + requires_dgrad: bool = False + requires_wgrad: bool = False + inp_shape: Optional[torch.Size] = None + + # --- Numerical / dtype config --- + activation_dtype: Optional[torch.dtype] = None + fp8: bool = False + fp8_recipe: Optional[Recipe] = None + backward_override: Optional[str] = None + is_weight_param_quantized: bool = False + custom: bool = False + debug: bool = False + + # --- Tensor / sequence parallelism --- + parallel_mode: Optional[str] = None + tp_group: Optional[Any] = None + tp_size: int = 1 + tensor_parallel: bool = False + sequence_parallel: bool = False + backward_input_needs_gather: bool = False + + # --- Userbuffers (comm + GEMM overlap) --- + ub_name: Optional[str] = None + ub_overlap_ag: bool = False + ub_overlap_rs_dgrad: bool = False + ub_bulk_dgrad: bool = False + ub_bulk_wgrad: bool = False + + # --- FSDP --- + fsdp_group: Optional[Any] = None + fsdp_shapes: Any = None + is_fsdp2: bool = False + + # --- Weight-grad scheduling / accumulation --- + is_first_microbatch: Optional[bool] = None + fuse_wgrad_accumulation: bool = False + wgrad_store: Optional[Any] = None + origin_weight_ref: Optional[Any] = None + origin_weight_overwrites_main_grad: bool = False + main_grad_func: Optional[Callable[[], torch.Tensor]] = None + + # --- FP8 reduce-and-update bookkeeping --- + reduce_and_update_bwd_fp8_tensors: bool = False + + # --- Misc --- + cpu_offloading: bool = False + owns_input: bool = False + + # --- Per-backward scratch state (populated inside _linear_backward) --- + ub_obj_gradout: Optional[Any] = None + + def setup_saved_tensors(self, ctx: torch.autograd.function.FunctionCtx) -> None: + """Pull saved tensors from ``ctx`` into the fields backward consumes.""" + ( + self.inputmat, + self.weight_fp8, + self.saved_weight, + self.bias, + ) = restore_from_func_ctx( + ctx + ) # pylint: disable=unbalanced-tuple-unpacking + + def _check_fp8_reduce_and_update(): """Check if this is the first FP8 module (for backward reduce-and-update).""" qstate = FP8GlobalStateManager.quantization_state @@ -91,54 +248,39 @@ def _check_fp8_reduce_and_update(): def _linear_forward_impl( - weight: torch.Tensor, - weight_workspace: Optional[torch.Tensor], - inp: torch.Tensor, - bias: Optional[torch.Tensor], - non_tensor_args: Tuple, - input_quantizer: Optional[Quantizer], - weight_quantizer: Optional[Quantizer], - output_quantizer: Optional[Quantizer], -) -> Tuple: + args: LinearFwdArgs, +) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple], None, Optional[Dict]]: """Forward implementation for the linear layer. - Returns (out, tensors_to_save, tensor_objects, ctx_attrs) where the last - three are None when gradients are disabled. + Returns ``(out, new_weight_workspace, tensors_to_save_from_forward, None, + ctx_attrs)``. ``new_weight_workspace`` is the freshly produced FP8 weight + workspace (returned alongside ``out`` so the caller can refresh its + cache). The last three are ``None`` when gradients are disabled. """ - ( - is_first_microbatch, - fp8, - fp8_calibration, - _wgrad_store, - _fuse_wgrad_accumulation, - cpu_offloading, - tp_group, - tp_size, - sequence_parallel, - tensor_parallel, - activation_dtype, - parallel_mode, - is_grad_enabled, - ub_overlap_rs_fprop, - _ub_overlap_ag_dgrad, - ub_overlap_ag_fprop, - _ub_overlap_rs_dgrad, - _ub_bulk_dgrad, - _ub_bulk_wgrad, - ub_name, - _fp8_output, - fsdp_group, - cache_weight, - skip_fp8_weight_update, - symmetric_ar_type, - save_original_input, - debug, - backward_override, - custom, - backward_input_needs_gather, - is_fsdp2, - ) = non_tensor_args + weight = args.weight + inp = args.inp + bias = args.bias + input_quantizer = args.input_quantizer + weight_quantizer = args.weight_quantizer + output_quantizer = args.output_quantizer + is_first_microbatch = args.is_first_microbatch + fp8 = args.fp8 + cpu_offloading = args.cpu_offloading + tp_group = args.tp_group + sequence_parallel = args.sequence_parallel + activation_dtype = args.activation_dtype + parallel_mode = args.parallel_mode + is_grad_enabled = args.is_grad_enabled + ub_overlap_rs_fprop = args.ub_overlap_rs_fprop + ub_overlap_ag_fprop = args.ub_overlap_ag_fprop + ub_name = args.ub_name + fsdp_group = args.fsdp_group + symmetric_ar_type = args.symmetric_ar_type + save_original_input = args.save_original_input + debug = args.debug + backward_override = args.backward_override + is_fsdp2 = args.is_fsdp2 if backward_override == "high_precision": save_original_input = True @@ -189,7 +331,7 @@ def _linear_forward_impl( if fp8 or debug: if input_quantizer is None: raise ValueError("Missing quantizer for input tensor") - if not isinstance(inputmat, QuantizedTensorStorage) and not custom: + if not isinstance(inputmat, QuantizedTensorStorage) and not args.custom: own_quantized_input = True input_quantizer.set_usage( rowwise=True, @@ -280,12 +422,12 @@ def _linear_forward_impl( weightmat, new_weight_workspace = quantize_weight( tensor=weight, quantizer=weight_quantizer, - workspace=weight_workspace, + workspace=args.weight_workspace, update_workspace=update_ws, - skip_update_flag=skip_fp8_weight_update, + skip_update_flag=args.skip_fp8_weight_update, fsdp_group=fsdp_group, workspace_dtype=activation_dtype, - cache=cache_weight, + cache=args.cache_weight, ) weightmat.update_usage(rowwise_usage=True) @@ -303,7 +445,7 @@ def _linear_forward_impl( bias = cast_if_needed(bias, bias_dtype) if bias is not None else bias # Calibrate quantizers if needed - if not fp8 and fp8_calibration: + if not fp8 and args.fp8_calibration: if input_quantizer is not None: input_quantizer.calibrate(inputmat_total) if weight_quantizer is not None: @@ -363,12 +505,12 @@ def _linear_forward_impl( out = None if ub_overlap_rs_fprop: out = reduce_scatter_out - elif parallel_mode == "row" and tp_size > 1: + elif parallel_mode == "row" and args.tp_size > 1: nvtx_range_push(f"{nvtx_label}.row_parallel_comm") out = gemm_out if sequence_parallel: out, _ = reduce_scatter_along_first_dim(out, tp_group) - elif tensor_parallel: + elif args.tensor_parallel: if symmetric_ar_type is not None: out, _ = symmetric_all_reduce(out, tp_group, all_reduce_type=symmetric_ar_type) else: @@ -381,8 +523,7 @@ def _linear_forward_impl( # ------------------------------------------------------ # Prepare backward state - tensors_to_save = None - tensor_objects = None + tensors_to_save_from_forward = None ctx_attrs = None if is_grad_enabled: @@ -398,7 +539,8 @@ def _linear_forward_impl( if backward_override is not None: inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) elif ( - backward_input_needs_gather and weight_quantizer.supports_only_rowwise_all_gather() + args.backward_input_needs_gather + and weight_quantizer.supports_only_rowwise_all_gather() ): # All-gather is not supported with FP8 column-wise data inputmat.update_usage(rowwise_usage=True, columnwise_usage=False) @@ -434,226 +576,224 @@ def _linear_forward_impl( wt_save = weightmat if is_fsdp2 and weightmat is not weight: wt_save = None - tensors_to_save, tensor_objects = prepare_for_saving( - saved_inputmat, - wt_save, - weight, - bias, - ) - owns_input = saved_inputmat is not inp + # Dedup save slots that alias forward inputs; ``_linear_setup_ctx`` + # rebuilds the refs from ``inp`` / ``weight`` / ``bias``. + # Needed for torch.compile to work correctly. + saved_tensor_aliases = ( + "inp" if saved_inputmat is inp else None, + "weight" if wt_save is weight else None, + "weight", # ``saved_weight`` slot is always the weight parameter + "bias" if bias is not None else None, + ) + tensors_to_save_from_forward = ( + None if saved_tensor_aliases[0] is not None else saved_inputmat, + None if saved_tensor_aliases[1] is not None else wt_save, + None, + None if saved_tensor_aliases[3] is not None else bias, + ) ctx_attrs = { - "weight_quantizer": weight_quantizer, "fsdp_shapes": fsdp_shapes, - "owns_input": owns_input, - "is_fsdp2": is_fsdp2, + "saved_tensor_aliases": saved_tensor_aliases, } - return out, new_weight_workspace, tensors_to_save, tensor_objects, ctx_attrs + return out, new_weight_workspace, tensors_to_save_from_forward, None, ctx_attrs def _linear_setup_ctx( - ctx, - tensors_to_save, - tensor_objects, - ctx_attrs, - inp, - weight, - bias, - non_tensor_args, - input_quantizer, - grad_input_quantizer, - grad_weight_quantizer, - grad_output_quantizer, -): - """Save forward state into autograd context for backward pass.""" - ctx.save_for_backward(*tensors_to_save) - ctx.tensor_objects = tensor_objects - - ( - is_first_microbatch, - fp8, - _fp8_calibration, - wgrad_store, - fuse_wgrad_accumulation, - cpu_offloading, - tp_group, - tp_size, - sequence_parallel, - tensor_parallel, - activation_dtype, - parallel_mode, - _is_grad_enabled, - _ub_overlap_rs_fprop, - ub_overlap_ag_dgrad, - _ub_overlap_ag_fprop, - ub_overlap_rs_dgrad, - ub_bulk_dgrad, - ub_bulk_wgrad, - ub_name, - _fp8_output, - fsdp_group, - _cache_weight, - _skip_fp8_weight_update, - _symmetric_ar_type, - _save_original_input, - debug, - backward_override, - custom, - backward_input_needs_gather, - _is_fsdp2, - ) = non_tensor_args - - # Values derived from input tensors - ctx.use_bias = bias is not None - ctx.requires_dgrad = inp.requires_grad - ctx.requires_wgrad = weight.requires_grad - ctx.inp_shape = inp.shape + bwd_args: LinearBwdArgs, + fwd_args: LinearFwdArgs, + out: torch.Tensor, + ctx_attrs: Dict, + tensors_to_save_from_forward: Tuple[Any, ...], +) -> Tuple[Any, ...]: + """Populate ``bwd_args`` from forward state. + + Returns the merged list of tensors that should be passed through + ``prepare_for_saving`` by the caller (``_Linear.forward``). Keeping the + ``prepare_for_saving`` call out of here lets callers stitch in extra + tensors (e.g. the original ``weight`` parameter so backward can reuse it + for FSDP2 re-quantization) without having to mutate the structured + metadata returned by ``prepare_for_saving``. + """ + del out # No-op; kept for symmetry with the compile-time helper signature. + + inp = fwd_args.inp + weight = fwd_args.weight + bias = fwd_args.bias + + backward_override = fwd_args.backward_override + fp8 = fwd_args.fp8 + fuse_wgrad_accumulation = fwd_args.fuse_wgrad_accumulation # Quantizers - ctx.input_quantizer = input_quantizer - ctx.grad_input_quantizer = grad_input_quantizer - ctx.grad_weight_quantizer = grad_weight_quantizer - ctx.grad_output_quantizer = grad_output_quantizer - - # Values from non_tensor_args - ctx.activation_dtype = activation_dtype - ctx.fp8 = fp8 - ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None - ctx.backward_override = backward_override - ctx.is_weight_param_quantized = isinstance(weight, QuantizedTensorStorage) - ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation - ctx.cpu_offloading = cpu_offloading - ctx.is_first_microbatch = is_first_microbatch - ctx.sequence_parallel = sequence_parallel - ctx.tensor_parallel = tensor_parallel - ctx.parallel_mode = parallel_mode - ctx.tp_group = tp_group - ctx.tp_size = tp_size - ctx.ub_name = ub_name - ctx.fsdp_group = fsdp_group - ctx.debug = debug - ctx.wgrad_store = wgrad_store - ctx.ub_overlap_ag = ub_overlap_ag_dgrad - - ctx.ub_overlap_rs_dgrad = ub_overlap_rs_dgrad - ctx.ub_bulk_dgrad = ub_bulk_dgrad - ctx.ub_bulk_wgrad = ub_bulk_wgrad - - # Derived values - ctx.backward_input_needs_gather = backward_input_needs_gather - ctx.custom = custom - - # main_grad_func setup - if fuse_wgrad_accumulation and weight.requires_grad: - ctx.origin_weight_ref = weakref.ref(weight) - ctx.origin_weight_overwrites_main_grad = getattr(weight, "overwrite_main_grad", False) + bwd_args.input_quantizer = fwd_args.input_quantizer + bwd_args.weight_quantizer = ( + weight._quantizer if isinstance(weight, QuantizedTensor) else fwd_args.weight_quantizer + ) + bwd_args.grad_input_quantizer = fwd_args.grad_input_quantizer + bwd_args.grad_weight_quantizer = fwd_args.grad_weight_quantizer + bwd_args.grad_output_quantizer = fwd_args.grad_output_quantizer + + # Differentiability summary + bwd_args.use_bias = bias is not None + bwd_args.requires_dgrad = fwd_args.input_requires_grad + bwd_args.requires_wgrad = fwd_args.weight_requires_grad + bwd_args.inp_shape = inp.shape + + # Numerical / dtype config + bwd_args.activation_dtype = fwd_args.activation_dtype + bwd_args.fp8 = fp8 + bwd_args.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + bwd_args.backward_override = backward_override + bwd_args.is_weight_param_quantized = isinstance(weight, QuantizedTensorStorage) + bwd_args.custom = fwd_args.custom + bwd_args.debug = fwd_args.debug + + # Tensor / sequence parallelism + bwd_args.parallel_mode = fwd_args.parallel_mode + bwd_args.tp_group = fwd_args.tp_group + bwd_args.tp_size = fwd_args.tp_size + bwd_args.tensor_parallel = fwd_args.tensor_parallel + bwd_args.sequence_parallel = fwd_args.sequence_parallel + bwd_args.backward_input_needs_gather = fwd_args.backward_input_needs_gather + + # Userbuffers + bwd_args.ub_name = fwd_args.ub_name + bwd_args.ub_overlap_ag = fwd_args.ub_overlap_ag_dgrad + bwd_args.ub_overlap_rs_dgrad = fwd_args.ub_overlap_rs_dgrad + bwd_args.ub_bulk_dgrad = fwd_args.ub_bulk_dgrad + bwd_args.ub_bulk_wgrad = fwd_args.ub_bulk_wgrad + + # FSDP + bwd_args.fsdp_group = fwd_args.fsdp_group + bwd_args.fsdp_shapes = ctx_attrs["fsdp_shapes"] + bwd_args.is_fsdp2 = fwd_args.is_fsdp2 + + # Weight-grad scheduling / accumulation + bwd_args.is_first_microbatch = fwd_args.is_first_microbatch + bwd_args.fuse_wgrad_accumulation = fuse_wgrad_accumulation + bwd_args.wgrad_store = fwd_args.wgrad_store + if fuse_wgrad_accumulation and fwd_args.weight_requires_grad: + bwd_args.origin_weight_ref = weakref.ref(weight) + bwd_args.origin_weight_overwrites_main_grad = getattr(weight, "overwrite_main_grad", False) if hasattr(weight, "__fsdp_param__"): - ctx.main_grad_func = weight.get_main_grad + bwd_args.main_grad_func = weight.get_main_grad else: - ctx.main_grad_func = lambda: weight.main_grad + bwd_args.main_grad_func = lambda: weight.main_grad - # Forward-computed values that can't be derived here - ctx.weight_quantizer = ctx_attrs["weight_quantizer"] - ctx.fsdp_shapes = ctx_attrs["fsdp_shapes"] - ctx.owns_input = ctx_attrs["owns_input"] - ctx.is_fsdp2 = ctx_attrs["is_fsdp2"] + # Misc + bwd_args.cpu_offloading = fwd_args.cpu_offloading - # backward overrides if backward_override is not None: - ctx.fp8 = False - ctx.debug = False - ctx.ub_overlap_ag = False - ctx.ub_overlap_rs_dgrad = False - ctx.ub_bulk_dgrad = False - ctx.ub_bulk_wgrad = False - ctx.grad_input_quantizer = None - ctx.grad_weight_quantizer = None - ctx.grad_output_quantizer = None - - -def _linear_backward( - ctx, - grad_output: torch.Tensor, - input_quantizer: Optional[Quantizer], - weight_quantizer: Optional[Quantizer], - grad_input_quantizer: Optional[Quantizer], - grad_weight_quantizer: Optional[Quantizer], - grad_output_quantizer: Optional[Quantizer], -) -> Tuple[Union[torch.Tensor, None], ...]: - """Backward implementation for the linear layer.""" + bwd_args.fp8 = False + bwd_args.debug = False + bwd_args.ub_overlap_ag = False + bwd_args.ub_overlap_rs_dgrad = False + bwd_args.ub_bulk_dgrad = False + bwd_args.ub_bulk_wgrad = False + bwd_args.grad_input_quantizer = None + bwd_args.grad_weight_quantizer = None + bwd_args.grad_output_quantizer = None + + saved_inputmat, wt_save, saved_weight, saved_bias = tensors_to_save_from_forward + inputmat_alias, wt_save_alias, saved_weight_alias, bias_alias = ctx_attrs[ + "saved_tensor_aliases" + ] + bwd_args.owns_input = inputmat_alias != "inp" + if inputmat_alias == "inp": + saved_inputmat = inp + if wt_save_alias == "weight": + wt_save = weight + if saved_weight_alias == "weight": + saved_weight = weight + if bias_alias == "bias": + saved_bias = bias + return (saved_inputmat, wt_save, saved_weight, saved_bias) + + +def _linear_backward(args: LinearBwdArgs) -> Tuple[Union[torch.Tensor, None], ...]: + """Backward implementation for the linear layer. + + Caller must have populated ``args.grad_output`` and run + ``args.setup_saved_tensors(ctx)`` before invocation. + """ + bwd_args = args + grad_output = args.grad_output + assert grad_output is not None + inputmat = args.inputmat + weight_fp8 = args.weight_fp8 + saved_weight = args.saved_weight + bias = args.bias + input_quantizer = args.input_quantizer + weight_quantizer = args.weight_quantizer + grad_input_quantizer = args.grad_input_quantizer + grad_weight_quantizer = args.grad_weight_quantizer + grad_output_quantizer = args.grad_output_quantizer # NVTX label for profiling nvtx_label = "transformer_engine._Linear.backward" - if ctx.ub_name is not None: - nvtx_label = f"{nvtx_label}.{ctx.ub_name}" + if bwd_args.ub_name is not None: + nvtx_label = f"{nvtx_label}.{bwd_args.ub_name}" with get_nvtx_range_context("_Linear_backward"): - ( - inputmat, - weight_fp8, - saved_weight, - bias, - ) = restore_from_func_ctx( # pylint: disable=unbalanced-tuple-unpacking - ctx - ) - origin_weight_python_object = None - origin_weight_overwrites_main_grad = getattr( - ctx, "origin_weight_overwrites_main_grad", False - ) + origin_weight_overwrites_main_grad = bwd_args.origin_weight_overwrites_main_grad main_grad = None - if ctx.fuse_wgrad_accumulation and ctx.requires_wgrad: - origin_weight_ref = ctx.origin_weight_ref - ctx.origin_weight_ref = None + if bwd_args.fuse_wgrad_accumulation and bwd_args.requires_wgrad: + origin_weight_ref = bwd_args.origin_weight_ref + bwd_args.origin_weight_ref = None origin_weight_python_object = ( origin_weight_ref() if origin_weight_ref is not None else None ) assert ( origin_weight_python_object is not None ), "weight was removed while fuse_wgrad_accumulation=True" - main_grad = ctx.main_grad_func() + main_grad = bwd_args.main_grad_func() origin_weight_python_object.main_grad = main_grad # Gather intermediate/activation tensors if needed - # NOTE: weight_fp8 = weight when ctx.fp8 == False and torch.disttributed.FSDP already + # NOTE: weight_fp8 = weight when bwd_args.fp8 == False and torch.disttributed.FSDP already # shards/unshards the base weights so we don't do it ourselves nvtx_range_push(f"{nvtx_label}.fsdp_gather") _fsdp_gather_tensors( - ctx.fsdp_group, - ctx.fsdp_shapes, + bwd_args.fsdp_group, + bwd_args.fsdp_shapes, inputmat, weight_fp8, ) nvtx_range_pop(f"{nvtx_label}.fsdp_gather") # Configure Userbuffers communication (comm+GEMM overlap) - ctx.ub_obj_gradout = None + bwd_args.ub_obj_gradout = None ub_obj_dgrad = None ub_obj_wgrad = None ub_type_dgrad = None ub_type_wgrad = None - dgrad_shape = [reduce(multiply_op, ctx.inp_shape[:-1]), ctx.inp_shape[-1]] - if ctx.ub_overlap_ag: + dgrad_shape = [ + reduce(multiply_op, bwd_args.inp_shape[:-1]), + bwd_args.inp_shape[-1], + ] + if bwd_args.ub_overlap_ag: # Overlap grad_output all-gather with dgrad compute - ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8) - ub_obj_dgrad = ctx.ub_obj_gradout + bwd_args.ub_obj_gradout = get_ub(bwd_args.ub_name + "_dgrad", bwd_args.fp8) + ub_obj_dgrad = bwd_args.ub_obj_gradout ub_type_dgrad = tex.CommOverlapType.AG - elif ctx.ub_overlap_rs_dgrad: + elif bwd_args.ub_overlap_rs_dgrad: # Overlap dgrad reduce-scatter with dgrad compute - ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8) - ub_obj_dgrad = ctx.ub_obj_gradout + bwd_args.ub_obj_gradout = get_ub(bwd_args.ub_name + "_dgrad", bwd_args.fp8) + ub_obj_dgrad = bwd_args.ub_obj_gradout ub_type_dgrad = tex.CommOverlapType.RS else: - if ctx.ub_bulk_dgrad: + if bwd_args.ub_bulk_dgrad: # Overlap inputmat all-gather with dgrad compute - ctx.ub_obj_gradout = get_ub(ctx.ub_name + "_dgrad", ctx.fp8) - ub_obj_dgrad = ctx.ub_obj_gradout + bwd_args.ub_obj_gradout = get_ub(bwd_args.ub_name + "_dgrad", bwd_args.fp8) + ub_obj_dgrad = bwd_args.ub_obj_gradout ub_type_dgrad = tex.CommOverlapType.AG - if ctx.ub_bulk_wgrad: + if bwd_args.ub_bulk_wgrad: # Overlap dgrad reduce-scatter with wgrad compute - ub_obj_wgrad = get_ub(ctx.ub_name + "_wgrad", ctx.fp8) + ub_obj_wgrad = get_ub(bwd_args.ub_name + "_wgrad", bwd_args.fp8) ub_type_wgrad = tex.CommOverlapType.RS # -------------------------------------------------- @@ -670,7 +810,7 @@ def _linear_backward( if grad_output_quantizer is not None: quantizer = grad_output_quantizer quantizer.set_usage(rowwise=True, columnwise=True) - if ctx.ub_overlap_ag: + if bwd_args.ub_overlap_ag: # Userbuffers only supports communication for one # tensor usage at a time. Configure quantizer with # usage for only dgrad GEMM. @@ -680,20 +820,28 @@ def _linear_backward( # on whether wgrad calculations will be performed. # NOTE: If requires_dgrad is False, disabling `rowwise` quantization and keeping `columnwise` quantization # results in `Assertion failed: output_tensor->has_data(). Quantizing in only the columnwise direction not supported yet!` - # NOTE: For `ctx.bias is True`, selected quantize kernel errors with + # NOTE: For `bias is True`, selected quantize kernel errors with # `cast_kernels.cuh:1322 in function fp8_quantize_arch_l_100: Not implemented scaling mode or fusion: NVTE_DELAYED_TENSOR_SCALING or IS_DBIAS=true on GPU with compute capability < 10.0.` - if not ctx.use_bias and not ctx.requires_wgrad and grad_output_quantizer is not None: + if ( + not bwd_args.use_bias + and not bwd_args.requires_wgrad + and grad_output_quantizer is not None + ): grad_output_quantizer.set_usage(columnwise=False) - # Prepare grad output tensor + # Prepare grad output tensor. + # ``grad_output_preprocess`` accesses a small set of attributes + # (sequence_parallel, fp8, backward_override, debug, ub_overlap_ag, + # tp_group, ub_obj_gradout, use_bias). ``LinearBwdArgs`` exposes the + # same names so we can pass it directly. nvtx_range_push(f"{nvtx_label}.grad_output_preprocess") ( grad_output, grad_bias, ) = TransformerEngineBaseModule.grad_output_preprocess( - ctx, + bwd_args, grad_output, - ctx.parallel_mode == "row", + bwd_args.parallel_mode == "row", grad_output_quantizer, ) nvtx_range_pop(f"{nvtx_label}.grad_output_preprocess") @@ -710,12 +858,12 @@ def _linear_backward( # -------------------------------------------------- inputmat_total = None inputmat_total_work = None - if ctx.requires_wgrad: - if ctx.fp8 or ctx.debug: + if bwd_args.requires_wgrad: + if bwd_args.fp8 or bwd_args.debug: if isinstance(inputmat, QuantizedTensorStorage): # Input tensor is already quantized pass - elif ctx.debug or ctx.custom: + elif bwd_args.debug or bwd_args.custom: # Debug quantizer will be applied immediately before wgrad GEMM pass else: @@ -725,19 +873,19 @@ def _linear_backward( # All-gather is not supported with FP8 column-wise data quantizer.set_usage( rowwise=True, - columnwise=not ctx.backward_input_needs_gather, + columnwise=not bwd_args.backward_input_needs_gather, ) else: quantizer.set_usage(rowwise=False, columnwise=True) inputmat = quantizer(inputmat) else: if isinstance(inputmat, QuantizedTensorStorage): - inputmat = inputmat.dequantize(dtype=ctx.activation_dtype) + inputmat = inputmat.dequantize(dtype=bwd_args.activation_dtype) else: - inputmat = cast_if_needed(inputmat, ctx.activation_dtype) - if ctx.backward_input_needs_gather: + inputmat = cast_if_needed(inputmat, bwd_args.activation_dtype) + if bwd_args.backward_input_needs_gather: quantizer = None - if ctx.fp8 or ctx.debug: + if bwd_args.fp8 or bwd_args.debug: quantizer = input_quantizer if quantizer.supports_only_rowwise_all_gather(): # If data is in FP8, we compute FP8 transposes manually @@ -745,18 +893,18 @@ def _linear_backward( else: # wgrad GEMM requires input with column-wise usage quantizer.set_usage(rowwise=False, columnwise=True) - if ctx.ub_bulk_dgrad: + if bwd_args.ub_bulk_dgrad: inputmat_total, _ = fill_userbuffers_buffer_for_all_gather( ub_obj_dgrad, inputmat, quantizer, - ctx.tp_group, + bwd_args.tp_group, ) else: nvtx_range_push(f"{nvtx_label}.column_parallel_comm_input") inputmat_total, inputmat_total_work = gather_along_first_dim( inputmat, - ctx.tp_group, + bwd_args.tp_group, async_op=True, quantizer=quantizer, ) @@ -773,7 +921,7 @@ def _linear_backward( dgrad = None dgrad_work = None - if ctx.requires_dgrad: + if bwd_args.requires_dgrad: # FSDP2: Re-create workspace from all-gathered weight when # workspace was not saved. (Issue #2681) @@ -784,15 +932,15 @@ def _linear_backward( # saved weight is already set to right usages by # fsdp2 quantized-tensor hooks when workspace was not saved. weight_fp8 = saved_weight - elif ctx.weight_quantizer is not None: - ctx.weight_quantizer.set_usage(rowwise=True, columnwise=True) - weight_fp8 = ctx.weight_quantizer(saved_weight) + elif bwd_args.weight_quantizer is not None: + bwd_args.weight_quantizer.set_usage(rowwise=True, columnwise=True) + weight_fp8 = bwd_args.weight_quantizer(saved_weight) # Make sure required data is available if isinstance(grad_output, QuantizedTensorStorage): grad_output.update_usage(rowwise_usage=True) if ( - ctx.fp8 + bwd_args.fp8 and weight_quantizer is not None and isinstance(weight_fp8, QuantizedTensorStorage) ): @@ -800,8 +948,8 @@ def _linear_backward( # Choose whether to use GEMM kernel with split accumulator use_split_accumulator = _2X_ACC_DGRAD - if ctx.fp8: - recipe = ctx.fp8_recipe + if bwd_args.fp8: + recipe = bwd_args.fp8_recipe if hasattr(recipe, "fp8_gemm_dgrad"): use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator @@ -812,11 +960,13 @@ def _linear_backward( # Output buffers for Userbuffers reduce-scatter gemm_out = None reduce_scatter_out = None - if ctx.ub_overlap_rs_dgrad: + if bwd_args.ub_overlap_rs_dgrad: reduce_scatter_out = torch.empty( - dgrad_shape, dtype=ctx.activation_dtype, device=grad_output_arg.device + dgrad_shape, + dtype=bwd_args.activation_dtype, + device=grad_output_arg.device, ) - elif ctx.ub_bulk_wgrad: + elif bwd_args.ub_bulk_wgrad: gemm_out = ub_obj_wgrad.get_buffer(local_chunk=False) # dgrad GEMM @@ -824,15 +974,15 @@ def _linear_backward( nvtx_range_push(f"{nvtx_label}.dgrad_gemm") weight_for_dgrad = weight_fp8 - if ctx.backward_override == "dequantized": + if bwd_args.backward_override == "dequantized": if isinstance(weight_for_dgrad, QuantizedTensorStorage): - weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) + weight_for_dgrad = weight_for_dgrad.dequantize(dtype=bwd_args.activation_dtype) else: - weight_for_dgrad = cast_if_needed(weight_for_dgrad, ctx.activation_dtype) - elif ctx.backward_override == "high_precision": + weight_for_dgrad = cast_if_needed(weight_for_dgrad, bwd_args.activation_dtype) + elif bwd_args.backward_override == "high_precision": weight_for_dgrad = saved_weight if isinstance(weight_for_dgrad, QuantizedTensorStorage): - weight_for_dgrad = weight_for_dgrad.dequantize(dtype=ctx.activation_dtype) + weight_for_dgrad = weight_for_dgrad.dequantize(dtype=bwd_args.activation_dtype) gemm_out, *_, reduce_scatter_out = general_gemm( weight_for_dgrad, grad_output, @@ -840,12 +990,12 @@ def _linear_backward( grad=True, quantization_params=grad_input_quantizer, out=gemm_out, - out_dtype=ctx.activation_dtype, + out_dtype=bwd_args.activation_dtype, use_split_accumulator=use_split_accumulator, ub=ub_obj_dgrad, ub_type=ub_type_dgrad, extra_output=reduce_scatter_out, - bulk_overlap=ctx.ub_bulk_dgrad, + bulk_overlap=bwd_args.ub_bulk_dgrad, ) nvtx_range_pop(f"{nvtx_label}.dgrad_gemm") @@ -854,26 +1004,26 @@ def _linear_backward( # and 2d block-scaled weights in TE managed memory. So we need to clear # it here. # (Issues #2681, #2717) - if getattr(ctx, "is_fsdp2", False) and isinstance(weight_fp8, QuantizedTensorStorage): + if bwd_args.is_fsdp2 and isinstance(weight_fp8, QuantizedTensorStorage): clear_columnwise_cache(weight_fp8) # Prepare grad input tensor # Note: Perform tensor-parallel communication - if ctx.ub_overlap_rs_dgrad: + if bwd_args.ub_overlap_rs_dgrad: dgrad = reduce_scatter_out - elif ctx.ub_bulk_wgrad: + elif bwd_args.ub_bulk_wgrad: dgrad = ub_obj_wgrad.get_buffer(local_chunk=True) - elif ctx.parallel_mode == "column" and ctx.tp_size > 1: + elif bwd_args.parallel_mode == "column" and bwd_args.tp_size > 1: nvtx_range_push(f"{nvtx_label}.column_parallel_comm_dgrad") dgrad = gemm_out - if ctx.sequence_parallel: + if bwd_args.sequence_parallel: dgrad, dgrad_work = reduce_scatter_along_first_dim( dgrad, - ctx.tp_group, + bwd_args.tp_group, async_op=True, ) else: - dgrad, dgrad_work = allreduce(dgrad, ctx.tp_group, async_op=True) + dgrad, dgrad_work = allreduce(dgrad, bwd_args.tp_group, async_op=True) nvtx_range_pop(f"{nvtx_label}.column_parallel_comm_dgrad") else: dgrad = gemm_out @@ -887,7 +1037,7 @@ def _linear_backward( # -------------------------------------------------- wgrad = None - if ctx.requires_wgrad: + if bwd_args.requires_wgrad: # Prepare input tensor # Note: Synchronize tensor-parallel communication and @@ -895,7 +1045,7 @@ def _linear_backward( if inputmat_total_work is not None: inputmat_total_work.wait() inputmat_total_work = None - if ctx.fp8 or ctx.debug: + if bwd_args.fp8 or bwd_args.debug: if isinstance(inputmat_total, QuantizedTensorStorage): inputmat_total.update_usage(columnwise_usage=True) else: @@ -905,7 +1055,7 @@ def _linear_backward( # Prepare grad output tensor # Note: Synchronize tensor-parallel communication and # make sure required data is available - if ctx.ub_overlap_ag and isinstance(grad_output_quantizer, MXFP8Quantizer): + if bwd_args.ub_overlap_ag and isinstance(grad_output_quantizer, MXFP8Quantizer): # UB does not support pipelined overlapping grad output # all-gather with wgrad GEMM. Also, we can't # convert row-scaled MXFP8 to column-scaled, so we @@ -917,7 +1067,7 @@ def _linear_backward( dgrad_send_stream, dgrad_recv_stream = ub_obj_dgrad.get_communication_stream() # This object is separate from the ub_obj_wgrad object which is passed to the GEMM - ub_obj_overlap_wgrad = get_ub(ctx.ub_name + "_wgrad", ctx.fp8) + ub_obj_overlap_wgrad = get_ub(bwd_args.ub_name + "_wgrad", bwd_args.fp8) grad_output_quantizer.set_usage(rowwise=False, columnwise=True) @@ -929,7 +1079,7 @@ def _linear_backward( ub_obj_overlap_wgrad, grad_output_arg, grad_output_quantizer, - ctx.tp_group, + bwd_args.tp_group, ) # Allgather grad_outputs[0] using the dgrad streams so we can overlap with the fc2_dgrad gemm @@ -937,7 +1087,7 @@ def _linear_backward( ub_obj_overlap_wgrad, dgrad_send_stream, dgrad_recv_stream ) - if ctx.fp8 or ctx.debug: + if bwd_args.fp8 or bwd_args.debug: if isinstance(grad_output, QuantizedTensorStorage): grad_output.update_usage(columnwise_usage=True) else: @@ -946,31 +1096,35 @@ def _linear_backward( # Figure out whether to use split accumulator use_split_accumulator = _2X_ACC_WGRAD - if ctx.fp8: - recipe = ctx.fp8_recipe + if bwd_args.fp8: + recipe = bwd_args.fp8_recipe if hasattr(recipe, "fp8_gemm_wgrad"): use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator # Figure out whether to output wgrad GEMM directly into main grad - if ctx.is_first_microbatch is not None: + if bwd_args.is_first_microbatch is not None: accumulate_wgrad_into_param_main_grad = ( - ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch + bwd_args.fuse_wgrad_accumulation and not bwd_args.is_first_microbatch ) else: - accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation + accumulate_wgrad_into_param_main_grad = bwd_args.fuse_wgrad_accumulation # Output buffer for overlapping FP8 grad input # reduce-scatter with wgrad GEMM reduce_scatter_out = None - if ctx.ub_bulk_wgrad and ub_obj_wgrad.is_fp8_ubuf(): + if bwd_args.ub_bulk_wgrad and ub_obj_wgrad.is_fp8_ubuf(): reduce_scatter_out = torch.empty( - dgrad_shape, dtype=ctx.activation_dtype, device=grad_output_arg.device + dgrad_shape, + dtype=bwd_args.activation_dtype, + device=grad_output_arg.device, ) # Arguments to include in wgrad GEMM closure wgrad_gemm_kwargs = { "out_dtype": ( - main_grad.dtype if ctx.fuse_wgrad_accumulation else ctx.activation_dtype + main_grad.dtype + if bwd_args.fuse_wgrad_accumulation + else bwd_args.activation_dtype ), "quantization_params": grad_weight_quantizer, "accumulate": ( @@ -979,14 +1133,14 @@ def _linear_backward( else False ), "layout": "NT", - "out": main_grad if ctx.fuse_wgrad_accumulation else None, - "bias": (bias if (grad_bias is None and not ctx.fp8) else None), + "out": main_grad if bwd_args.fuse_wgrad_accumulation else None, + "bias": (bias if (grad_bias is None and not bwd_args.fp8) else None), "use_split_accumulator": use_split_accumulator, "grad": True, "ub": ub_obj_wgrad, "ub_type": ub_type_wgrad, "extra_output": reduce_scatter_out, - "bulk_overlap": ctx.ub_bulk_wgrad, + "bulk_overlap": bwd_args.ub_bulk_wgrad, } def wgrad_gemm( @@ -1007,7 +1161,7 @@ def wgrad_gemm( return dw, db # Choose whether to call wgrad GEMM now or delay - if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): + if bwd_args.wgrad_store is not None and bwd_args.wgrad_store.delay_wgrad_compute(): if ( wgrad_gemm_kwargs["ub"] is not None or wgrad_gemm_kwargs["ub_type"] is not None @@ -1018,7 +1172,7 @@ def wgrad_gemm( "Delayed weight grad computation is not supported " "with Userbuffers (tensor-parallel communication overlapping)" ) - ctx.wgrad_store.put([inputmat_total, grad_output], wgrad_gemm) + bwd_args.wgrad_store.put([inputmat_total, grad_output], wgrad_gemm) else: # Call wgrad GEMM now @@ -1030,18 +1184,18 @@ def wgrad_gemm( del grad_bias_ # Deallocate tensors if permitted - if ctx.owns_input: + if bwd_args.owns_input: # Input tensor is internal clear_tensor_data(inputmat_total) - elif ctx.backward_input_needs_gather: + elif bwd_args.backward_input_needs_gather: # Gathered input tensor is internal clear_tensor_data(inputmat_total) - if ctx.parallel_mode == "row" and ctx.sequence_parallel: + if bwd_args.parallel_mode == "row" and bwd_args.sequence_parallel: # Gathered grad output tensor is internal clear_tensor_data(grad_output) # Update grad input if overlapping reduce-scatter with wgrad GEMM - if ctx.ub_bulk_wgrad: + if bwd_args.ub_bulk_wgrad: if ub_obj_wgrad.is_fp8_ubuf(): dgrad = reduce_scatter_out else: @@ -1052,7 +1206,7 @@ def wgrad_gemm( # -------------------------------------------------- # Don't return grad bias if not needed - if not ctx.use_bias: + if not bwd_args.use_bias: grad_bias = None # Make sure all tensor-parallel communication is finished @@ -1063,9 +1217,9 @@ def wgrad_gemm( dgrad_work.wait() dgrad_work = None - if ctx.requires_wgrad: + if bwd_args.requires_wgrad: # Handle custom DDP from mcore. - if ctx.fuse_wgrad_accumulation and hasattr( + if bwd_args.fuse_wgrad_accumulation and hasattr( origin_weight_python_object, "grad_added_to_main_grad" ): origin_weight_python_object.grad_added_to_main_grad = True @@ -1080,26 +1234,18 @@ def wgrad_gemm( list(main_grad.shape), origin_weight_python_object.dtype, ) - elif ctx.fuse_wgrad_accumulation: + elif bwd_args.fuse_wgrad_accumulation: wgrad = None else: wgrad = None # Scatter fp8 weight buffers - if ctx.fp8 and not ctx.is_weight_param_quantized: - _fsdp_scatter_tensors(ctx.fsdp_group, weight_fp8) + if bwd_args.fp8 and not bwd_args.is_weight_param_quantized: + _fsdp_scatter_tensors(bwd_args.fsdp_group, weight_fp8) return ( wgrad, - None, # weight_workspace - dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, + dgrad.view(bwd_args.inp_shape) if bwd_args.requires_dgrad else None, grad_bias, - None, - None, - None, - None, - None, - None, - None, ) @@ -1112,73 +1258,75 @@ class _Linear(torch.autograd.Function): def forward( ctx, weight: torch.Tensor, - weight_workspace: Optional[torch.Tensor], inp: torch.Tensor, bias: Optional[torch.Tensor], - non_tensor_args: Tuple, - input_quantizer: Optional[Quantizer], - weight_quantizer: Optional[Quantizer], - output_quantizer: Optional[Quantizer], - grad_input_quantizer: Optional[Quantizer], - grad_weight_quantizer: Optional[Quantizer], - grad_output_quantizer: Optional[Quantizer], + fwd_args: LinearFwdArgs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: - """Forward pass: compute linear output and set up autograd context.""" - out, new_weight_workspace, tensors_to_save, tensor_objects, ctx_attrs = ( - _linear_forward_impl( - weight, - weight_workspace, - inp, - bias, - non_tensor_args, - input_quantizer, - weight_quantizer, - output_quantizer, - ) - ) + """Forward pass: compute linear output and set up autograd context. + + ``weight``, ``inp`` and ``bias`` are positional Tensor arguments so + autograd tracks them; they are immediately re-attached to ``fwd_args`` + so every downstream helper can be invoked with a single argument. + + ``weight_workspace`` is intentionally NOT a positional input: it is a + non-differentiable cached tensor passed in via + ``fwd_args.weight_workspace`` and the freshly produced workspace is + returned as a separate output so the module can refresh its cache. + """ + fwd_args.weight = weight + fwd_args.inp = inp + fwd_args.bias = bias + ( + out, + new_weight_workspace, + tensors_to_save_from_forward, + _, + ctx_attrs, + ) = _linear_forward_impl(fwd_args) if ctx is not None: - _linear_setup_ctx( - ctx, - tensors_to_save, - tensor_objects, + bwd_args = LinearBwdArgs() + tensors_to_save_from_setup = _linear_setup_ctx( + bwd_args, + fwd_args, + out, ctx_attrs, - inp, - weight, - bias, - non_tensor_args, - input_quantizer=input_quantizer, - grad_input_quantizer=grad_input_quantizer, - grad_weight_quantizer=grad_weight_quantizer, - grad_output_quantizer=grad_output_quantizer, + tensors_to_save_from_forward, ) - fp8 = non_tensor_args[1] - if fp8 and requires_grad(inp, weight, bias): - ctx.reduce_and_update_bwd_fp8_tensors = _check_fp8_reduce_and_update() - else: - ctx.reduce_and_update_bwd_fp8_tensors = False - if ctx.backward_override is not None: - ctx.reduce_and_update_bwd_fp8_tensors = False + tensors_to_save, tensor_objects = prepare_for_saving(*tensors_to_save_from_setup) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects + ctx.backward_objects = bwd_args + if fwd_args.fp8 and ( + fwd_args.input_requires_grad + or fwd_args.weight_requires_grad + or fwd_args.bias_requires_grad + ): + bwd_args.reduce_and_update_bwd_fp8_tensors = _check_fp8_reduce_and_update() + if fwd_args.backward_override is not None: + bwd_args.reduce_and_update_bwd_fp8_tensors = False return out, new_weight_workspace @staticmethod def backward( - ctx, grad_output: torch.Tensor, _grad_weight_workspace + ctx, + grad_output: torch.Tensor, + _grad_weight_workspace, ) -> Tuple[Union[torch.Tensor, None], ...]: """Backward pass: compute gradients and reduce FP8 scaling factors.""" + bwd_args: LinearBwdArgs = ctx.backward_objects + bwd_args.grad_output = grad_output + bwd_args.setup_saved_tensors(ctx) nvtx_label = "transformer_engine._Linear.backward" - if ctx.ub_name is not None: - nvtx_label = f"{nvtx_label}.{ctx.ub_name}" - result = _linear_backward( - ctx, - grad_output, - input_quantizer=ctx.input_quantizer, - weight_quantizer=ctx.weight_quantizer, - grad_input_quantizer=ctx.grad_input_quantizer, - grad_weight_quantizer=ctx.grad_weight_quantizer, - grad_output_quantizer=ctx.grad_output_quantizer, - ) - if ctx.reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): + if bwd_args.ub_name is not None: + nvtx_label = f"{nvtx_label}.{bwd_args.ub_name}" + result = _linear_backward(bwd_args) + (None,) # fwd_args grad slot + reduce_and_update_bwd_fp8_tensors = bwd_args.reduce_and_update_bwd_fp8_tensors + # Drop all references held by bwd_args (saved tensors, quantizers, weakrefs, + # main_grad closure) so they don't outlive backward via ctx under retain_graph. + ctx.backward_objects = None + del bwd_args + if reduce_and_update_bwd_fp8_tensors and not is_graph_capturing(): nvtx_range_push(f"{nvtx_label}.reduce_and_update_fp8_tensors") FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) nvtx_range_pop(f"{nvtx_label}.reduce_and_update_fp8_tensors") @@ -1685,52 +1833,74 @@ def forward( ub_bulk_dgrad = self.ub_bulk_dgrad ub_bulk_wgrad = self.ub_bulk_wgrad - non_tensor_args = ( - is_first_microbatch, - self.fp8, - self.fp8_calibration, - self.wgrad_store, - self.fuse_wgrad_accumulation, - is_cpu_offload_enabled(), - self.tp_group, - self.tp_size, - self.sequence_parallel, - self.tp_size > 1, - self.activation_dtype, - self.parallel_mode, - is_grad_enabled, - ub_overlap_rs_fprop, - ub_overlap_ag_dgrad, - ub_overlap_ag_fprop, - ub_overlap_rs_dgrad, - ub_bulk_dgrad, - ub_bulk_wgrad, - self.ub_name, - fp8_output, - self.fsdp_group, - cache_name is not None, - skip_fp8_weight_update, - self.symmetric_ar_type, - self.save_original_input, - debug, - backward_override, - custom, - backward_input_needs_gather, - self.is_fsdp2, + linear_bias_tensor = ( + bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None + ) + wgrad_store = self.wgrad_store if self.wgrad_store.delay_wgrad_compute() else None + fwd_args = LinearFwdArgs( + # tensors + weight=weight_tensor, + inp=inp, + bias=linear_bias_tensor, + weight_workspace=weight_workspace, + # requires_grad flags + input_requires_grad=inp.requires_grad, + weight_requires_grad=weight_tensor.requires_grad, + bias_requires_grad=( + linear_bias_tensor.requires_grad if linear_bias_tensor is not None else False + ), + # quantizers + input_quantizer=input_quantizer, + weight_quantizer=weight_quantizer, + output_quantizer=output_quantizer, + grad_input_quantizer=grad_input_quantizer, + grad_weight_quantizer=grad_weight_quantizer, + grad_output_quantizer=grad_output_quantizer, + # numerical / dtype config + activation_dtype=self.activation_dtype, + fp8=self.fp8, + fp8_calibration=self.fp8_calibration, + fp8_output=fp8_output, + save_original_input=self.save_original_input, + backward_override=backward_override, + custom=custom, + debug=debug, + # weight-workspace caching + is_first_microbatch=is_first_microbatch, + cache_weight=cache_name is not None, + skip_fp8_weight_update=skip_fp8_weight_update, + # tensor / sequence parallelism + parallel_mode=self.parallel_mode, + tp_group=self.tp_group, + tp_size=self.tp_size, + tensor_parallel=self.tp_size > 1, + sequence_parallel=self.sequence_parallel, + symmetric_ar_type=self.symmetric_ar_type, + backward_input_needs_gather=backward_input_needs_gather, + # userbuffers + ub_name=self.ub_name, + ub_overlap_ag_fprop=ub_overlap_ag_fprop, + ub_overlap_rs_fprop=ub_overlap_rs_fprop, + ub_overlap_ag_dgrad=ub_overlap_ag_dgrad, + ub_overlap_rs_dgrad=ub_overlap_rs_dgrad, + ub_bulk_dgrad=ub_bulk_dgrad, + ub_bulk_wgrad=ub_bulk_wgrad, + # FSDP + fsdp_group=self.fsdp_group, + is_fsdp2=self.is_fsdp2, + # weight-grad scheduling + fuse_wgrad_accumulation=self.fuse_wgrad_accumulation, + wgrad_store=wgrad_store, + # misc + cpu_offloading=is_cpu_offload_enabled(), + is_grad_enabled=is_grad_enabled, ) out, new_weight_workspace = linear_fn( *autograd_ctx, weight_tensor, - weight_workspace, inp, - bias_tensor if (self.apply_bias and not self.gemm_bias_unfused_add) else None, - non_tensor_args, - input_quantizer, - weight_quantizer, - output_quantizer, - grad_input_quantizer, - grad_weight_quantizer, - grad_output_quantizer, + linear_bias_tensor, + fwd_args, ) if new_weight_workspace is not None and cache_name is not None: From 6cdd7115e1c1e6605feaff6ebe111ea0466cd71a Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Mon, 11 May 2026 21:25:17 -0700 Subject: [PATCH 050/180] [PyTorch] CPU overhead optimizations for te autocast (#2957) * cpu optimizations for te autocast Signed-off-by: Varun Thumbe * address some review comments Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * address review comments Signed-off-by: Varun Thumbe * clean comments Signed-off-by: Varun Thumbe --------- Signed-off-by: Varun Thumbe Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/common/recipe/__init__.py | 70 +++++++++++--- transformer_engine/pytorch/quantization.py | 97 +++++++++++++------- 2 files changed, 118 insertions(+), 49 deletions(-) diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index 959966369..b773a81d1 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -4,6 +4,7 @@ """This module provides predefined FP8 recipes.""" from __future__ import annotations +import abc import os from enum import Enum from typing import Any, Literal, Optional, Union, Callable, NamedTuple @@ -60,6 +61,16 @@ class MMParams: use_split_accumulator: bool = True + def __post_init__(self) -> None: + object.__setattr__( + self, + "_cached_repr", + f"MMParams(use_split_accumulator={self.use_split_accumulator})", + ) + + def __repr__(self) -> str: + return self._cached_repr + @dataclass(frozen=True) class QParams: @@ -76,21 +87,50 @@ class QParams: stochastic_rounding: bool = False fp4_2d_quantization: bool = False - def __repr__(self) -> str: - return ( + def __post_init__(self) -> None: + object.__setattr__( + self, + "_cached_repr", f"Qparams(\npower_2_scale={self.power_2_scale},\n" f"amax_epsilon={self.amax_epsilon},\n" f"random_hadamard_transform={self.random_hadamard_transform},\n" f"stochastic_rounding={self.stochastic_rounding},\n" - f"fp4_2d_quantization={self.fp4_2d_quantization}\n)" + f"fp4_2d_quantization={self.fp4_2d_quantization}\n)", ) + def __repr__(self) -> str: + return self._cached_repr + class Recipe: """ Base recipe class. """ + # Cached string representation. Lazily populated by ``__repr__`` in + # subclasses and invalidated by ``__setattr__`` whenever any attribute + # changes. This makes repeated ``str(recipe)`` calls much cheaper + _cached_repr: Optional[str] = None + + def __setattr__(self, name: str, value: Any) -> None: + # Invalidate the cached repr on any attribute mutation. + if name != "_cached_repr": + object.__setattr__(self, "_cached_repr", None) + object.__setattr__(self, name, value) + + @abc.abstractmethod + def _make_repr(self) -> str: + """Build the string representation for this recipe. + + Subclasses must override this method. The result is cached by + ``__repr__`` and reused until any attribute is mutated. + """ + + def __repr__(self) -> str: + if self._cached_repr is None: + self._cached_repr = self._make_repr() + return self._cached_repr + @classmethod def nvfp4(cls): """Whether the given recipe is NVFP4 1D block scaling.""" @@ -127,7 +167,7 @@ def custom(cls): return issubclass(cls, CustomRecipe) -@dataclass() +@dataclass(repr=False) class DelayedScaling(Recipe): """ Use the delayed scaling factor strategy. Use scale factor from previous @@ -227,7 +267,7 @@ def __post_init__(self) -> None: self.backward_override is None ), "Delayed scaling only supports backward_override=None." - def __repr__(self) -> str: + def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " f"margin={self.margin}, " @@ -240,7 +280,7 @@ def __repr__(self) -> str: ) -@dataclass() +@dataclass(repr=False) class Float8CurrentScaling(Recipe): """ Use the per-tensor current scaling factor strategy. @@ -275,7 +315,7 @@ def __post_init__(self) -> None: self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." - def __repr__(self) -> str: + def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " f"format={str(self.fp8_format).split('.')[1]}, " @@ -291,7 +331,7 @@ def __repr__(self) -> str: ) -@dataclass() +@dataclass(repr=False) class MXFP8BlockScaling(Recipe): """ Use the MXFP8 scaling factor strategy. @@ -333,7 +373,7 @@ def __post_init__(self) -> None: self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." - def __repr__(self) -> str: + def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " f"margin={self.margin}, " @@ -342,7 +382,7 @@ def __repr__(self) -> str: ) -@dataclass() +@dataclass(repr=False) class Float8BlockScaling(Recipe): """ Use block-wise scaling for FP8 tensors. @@ -414,7 +454,7 @@ def __post_init__(self) -> None: self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." - def __repr__(self) -> str: + def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " f"format={str(self.fp8_format).split('.')[1]}, " @@ -433,7 +473,7 @@ def __repr__(self) -> str: ) -@dataclass() +@dataclass(repr=False) class NVFP4BlockScaling(Recipe): """ Use the NVFP4 scaling strategy. @@ -531,7 +571,7 @@ def __post_init__(self) -> None: fp4_2d_quantization=False, ) - def __repr__(self) -> str: + def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " f"fp4_format={str(self.fp4_format).split('.')[1]}, " @@ -546,7 +586,7 @@ def __repr__(self) -> str: ) -@dataclass() +@dataclass(repr=False) class CustomRecipe(Recipe): """ Custom recipe that allows users to provide quantizer factories. @@ -608,7 +648,7 @@ def __post_init__(self) -> None: self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." - def __repr__(self) -> str: + def _make_repr(self) -> str: return ( f"recipe_type={self.__class__.__name__}, " f"qfactory={self.qfactory}, " diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 82b827437..0c4072351 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -686,9 +686,8 @@ def reduce_and_update_fp8_tensors( amax_history, scale, get_fp8_max(recipe, forward), recipe ) - @classmethod + @staticmethod def get_unique_autocast_key( - cls, recipe: Optional[Recipe] = None, group: Optional[dist_group_type] = None, ): @@ -697,7 +696,11 @@ def get_unique_autocast_key( Object identity is sufficient since autocast contexts never outlive a single training session. """ - return str((str(recipe), id(group) if group is not None else None)) + recipe_repr = recipe.__dict__.get("_cached_repr") if recipe is not None else None + if recipe_repr is None: + recipe_repr = str(recipe) + group_id = id(group) if group is not None else None + return f"recipe={recipe_repr},group={group_id}" @classmethod def autocast_enter( @@ -911,14 +914,13 @@ def quantized_model_init( qstate.high_precision_init_val = _high_precision_init_val -@contextmanager def fp8_autocast( enabled: bool = True, calibrating: bool = False, fp8_recipe: Optional[Recipe] = None, fp8_group: Optional[dist_group_type] = None, _graph: bool = False, -) -> None: +) -> "autocast": """ .. warning:: @@ -934,25 +936,16 @@ def fp8_autocast( stacklevel=2, ) - # Call new implementation. - with autocast( + return autocast( enabled=enabled, calibrating=calibrating, recipe=fp8_recipe, amax_reduction_group=fp8_group, _graph=_graph, - ): - yield + ) -@contextmanager -def autocast( - enabled: bool = True, - calibrating: bool = False, - recipe: Optional["Recipe"] = None, - amax_reduction_group: Optional["dist_group_type"] = None, - _graph: bool = False, -) -> None: +class autocast: """ Context manager for quantization schemes like FP8 or FP4. @@ -991,24 +984,60 @@ def autocast( are reduced at the end of each training step. """ - if enabled: - check_recipe_support(recipe) - - # Save current state so we always restore it on exit. - fp8_state = FP8GlobalStateManager.get_autocast_state() - - FP8GlobalStateManager.autocast_enter( - enabled=enabled, - calibrating=calibrating, - fp8_recipe=recipe, - fp8_group=amax_reduction_group, - _graph=_graph, + # Class-based context manager (instead of ``@contextmanager`` from contextlib) + # to avoid overheads. + __slots__ = ( + "_enabled", + "_calibrating", + "_recipe", + "_amax_reduction_group", + "_graph", + "_fp8_state", ) - try: - yield - finally: - FP8GlobalStateManager.set_autocast_state(fp8_state) - FP8GlobalStateManager.autocast_exit(enabled, _graph=_graph) + + def __init__( + self, + enabled: bool = True, + calibrating: bool = False, + recipe: Optional["Recipe"] = None, + amax_reduction_group: Optional["dist_group_type"] = None, + _graph: bool = False, + ) -> None: + self._enabled = enabled + self._calibrating = calibrating + self._recipe = recipe + self._amax_reduction_group = amax_reduction_group + self._graph = _graph + self._fp8_state = None + + def __enter__(self) -> "autocast": + # Disallow nested re-entry of the same instance. + if self._fp8_state is not None: + raise RuntimeError( + "autocast context manager cannot be entered more than once concurrently" + ) + if self._enabled: + check_recipe_support(self._recipe) + # Save current state so we always restore it on exit. + self._fp8_state = FP8GlobalStateManager.get_autocast_state() + FP8GlobalStateManager.autocast_enter( + enabled=self._enabled, + calibrating=self._calibrating, + fp8_recipe=self._recipe, + fp8_group=self._amax_reduction_group, + _graph=self._graph, + ) + return self + + def __exit__(self, exc_type, exc_val, exc_tb) -> None: + try: + FP8GlobalStateManager.set_autocast_state(self._fp8_state) + FP8GlobalStateManager.autocast_exit(self._enabled, _graph=self._graph) + finally: + # Clear the saved state so the instance can be entered again + # sequentially (and so a failure inside the restore path does not + # permanently mark the instance as "active"). + self._fp8_state = None def _update_amax_history(amax_history: torch.Tensor) -> torch.Tensor: From d5e7087db845d4e963bacad7be058b6c6d88c00e Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Tue, 12 May 2026 11:38:43 -0700 Subject: [PATCH 051/180] Disable the RHT fusion for non-SM100 family devices (#2968) * Disable the RHT fusion for non-SM100 family devices Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix the compilation error Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Przemek Tredak Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/pytorch/csrc/quantizer.cpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 2b29f260e..82dfe4d22 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -7,6 +7,7 @@ #include #include "common.h" +#include "common/util/cuda_runtime.h" #include "common/util/system.h" #include "pybind.h" #include "torch/torch.h" @@ -2264,7 +2265,8 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou // Restriction for the RHT cast fusion kernel because we are using MMA hardware for computing RHT bool eligible_for_rht_cast_fusion = - input.dtype() == DType::kBFloat16 && rows % 64 == 0 && cols % 128 == 0; + input.dtype() == DType::kBFloat16 && rows % 64 == 0 && cols % 128 == 0 && + transformer_engine::cuda::sm_arch() >= 100 && transformer_engine::cuda::sm_arch() <= 110; // Stochastic rounding // When both rowwise and columnwise quantization are used with RHT, From cb59ef1567d5a5a0bb97f567ca499397d4762400 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Tue, 12 May 2026 11:41:33 -0700 Subject: [PATCH 052/180] [PyTorch] Expose function to bulk-allocate tensors backed by the same buffer (#2900) * [PyTorch] Add bulk_allocate utility and use it in quantized tensor allocators Introduces transformer_engine/pytorch/csrc/extensions/allocate.cpp with a general-purpose bulk_allocate function: given parallel lists of shapes, dtypes, and per-tensor byte alignments, it computes a packed layout, does a single CUDA allocation, and returns at::from_blob views whose deleters keep the backing buffer alive. The three internal bulk_allocate_*_tensors helpers in cast.cpp are refactored to call bulk_allocate instead of each owning a copy of the make_torch_view lambda and the offset-computation loops (~120 lines removed). The new function is also exposed via pybind11 so Python can allocate packed CUDA buffers directly without going through a quantizer. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Tim Moon * Bulk-allocate wgrads in grouped linear impls Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Apply review suggestions Make optional args for device and alignment. Handle case where base data_ptr is unaligned. Align grouped linear wgrad buffers to 256B. Signed-off-by: Tim Moon * Nits from Claude Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix incorrect call to `bulk_allocate` Signed-off-by: Tim Moon * Fix ambiguous return type Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use c10::Device consistently Signed-off-by: Tim Moon --------- Signed-off-by: Tim Moon Co-authored-by: Claude Sonnet 4.6 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/pytorch/csrc/extensions.h | 10 + .../pytorch/csrc/extensions/allocate.cpp | 86 +++++ .../pytorch/csrc/extensions/cast.cpp | 340 ++++++------------ .../pytorch/csrc/extensions/pybind.cpp | 7 + .../pytorch/module/grouped_linear.py | 11 +- .../pytorch/ops/basic/grouped_linear.py | 10 +- .../pytorch/ops/fused/backward_grouped_mlp.py | 9 +- 7 files changed, 229 insertions(+), 244 deletions(-) create mode 100644 transformer_engine/pytorch/csrc/extensions/allocate.cpp diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 4a2ea7412..9b10a9c5a 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -309,6 +309,16 @@ std::vector rmsnorm_fwd(const py::handle &input, const py::handle &w py::object ln_out, py::handle quantizer, DType otype, const int sm_margin, const bool zero_centered_gamma); +/*************************************************************************************************** + * Memory allocation + **************************************************************************************************/ + +// Allocates tensors all backed by a single contiguous buffer. +std::vector bulk_allocate(const std::vector> &shapes, + const std::vector &dtypes, + std::optional device = std::nullopt, + std::optional> alignments = std::nullopt); + /*************************************************************************************************** * Cast **************************************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/allocate.cpp b/transformer_engine/pytorch/csrc/extensions/allocate.cpp new file mode 100644 index 000000000..f972f8a2d --- /dev/null +++ b/transformer_engine/pytorch/csrc/extensions/allocate.cpp @@ -0,0 +1,86 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include + +#include "../extensions.h" + +namespace transformer_engine { +namespace pytorch { + +std::vector bulk_allocate(const std::vector> &shapes, + const std::vector &dtypes, + std::optional device, + std::optional> alignments) { + // Check shapes and dtypes + const size_t n = shapes.size(); + NVTE_CHECK(dtypes.size() == n, "Got ", shapes.size(), " shapes and ", dtypes.size(), " dtypes."); + NVTE_CHECK(!alignments || alignments->size() == n, "Got ", shapes.size(), " shapes and ", + alignments->size(), " alignments."); + + // Return immediately if no tensors are needed + if (n == 0) return {}; + + // Set defaults for optional arguments + if (!device) { + device = c10::Device(c10::kCUDA); + } + if (!alignments) { + alignments = std::vector{}; + alignments->reserve(n); + for (const auto &dtype : dtypes) { + alignments->push_back(c10::elementSize(dtype)); + } + } + + // Compute offsets in base buffer + std::vector byte_sizes(n); + std::vector offsets(n); + size_t base_byte_size = 0; + size_t base_alignment = 1; + for (size_t i = 0; i < n; ++i) { + byte_sizes[i] = product(shapes[i]) * at::elementSize(dtypes[i]); + offsets[i] = roundup(base_byte_size, (*alignments)[i]); + base_byte_size = offsets[i] + byte_sizes[i]; + base_alignment = std::max(base_alignment, (*alignments)[i]); + } + if (base_alignment > 1) { + // Pad in case data pointer is not aligned + base_byte_size += base_alignment; + } + + // Allocate base buffer + auto base_buffer = std::make_shared( + at::empty({static_cast(base_byte_size)}, at::device(*device).dtype(torch::kUInt8))); + uint8_t *base_ptr = base_buffer->data_ptr(); + base_ptr = + reinterpret_cast(roundup(reinterpret_cast(base_ptr), base_alignment)); + + // Create views into base buffer + std::vector out; + out.reserve(n); + std::vector shape_int64; + for (size_t i = 0; i < n; ++i) { + shape_int64.assign(shapes[i].begin(), shapes[i].end()); + if (byte_sizes[i] == 0) { + // Work around problems with from_blob when constructing an + // empty tensor. Passing a null pointer fails because it checks + // that the pointer is on GPU. Passing a non-null pointer can + // cause bugs in TE kernels. + out.emplace_back(at::empty(shape_int64, at::device(*device).dtype(dtypes[i]))); + } else { + // Construct tensor with custom deleter to keep base buffer alive + out.emplace_back(at::from_blob( + base_ptr + offsets[i], shape_int64, [base_buffer](void *) {}, + at::device(*device).dtype(dtypes[i]))); + } + } + return out; +} + +} // namespace pytorch +} // namespace transformer_engine diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 00f4383ab..3ada2459c 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -495,60 +495,30 @@ std::tuple, std::vector> bulk_allocate_fp const auto scaling_mode = quantizer_cpp_list[0]->get_scaling_mode(); const auto is_2D_scaled = scaling_mode == NVTE_BLOCK_SCALING_2D; const auto fp8_dtype = quantizer_cpp_list[0]->dtype; - constexpr size_t fp8_elem_size = 1; - constexpr size_t scale_elem_size = 4; - - // Helper function to construct tensor view - // Note: Deleter holds a shared_ptr for the buffer, so the buffer - // will survive until all views are deleted. - auto make_torch_view = [](std::shared_ptr &buffer, const std::vector &shape, - size_t offset, at::ScalarType dtype) -> at::Tensor { - std::vector shape_int64(shape.begin(), shape.end()); - bool is_empty_shape = product(shape) == 0; - if (buffer->data_ptr() == nullptr || is_empty_shape) { - return at::empty(shape_int64, at::device(at::kCUDA).dtype(dtype)); - } - return at::from_blob( - buffer->data_ptr() + offset, shape_int64, - [buffer](void *) {}, // deleter holds shared_ptr - at::device(at::kCUDA).dtype(dtype)); - }; // Allocate row-wise data std::vector rowwise_data_list, rowwise_scale_list; std::vector> rowwise_data_shapes, rowwise_scale_shapes; if (rowwise_usage) { - // Tensor sizes for (size_t i = 0; i < num_tensors; ++i) { rowwise_data_shapes.emplace_back(shape_list[i]); rowwise_scale_shapes.emplace_back( quantizer_cpp_list[i]->get_scale_shape(shape_list[i], false)); } - // Offsets in full buffer - size_t buffer_size = 0; - std::vector data_offsets, scale_offsets; - for (size_t i = 0; i < num_tensors; ++i) { - buffer_size = roundup(buffer_size, 256); // align to 256B - data_offsets.push_back(buffer_size); - buffer_size += product(rowwise_data_shapes[i]) * fp8_elem_size; - } - for (size_t i = 0; i < num_tensors; ++i) { - buffer_size = roundup(buffer_size, 16); // align to 16B - scale_offsets.push_back(buffer_size); - buffer_size += product(rowwise_scale_shapes[i]) * scale_elem_size; - } + // Bulk-allocate data and scale tensors + std::vector> shapes = rowwise_data_shapes; + std::vector dtypes(num_tensors, torch::kUInt8); + std::vector alignments(num_tensors, 256); + shapes.insert(shapes.end(), rowwise_scale_shapes.begin(), rowwise_scale_shapes.end()); + dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32); + alignments.insert(alignments.end(), num_tensors, 16); + auto tensors = bulk_allocate(shapes, dtypes, std::nullopt, alignments); - // Allocate full buffer - auto buffer = std::make_shared( - at::empty({(int64_t)buffer_size}, at::device(at::kCUDA).dtype(torch::kUInt8))); - - // Construct tensor views + // Split data and scale tensors for (size_t i = 0; i < num_tensors; ++i) { - rowwise_data_list.emplace_back( - make_torch_view(buffer, rowwise_data_shapes[i], data_offsets[i], torch::kUInt8)); - rowwise_scale_list.emplace_back( - make_torch_view(buffer, rowwise_scale_shapes[i], scale_offsets[i], torch::kFloat32)); + rowwise_data_list.emplace_back(std::move(tensors[i])); + rowwise_scale_list.emplace_back(std::move(tensors[num_tensors + i])); } } @@ -556,7 +526,6 @@ std::tuple, std::vector> bulk_allocate_fp std::vector columnwise_data_list, columnwise_scale_list; std::vector> columnwise_data_shapes, columnwise_scale_shapes; if (columnwise_usage) { - // Tensor sizes for (size_t i = 0; i < num_tensors; ++i) { columnwise_data_shapes.emplace_back(); auto &shape = columnwise_data_shapes.back(); @@ -568,30 +537,19 @@ std::tuple, std::vector> bulk_allocate_fp quantizer_cpp_list[i]->get_scale_shape(shape_list[i], true)); } - // Offsets in full buffer - size_t buffer_size = 0; - std::vector data_offsets, scale_offsets; - for (size_t i = 0; i < num_tensors; ++i) { - buffer_size = roundup(buffer_size, 256); // align to 256B - data_offsets.push_back(buffer_size); - buffer_size += product(columnwise_data_shapes[i]) * fp8_elem_size; - } - for (size_t i = 0; i < num_tensors; ++i) { - buffer_size = roundup(buffer_size, 16); // align to 16B - scale_offsets.push_back(buffer_size); - buffer_size += product(columnwise_scale_shapes[i]) * scale_elem_size; - } - - // Allocate full buffer - auto buffer = std::make_shared( - at::empty({(int64_t)buffer_size}, at::device(at::kCUDA).dtype(torch::kUInt8))); + // Bulk-allocate data and scale tensors + std::vector> shapes = columnwise_data_shapes; + std::vector dtypes(num_tensors, torch::kUInt8); + std::vector alignments(num_tensors, 256); + shapes.insert(shapes.end(), columnwise_scale_shapes.begin(), columnwise_scale_shapes.end()); + dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32); + alignments.insert(alignments.end(), num_tensors, 16); + auto tensors = bulk_allocate(shapes, dtypes, std::nullopt, alignments); - // Construct tensor views + // Split data and scale tensors for (size_t i = 0; i < num_tensors; ++i) { - columnwise_data_list.emplace_back( - make_torch_view(buffer, columnwise_data_shapes[i], data_offsets[i], torch::kUInt8)); - columnwise_scale_list.emplace_back( - make_torch_view(buffer, columnwise_scale_shapes[i], scale_offsets[i], torch::kFloat32)); + columnwise_data_list.push_back(tensors[i]); + columnwise_scale_list.push_back(tensors[num_tensors + i]); } } @@ -648,60 +606,29 @@ std::tuple, std::vector> bulk_allocate_mx const auto fp8_dtype = quantizer_cpp_list[0]->dtype; const bool with_gemm_swizzled_scales = quantizer_cpp_list[0]->optimize_for_gemm; - constexpr size_t fp8_elem_size = 1; - constexpr size_t scale_elem_size = 1; - - // Helper function to construct tensor view - // Note: Deleter holds a shared_ptr for the buffer, so the buffer - // will survive until all views are deleted. - auto make_torch_view = [](std::shared_ptr &buffer, const std::vector &shape, - size_t offset, at::ScalarType dtype) -> at::Tensor { - std::vector shape_int64(shape.begin(), shape.end()); - bool is_empty_shape = product(shape) == 0; - if (buffer->data_ptr() == nullptr || is_empty_shape) { - return at::empty(shape_int64, at::device(at::kCUDA).dtype(dtype)); - } - return at::from_blob( - buffer->data_ptr() + offset, shape_int64, - [buffer](void *) {}, // deleter holds shared_ptr - at::device(at::kCUDA).dtype(dtype)); - }; - // Allocate row-wise data std::vector rowwise_data_list, rowwise_scale_list; std::vector> rowwise_data_shapes, rowwise_scale_shapes; if (rowwise_usage) { - // Tensor sizes for (size_t i = 0; i < num_tensors; ++i) { rowwise_data_shapes.emplace_back(shape_list[i]); rowwise_scale_shapes.emplace_back( quantizer_cpp_list[i]->get_scale_shape(shape_list[i], false)); } - // Offsets in full buffer - size_t buffer_size = 0; - std::vector data_offsets, scale_offsets; - for (size_t i = 0; i < num_tensors; ++i) { - buffer_size = roundup(buffer_size, 256); // align to 256B - data_offsets.push_back(buffer_size); - buffer_size += product(rowwise_data_shapes[i]) * fp8_elem_size; - } - for (size_t i = 0; i < num_tensors; ++i) { - buffer_size = roundup(buffer_size, 16); // align to 16B - scale_offsets.push_back(buffer_size); - buffer_size += product(rowwise_scale_shapes[i]) * scale_elem_size; - } - - // Allocate full buffer - auto buffer = std::make_shared( - at::empty({(int64_t)buffer_size}, at::device(at::kCUDA).dtype(torch::kUInt8))); + // Bulk-allocate data and scale tensors + std::vector> shapes = rowwise_data_shapes; + std::vector dtypes(num_tensors, torch::kUInt8); + std::vector alignments(num_tensors, 256); + shapes.insert(shapes.end(), rowwise_scale_shapes.begin(), rowwise_scale_shapes.end()); + dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8); + alignments.insert(alignments.end(), num_tensors, 16); + auto tensors = bulk_allocate(shapes, dtypes, std::nullopt, alignments); - // Construct tensor views + // Split data and scale tensors for (size_t i = 0; i < num_tensors; ++i) { - rowwise_data_list.emplace_back( - make_torch_view(buffer, rowwise_data_shapes[i], data_offsets[i], torch::kUInt8)); - rowwise_scale_list.emplace_back( - make_torch_view(buffer, rowwise_scale_shapes[i], scale_offsets[i], torch::kUInt8)); + rowwise_data_list.emplace_back(std::move(tensors[i])); + rowwise_scale_list.emplace_back(std::move(tensors[num_tensors + i])); } } @@ -709,7 +636,6 @@ std::tuple, std::vector> bulk_allocate_mx std::vector columnwise_data_list, columnwise_scale_list; std::vector> columnwise_data_shapes, columnwise_scale_shapes; if (columnwise_usage) { - // Tensor sizes for (size_t i = 0; i < num_tensors; ++i) { // For MXFP8, the columnwise data doesn't need transpose // because of TN, NT, NN layout support in SM100 @@ -718,30 +644,19 @@ std::tuple, std::vector> bulk_allocate_mx quantizer_cpp_list[i]->get_scale_shape(shape_list[i], true)); } - // Offsets in full buffer - size_t buffer_size = 0; - std::vector data_offsets, scale_offsets; - for (size_t i = 0; i < num_tensors; ++i) { - buffer_size = roundup(buffer_size, 256); // align to 256B - data_offsets.push_back(buffer_size); - buffer_size += product(columnwise_data_shapes[i]) * fp8_elem_size; - } - for (size_t i = 0; i < num_tensors; ++i) { - buffer_size = roundup(buffer_size, 16); // align to 16B - scale_offsets.push_back(buffer_size); - buffer_size += product(columnwise_scale_shapes[i]) * scale_elem_size; - } - - // Allocate full buffer - auto buffer = std::make_shared( - at::empty({(int64_t)buffer_size}, at::device(at::kCUDA).dtype(torch::kUInt8))); + // Bulk-allocate data and scale tensors + std::vector> shapes = columnwise_data_shapes; + std::vector dtypes(num_tensors, torch::kUInt8); + std::vector alignments(num_tensors, 256); + shapes.insert(shapes.end(), columnwise_scale_shapes.begin(), columnwise_scale_shapes.end()); + dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8); + alignments.insert(alignments.end(), num_tensors, 16); + auto tensors = bulk_allocate(shapes, dtypes, std::nullopt, alignments); - // Construct tensor views + // Split data and scale tensors for (size_t i = 0; i < num_tensors; ++i) { - columnwise_data_list.emplace_back( - make_torch_view(buffer, columnwise_data_shapes[i], data_offsets[i], torch::kUInt8)); - columnwise_scale_list.emplace_back( - make_torch_view(buffer, columnwise_scale_shapes[i], scale_offsets[i], torch::kUInt8)); + columnwise_data_list.push_back(tensors[i]); + columnwise_scale_list.push_back(tensors[num_tensors + i]); } } @@ -808,103 +723,70 @@ std::tuple, std::vector, bool> bulk_alloc const auto scaling_mode = quantizer_cpp_list[0]->get_scaling_mode(); const auto fp4_dtype = quantizer_cpp_list[0]->dtype; const bool with_gemm_swizzled_scales = false; /// TODO (tmoon) Enable based on optimize_for_gemm; - constexpr size_t scale_elem_size = 1; - - // Helper function to construct tensor view - // Note: Deleter holds a shared_ptr for the buffer, so the buffer - // will survive until all views are deleted. - auto make_torch_view = [](std::shared_ptr &buffer, const std::vector &shape, - size_t offset, at::ScalarType dtype) -> at::Tensor { - std::vector shape_int64(shape.begin(), shape.end()); - bool is_empty_shape = product(shape) == 0; - if (buffer->data_ptr() == nullptr || is_empty_shape) { - return at::empty(shape_int64, at::device(at::kCUDA).dtype(dtype)); - } - return at::from_blob( - buffer->data_ptr() + offset, shape_int64, - [buffer](void *) {}, // deleter holds shared_ptr - at::device(at::kCUDA).dtype(dtype)); - }; - // Lambda function for converting std::vector shape to NVFP4 shape (last dim divided by 2) - auto to_fp4_shape = [](const std::vector &shape) { - std::vector fp4_shape(shape.begin(), shape.end()); - if (!fp4_shape.empty()) { - fp4_shape.back() /= 2; - } - return fp4_shape; + // Helper function to get size of byte buffer holding FP4 data (last dim divided by 2) + auto fp4_byte_shape = [](const std::vector &shape) -> std::vector { + NVTE_CHECK(!shape.empty()); + NVTE_CHECK(shape.back() % 2 == 0); + std::vector out(shape.begin(), shape.end()); + out.back() /= 2; + return out; }; - auto flat_first_dim = [](const std::vector &shape) -> size_t { - if (shape.empty()) { - return 1; - } - size_t rows = 1; - for (size_t i = 0; i + 1 < shape.size(); ++i) { - rows *= shape[i]; + + // Helper function to get size of amax buffer + auto amax_shape = [](const std::vector &shape, + bool row_scaled = false) -> std::vector { + if (row_scaled) { + const auto [rows, _] = get_2d_dims(shape); + return {rows}; } - return rows; + return {1}; }; // Allocate row-wise data std::vector rowwise_data_list, rowwise_scale_list, amax_rowwise_list; std::vector> rowwise_data_shapes, rowwise_scale_shapes; if (rowwise_usage) { - // Tensor sizes for (size_t i = 0; i < num_tensors; ++i) { rowwise_data_shapes.emplace_back(shape_list[i]); rowwise_scale_shapes.emplace_back( quantizer_cpp_list[i]->get_scale_shape(shape_list[i], false)); } - // Offsets in full buffer - size_t buffer_size = 0; - std::vector data_offsets, scale_offsets, amax_offsets; + // Check whether data and scales can be packed in contiguous + // buffer. Amaxes are not contiguous since they are aligned to + // 16B. for (size_t i = 0; i < num_tensors; ++i) { - // FP4 data is aligned to 256B - const auto offset = roundup(buffer_size, 256); - if (offset != buffer_size) { + if (product(rowwise_data_shapes[i]) / 2 % 256 != 0) { contiguous_data_and_scale = false; } - data_offsets.push_back(offset); - buffer_size = offset + (product(rowwise_data_shapes[i]) + 1) / 2; - } - for (size_t i = 0; i < num_tensors; ++i) { - // Scales are aligned to 16B - const auto offset = roundup(buffer_size, 16); - if (offset != buffer_size) { + if (product(rowwise_scale_shapes[i]) % 16 != 0) { contiguous_data_and_scale = false; } - scale_offsets.push_back(offset); - buffer_size = offset + product(rowwise_scale_shapes[i]) * scale_elem_size; } + + // Bulk-allocate tensors data, scale, and amax tensors + std::vector> shapes; + for (size_t i = 0; i < num_tensors; ++i) { + shapes.emplace_back(fp4_byte_shape(rowwise_data_shapes[i])); + } + std::vector dtypes(num_tensors, torch::kUInt8); + std::vector alignments(num_tensors, 256); + shapes.insert(shapes.end(), rowwise_scale_shapes.begin(), rowwise_scale_shapes.end()); + dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8); + alignments.insert(alignments.end(), num_tensors, 16); for (size_t i = 0; i < num_tensors; ++i) { - // Amaxes (FP32) are aligned to 16B - // Note: Multi-quantize kernel does not require contiguous amaxes. - const auto offset = roundup(buffer_size, 16); - amax_offsets.push_back(offset); - size_t amax_size = 4; - if (row_scaled_nvfp4) { - amax_size *= flat_first_dim(rowwise_data_shapes[i]); - } - buffer_size = offset + amax_size; + shapes.emplace_back(amax_shape(rowwise_data_shapes[i], row_scaled_nvfp4)); } + dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32); + alignments.insert(alignments.end(), num_tensors, 16); + auto tensors = bulk_allocate(shapes, dtypes, std::nullopt, alignments); - // Allocate full buffer - auto buffer = std::make_shared( - at::empty({(int64_t)buffer_size}, at::device(at::kCUDA).dtype(torch::kUInt8))); - - // Construct tensor views + // Split data, scale, and amax tensors for (size_t i = 0; i < num_tensors; ++i) { - rowwise_data_list.emplace_back(make_torch_view(buffer, to_fp4_shape(rowwise_data_shapes[i]), - data_offsets[i], torch::kUInt8)); - rowwise_scale_list.emplace_back( - make_torch_view(buffer, rowwise_scale_shapes[i], scale_offsets[i], torch::kUInt8)); - std::vector amax_shape{1}; - if (row_scaled_nvfp4) { - amax_shape = {flat_first_dim(rowwise_data_shapes[i])}; - } - amax_rowwise_list.emplace_back( - make_torch_view(buffer, amax_shape, amax_offsets[i], torch::kFloat32)); + rowwise_data_list.push_back(tensors[i]); + rowwise_scale_list.push_back(tensors[num_tensors + i]); + amax_rowwise_list.push_back(tensors[2 * num_tensors + i]); } } @@ -912,7 +794,6 @@ std::tuple, std::vector, bool> bulk_alloc std::vector columnwise_data_list, columnwise_scale_list, amax_columnwise_list; std::vector> columnwise_data_shapes, columnwise_scale_shapes; if (columnwise_usage) { - // Tensor sizes for (size_t i = 0; i < num_tensors; ++i) { // push the transposed shape into NVFP4 columnwise shape // NVFP4 on SM100 is TN only @@ -926,47 +807,40 @@ std::tuple, std::vector, bool> bulk_alloc quantizer_cpp_list[i]->get_scale_shape(shape_list[i], true)); } - // Offsets in full buffer - size_t buffer_size = 0; - std::vector data_offsets, scale_offsets, amax_offsets; + // Check whether data and scales can be packed in contiguous + // buffer. Amaxes are not contiguous since they are aligned to + // 16B. for (size_t i = 0; i < num_tensors; ++i) { - // FP4 data is aligned to 256B - const auto offset = roundup(buffer_size, 256); - if (offset != buffer_size) { + if (product(columnwise_data_shapes[i]) / 2 % 256 != 0) { contiguous_data_and_scale = false; } - data_offsets.push_back(offset); - buffer_size = offset + (product(columnwise_data_shapes[i]) + 1) / 2; - } - for (size_t i = 0; i < num_tensors; ++i) { - // Scales are aligned to 16B - const auto offset = roundup(buffer_size, 16); - if (offset != buffer_size) { + if (product(columnwise_scale_shapes[i]) % 16 != 0) { contiguous_data_and_scale = false; } - scale_offsets.push_back(offset); - buffer_size = offset + product(columnwise_scale_shapes[i]) * scale_elem_size; } + + // Bulk-allocate tensors data, scale, and amax tensors + std::vector> shapes; for (size_t i = 0; i < num_tensors; ++i) { - // Amaxes (FP32) are aligned to 16B - // Note: Multi-quantize kernel does not require contiguous amaxes. - const auto offset = roundup(buffer_size, 16); - amax_offsets.push_back(offset); - buffer_size = offset + 4; + shapes.emplace_back(fp4_byte_shape(columnwise_data_shapes[i])); } + std::vector dtypes(num_tensors, torch::kUInt8); + std::vector alignments(num_tensors, 256); + shapes.insert(shapes.end(), columnwise_scale_shapes.begin(), columnwise_scale_shapes.end()); + dtypes.insert(dtypes.end(), num_tensors, torch::kUInt8); + alignments.insert(alignments.end(), num_tensors, 16); + for (size_t i = 0; i < num_tensors; ++i) { + shapes.emplace_back(amax_shape(columnwise_data_shapes[i])); + } + dtypes.insert(dtypes.end(), num_tensors, torch::kFloat32); + alignments.insert(alignments.end(), num_tensors, 16); + auto tensors = bulk_allocate(shapes, dtypes, std::nullopt, alignments); - // Allocate full buffer - auto buffer = std::make_shared( - at::empty({(int64_t)buffer_size}, at::device(at::kCUDA).dtype(torch::kUInt8))); - - // Construct tensor views + // Split data, scale, and amax tensors for (size_t i = 0; i < num_tensors; ++i) { - columnwise_data_list.emplace_back(make_torch_view( - buffer, to_fp4_shape(columnwise_data_shapes[i]), data_offsets[i], torch::kUInt8)); - columnwise_scale_list.emplace_back( - make_torch_view(buffer, columnwise_scale_shapes[i], scale_offsets[i], torch::kUInt8)); - amax_columnwise_list.emplace_back( - make_torch_view(buffer, std::vector{1}, amax_offsets[i], torch::kFloat32)); + columnwise_data_list.push_back(tensors[i]); + columnwise_scale_list.push_back(tensors[num_tensors + i]); + amax_columnwise_list.push_back(tensors[2 * num_tensors + i]); } } diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index eb7576d90..a813f3119 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -352,6 +352,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Partial cast from master weights for fp8 block scaling", py::arg("inp"), py::arg("out"), py::arg("scale"), py::arg("h"), py::arg("w"), py::arg("start_offset"), py::arg("block_len"), py::arg("out_dtype"), py::call_guard()); + // NVFP4 2D m.def("nvfp4_2d_compute_partial_amax", &transformer_engine::pytorch::nvfp4_2d_compute_partial_amax, @@ -404,6 +405,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "In-place swizzle of grouped tensor scales for GEMM", py::arg("tensor"), py::arg("rowwise"), py::arg("columnwise")); + // Tensor allocation + m.def("bulk_allocate", &transformer_engine::pytorch::bulk_allocate, + "Allocate tensors backed by a single contiguous buffer", py::arg("shapes"), + py::arg("dtypes"), py::arg("device") = py::none(), py::arg("alignments") = py::none(), + py::call_guard()); + // attention kernels m.def("fa_prepare_fwd", &transformer_engine::pytorch::fa_prepare_fwd, "Prepare QKV for Flash Attention", py::call_guard()); diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index e950f2657..627144345 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -496,10 +496,13 @@ def backward( if ctx.fuse_wgrad_accumulation: wgrad_list = main_grads else: - wgrad_list = [ - torch.empty(w.size(), dtype=ctx.activation_dtype, device=ctx.device) - for w in weights - ] + weight_shape = list(weights[0].size()) + wgrad_list = tex.bulk_allocate( + [weight_shape] * ctx.num_gemms, + [ctx.activation_dtype] * ctx.num_gemms, + ctx.device, + [256] * ctx.num_gemms, # alignment + ) if ctx.save_original_input: inp = inputmats[0] diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index e698c2697..1f00d9228 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -1393,10 +1393,12 @@ def _fuser_backward_split_quantize( ] accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: - grad_weights = [ - torch.empty(weight_shape, dtype=ctx.dtype, device=device) - for _ in range(num_groups) - ] + grad_weights = tex.bulk_allocate( + [weight_shape] * num_groups, + [ctx.dtype] * num_groups, + device, + [256] * num_groups, # alignment + ) final_weight_grads = list(grad_weights) # Perform dgrad GEMMs diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 320c7c39e..a11d0505c 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -194,9 +194,12 @@ def _compute_grad_params( w_list = [get_main_grad_from_param(w, op_label=op_label) for w in weights] accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: - w_list = [ - torch.empty(weight_shape, dtype=dtype, device=device) for _ in range(num_groups) - ] + w_list = tex.bulk_allocate( + [weight_shape] * num_groups, + [dtype] * num_groups, + device, + [256] * num_groups, # alignment + ) wgrad_output = w_list if ctx.weight_requires_grad: From c3fd8f88a54b824462f0c25af86019e0bd021350 Mon Sep 17 00:00:00 2001 From: zhujian Date: Wed, 13 May 2026 02:48:11 +0800 Subject: [PATCH 053/180] fix(CP, FA): the conditional logic in the FA version contains a vulnerability when processing the output of Flash Attn forward pass (#2825) fix(CP, FA): when processing the output of Flash Attn forward pass, the conditional logic in the FA version contains a vulnerability, fix it Signed-off-by: zhujian --- .../dot_product_attention/context_parallel.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 995ecf31b..3db0417bd 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -1064,11 +1064,10 @@ def cp_p2p_fwd_flash_attn( **fa_forward_kwargs, ) rng_states = None - if not fa_utils.v2_7_0_plus: + if not use_flash_attn_3 and not fa_utils.v2_7_0_plus: out_per_step = fa_outputs[4] softmax_lse_per_step = fa_outputs[5] - if not use_flash_attn_3: - rng_states = fa_outputs[7] + rng_states = fa_outputs[7] else: out_per_step = fa_outputs[0] softmax_lse_per_step = fa_outputs[1] @@ -3255,11 +3254,10 @@ def forward( causal=causal, **fa_forward_kwargs, ) - if not fa_utils.v2_7_0_plus: + if not use_flash_attn_3 and not fa_utils.v2_7_0_plus: out_per_step[i] = fa_outputs[4] softmax_lse_per_step[i] = fa_outputs[5] - if not use_flash_attn_3: - rng_states[i] = fa_outputs[7] + rng_states[i] = fa_outputs[7] else: out_per_step[i] = fa_outputs[0] softmax_lse_per_step[i] = fa_outputs[1] @@ -4086,9 +4084,9 @@ def forward( causal=causal, **fa_forward_kwargs, ) - if not fa_utils.v2_7_0_plus: + if not use_flash_attn_3 and not fa_utils.v2_7_0_plus: out_, softmax_lse = fa_outputs[4], fa_outputs[5] - rng_state = fa_outputs[7] if not use_flash_attn_3 else None + rng_state = fa_outputs[7] else: out_, softmax_lse = fa_outputs[0], fa_outputs[1] rng_state = fa_outputs[3] if not use_flash_attn_3 else None From 1800fe374f87120aa01d86f42e4712a7c8d3ec44 Mon Sep 17 00:00:00 2001 From: Oleg Goncharov <64355998+Oleg-Goncharov@users.noreply.github.com> Date: Tue, 12 May 2026 20:48:48 +0200 Subject: [PATCH 054/180] [Common] Use specialized unfused MXFP8 cast kernels by default (#2958) * Use fast unfused cast mxfp8 kernels by default Signed-off-by: Oleg Goncharov * Removed dead code Signed-off-by: Oleg Goncharov * Use fast kernel for full 32-element chunks only Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Signed-off-by: Oleg Goncharov * Fixed grid size overflow Signed-off-by: Oleg Goncharov * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Oleg Goncharov Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Przemyslaw Tredak --- .../common/cast/mxfp8/quantize_mxfp8.cuh | 40 +++++++++++++------ .../cast/mxfp8/specialized/quantize_mxfp8.cuh | 20 ++-------- 2 files changed, 31 insertions(+), 29 deletions(-) diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index a0ae7dde8..1549a292d 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -643,15 +643,35 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, TRANSFORMER_ENGINE_SWITCH_CONDITION( with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, + // The specialized rowwise cast-only kernel vectorizes full 128-element chunks. + // Shapes with a partial row tail (for example, N=48) must use the generic kernel, + // otherwise the last chunk reads/writes past the logical end of the row. + using rowwise_traits = specialized::CastTraits; + using bidimensional_traits = specialized::CastTraits; + constexpr size_t max_grid_dim_y = 65535; + const bool rowwise_specialized_grid_fits = + ((rows + rowwise_traits::blockDimM - 1) / rowwise_traits::blockDimM) <= + max_grid_dim_y; + const bool bidimensional_specialized_grid_fits = + ((rows + bidimensional_traits::blockDIM::M - 1) / + bidimensional_traits::blockDIM::M) <= max_grid_dim_y; + + const bool is_full_rowwise_chunk = (cols % 128 == 0); + const bool scaling_type_has_specialized_support = + (scaling_type == ScalingType::ROWWISE && is_full_rowwise_chunk && + rowwise_specialized_grid_fits) || + (scaling_type == ScalingType::BIDIMENSIONAL && + bidimensional_specialized_grid_fits); + if (specialized::hasSpec() && - !WITH_GEMM_SWIZZLED_SCALES) { + !WITH_GEMM_SWIZZLED_SCALES && scaling_type_has_specialized_support) { switch (scaling_type) { case ScalingType::ROWWISE: { using traits = specialized::CastTraits; auto kernel = specialized::quantize_mxfp8_kernel_cast_only; - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - traits::smem); + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, traits::smem)); dim3 block(traits::threadLayout::num, traits::warpLayout::N, traits::warpLayout::M); @@ -664,16 +684,12 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, break; } - case ScalingType::COLWISE: { - NVTE_WARN("Colwise scaling will fallback to original kernel."); - break; - } case ScalingType::BIDIMENSIONAL: { using traits = specialized::CastTraits; auto kernel = specialized::quantize_mxfp8_kernel_cast_only; - cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, - traits::smem); + NVTE_CHECK_CUDA(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, traits::smem)); // TMA for loading, so that we don't need STS for transposing alignas(64) CUtensorMap tensor_map_input{}; constexpr size_t input_type_bit_size = TypeInfo::size; @@ -710,6 +726,7 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, NVTE_ERROR("Invalid scaling type."); } } + NVTE_CHECK_CUDA(cudaGetLastError()); return; } @@ -789,7 +806,6 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, scale_stride_colwise); - NVTE_CHECK_CUDA(cudaGetLastError()); break; } case ScalingType::COLWISE: { @@ -804,7 +820,6 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, scale_stride_colwise); - NVTE_CHECK_CUDA(cudaGetLastError()); break; } case ScalingType::BIDIMENSIONAL: { @@ -819,10 +834,9 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, tensor_map_output_colwise, scales_rowwise_ptr, scales_colwise_ptr, noop_ptr, workspace_ptr, amax_ptr, rows, cols, scale_stride_rowwise, scale_stride_colwise); - NVTE_CHECK_CUDA(cudaGetLastError()); break; } - } + } NVTE_CHECK_CUDA(cudaGetLastError()); if constexpr (IS_DBIAS) { common::reduce_dbias(workspace_ptr, dbias, dbias_rows, dbias_cols, stream); diff --git a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh index 41e62ac31..9459f0273 100644 --- a/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/specialized/quantize_mxfp8.cuh @@ -91,18 +91,6 @@ __device__ __forceinline__ e8m0_t to_e8m0(IType amax) { #endif // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) } // anonymous namespace -inline bool is_cast_only_enabled() { - static bool enabled = []() { - const char *env = std::getenv("ENABLE_CAST_ONLY"); - return env != nullptr && (env[0] == '1'); - }(); - return enabled; - - // // FIXME: when finish debugging, remove this - // const char* env = std::getenv("ENABLE_CAST_ONLY"); - // return env != nullptr && (env[0] == '1'); -} - template inline bool hasSpec() { return false; @@ -112,19 +100,19 @@ inline bool hasSpec() { // OType could be [fp8e5m2, fp8e4m3] template <> inline bool hasSpec() { - return is_cast_only_enabled(); + return true; } template <> inline bool hasSpec() { - return is_cast_only_enabled(); + return true; } template <> inline bool hasSpec() { - return is_cast_only_enabled(); + return true; } template <> inline bool hasSpec() { - return is_cast_only_enabled(); + return true; } template From f0ab81d95038c12f50b62fe32f66862e2a1e6a4f Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Tue, 12 May 2026 13:56:17 -0700 Subject: [PATCH 055/180] Build Docs fix (#2982) * Build Docs fix Signed-off-by: Varun Thumbe * fix doxygen warnings Signed-off-by: Varun Thumbe * class doc fix Signed-off-by: Varun Thumbe --------- Signed-off-by: Varun Thumbe --- docs/Doxyfile | 78 --------------- docs/api/pytorch.rst | 2 +- .../common/include/transformer_engine/gemm.h | 95 +++++++++++-------- .../transformer_engine/transformer_engine.h | 8 +- 4 files changed, 63 insertions(+), 120 deletions(-) diff --git a/docs/Doxyfile b/docs/Doxyfile index f17ffc297..2c593e459 100644 --- a/docs/Doxyfile +++ b/docs/Doxyfile @@ -93,14 +93,6 @@ ALLOW_UNICODE_NAMES = NO OUTPUT_LANGUAGE = English -# The OUTPUT_TEXT_DIRECTION tag is used to specify the direction in which all -# documentation generated by doxygen is written. Doxygen will use this -# information to generate all generated output in the proper direction. -# Possible values are: None, LTR, RTL and Context. -# The default value is: None. - -OUTPUT_TEXT_DIRECTION = None - # If the BRIEF_MEMBER_DESC tag is set to YES, doxygen will include brief member # descriptions after the members that are listed in the file and class # documentation (similar to Javadoc). Set to NO to disable this. @@ -263,12 +255,6 @@ TAB_SIZE = 2 ALIASES = -# This tag can be used to specify a number of word-keyword mappings (TCL only). -# A mapping has the form "name=value". For example adding "class=itcl::class" -# will allow you to use the command class in the itcl::class meaning. - -TCL_SUBST = - # Set the OPTIMIZE_OUTPUT_FOR_C tag to YES if your project consists of C sources # only. Doxygen will then generate output that is more tailored for C. For # instance, some of the names that are used will be different. The list of all @@ -1156,13 +1142,6 @@ CLANG_DATABASE_PATH = ALPHABETICAL_INDEX = YES -# The COLS_IN_ALPHA_INDEX tag can be used to specify the number of columns in -# which the alphabetical index list will be split. -# Minimum value: 1, maximum value: 20, default value: 5. -# This tag requires that the tag ALPHABETICAL_INDEX is set to YES. - -COLS_IN_ALPHA_INDEX = 5 - # In case all classes in a project start with a common prefix, all classes will # be put under the same header in the alphabetical index. The IGNORE_PREFIX tag # can be used to specify a prefix (or a list of prefixes) that should be ignored @@ -1290,15 +1269,6 @@ HTML_COLORSTYLE_SAT = 100 HTML_COLORSTYLE_GAMMA = 80 -# If the HTML_TIMESTAMP tag is set to YES then the footer of each generated HTML -# page will contain the date and time when the page was generated. Setting this -# to YES can help to show when doxygen was last run and thus if the -# documentation is up to date. -# The default value is: NO. -# This tag requires that the tag GENERATE_HTML is set to YES. - -HTML_TIMESTAMP = NO - # If the HTML_DYNAMIC_MENUS tag is set to YES then the generated HTML # documentation will contain a main index with vertical navigation menus that # are dynamically created via JavaScript. If disabled, the navigation index will @@ -1580,17 +1550,6 @@ EXT_LINKS_IN_WINDOW = NO FORMULA_FONTSIZE = 10 -# Use the FORMULA_TRANSPARENT tag to determine whether or not the images -# generated for formulas are transparent PNGs. Transparent PNGs are not -# supported properly for IE 6.0, but are supported on all modern browsers. -# -# Note that when changing this option you need to delete any form_*.png files in -# the HTML output directory before the changes have effect. -# The default value is: YES. -# This tag requires that the tag GENERATE_HTML is set to YES. - -FORMULA_TRANSPARENT = YES - # The FORMULA_MACROFILE can contain LaTeX \newcommand and \renewcommand commands # to create new LaTeX commands to be used in formulas as building blocks. See # the section "Including formulas" for details. @@ -1889,16 +1848,6 @@ LATEX_BATCHMODE = NO LATEX_HIDE_INDICES = NO -# If the LATEX_SOURCE_CODE tag is set to YES then doxygen will include source -# code with syntax highlighting in the LaTeX output. -# -# Note that which sources are shown also depends on other settings such as -# SOURCE_BROWSER. -# The default value is: NO. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_SOURCE_CODE = NO - # The LATEX_BIB_STYLE tag can be used to specify the style to use for the # bibliography, e.g. plainnat, or ieeetr. See # https://en.wikipedia.org/wiki/BibTeX and \cite for more info. @@ -1907,14 +1856,6 @@ LATEX_SOURCE_CODE = NO LATEX_BIB_STYLE = plain -# If the LATEX_TIMESTAMP tag is set to YES then the footer of each generated -# page will contain the date and time when the page was generated. Setting this -# to NO can help when comparing the output of multiple runs. -# The default value is: NO. -# This tag requires that the tag GENERATE_LATEX is set to YES. - -LATEX_TIMESTAMP = NO - # The LATEX_EMOJI_DIRECTORY tag is used to specify the (relative or absolute) # path from which the emoji images will be read. If a relative path is entered, # it will be relative to the LATEX_OUTPUT directory. If left blank the @@ -1979,16 +1920,6 @@ RTF_STYLESHEET_FILE = RTF_EXTENSIONS_FILE = -# If the RTF_SOURCE_CODE tag is set to YES then doxygen will include source code -# with syntax highlighting in the RTF output. -# -# Note that which sources are shown also depends on other settings such as -# SOURCE_BROWSER. -# The default value is: NO. -# This tag requires that the tag GENERATE_RTF is set to YES. - -RTF_SOURCE_CODE = NO - #--------------------------------------------------------------------------- # Configuration options related to the man page output #--------------------------------------------------------------------------- @@ -2085,15 +2016,6 @@ GENERATE_DOCBOOK = NO DOCBOOK_OUTPUT = docbook -# If the DOCBOOK_PROGRAMLISTING tag is set to YES, doxygen will include the -# program listings (including syntax highlighting and cross-referencing -# information) to the DOCBOOK output. Note that enabling this will significantly -# increase the size of the DOCBOOK output. -# The default value is: NO. -# This tag requires that the tag GENERATE_DOCBOOK is set to YES. - -DOCBOOK_PROGRAMLISTING = NO - #--------------------------------------------------------------------------- # Configuration options for the AutoGen Definitions output #--------------------------------------------------------------------------- diff --git a/docs/api/pytorch.rst b/docs/api/pytorch.rst index db8649800..99d850d04 100644 --- a/docs/api/pytorch.rst +++ b/docs/api/pytorch.rst @@ -38,7 +38,7 @@ PyTorch :members: reset, get_states, set_states, add, fork -.. autoapifunction:: transformer_engine.pytorch.autocast +.. autoapiclass:: transformer_engine.pytorch.autocast(enabled=True, calibrating=False, recipe=None, amax_reduction_group=None) .. autoapifunction:: transformer_engine.pytorch.quantized_model_init diff --git a/transformer_engine/common/include/transformer_engine/gemm.h b/transformer_engine/common/include/transformer_engine/gemm.h index 9fe692dd2..a99e0946e 100644 --- a/transformer_engine/common/include/transformer_engine/gemm.h +++ b/transformer_engine/common/include/transformer_engine/gemm.h @@ -108,7 +108,7 @@ void nvte_get_matmul_config_attribute(NVTEMatmulConfig config, NVTEMatmulConfigA /*! \brief Set an option in matrix multiplication configuration. * - * \param[in/out] config Matrix multiplication configuration. + * \param[in,out] config Matrix multiplication configuration. * \param[in] attr Option type. * \param[in] buf Memory address to read option value from. * \param[in] size_in_bytes Size of buf. @@ -298,39 +298,6 @@ void nvte_multi_tensor_gemm(const NVTETensor *A, const NVTETensor *B, NVTETensor bool accumulate, bool use_split_accumulator, int math_sm_count, cudaStream_t stream); -/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ -/*! \brief Grouped matrix multiplication: D = alpha * op(A) @ op(B) + beta * C - * - * \note Requires cuBLAS 13.2+ (CUDA 13.1+) and Blackwell (SM100) or newer GPU architecture. - * Will error at runtime if compiled with an older cuBLAS version or run on - * a pre-Blackwell GPU. - * - * Performs batched GEMM on a collection of matrices with potentially different shapes. - * All tensors in the group must have compatible dimensions for matrix multiplication. - * Uses NVTEGroupedTensor to efficiently handle collections of tensors with contiguous - * memory layout and shape metadata. - * - * \param[in] A Input grouped tensor A. - * \param[in] transa Whether to transpose A matrices. - * \param[in] B Input grouped tensor B. - * \param[in] transb Whether to transpose B matrices. - * \param[in] C Input grouped tensor C (can be NULL for beta=0). - * \param[out] D Output grouped tensor D. - * \param[in] alpha Scale multipliers for A @ B (NVTETensor with num_tensors elements). - * \param[in] beta Scale multipliers for C (NVTETensor with num_tensors elements). - * \param[in] workspace_setup Workspace tensor for pointer array setup. - * \param[in] workspace_cublas Workspace tensor for cuBLAS operations. - * \param[in] config Additional configuration (can be NULL for defaults). - * \param[in] stream CUDA stream for the operation. - * - * Requirements: - * - cuBLAS 13.2+ (CUDA 13.1+) - * - Blackwell (SM100) or newer GPU architecture - * - A, B, C (if provided), D must have the same num_tensors - * - For each i: D[i] = alpha[i] * op(A[i]) @ op(B[i]) + beta[i] * C[i] - * - Shape compatibility: if transa=false, transb=false: - * - A[i]: (M[i], K[i]), B[i]: (K[i], N[i]), D[i]: (M[i], N[i]) - */ /*! \brief Return the required size in bytes for the setup workspace of grouped GEMM. * * The setup workspace stores pointer arrays and per-matrix dimension arrays used @@ -385,6 +352,39 @@ void nvte_convert_int32_to_int64_with_multiplier(const int32_t *src, int64_t *ds void nvte_compute_grouped_tensor_offsets(const int64_t *first_dims, int64_t *offsets, size_t n_groups, int64_t last_dim, cudaStream_t stream); +/* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ +/*! \brief Grouped matrix multiplication: D = alpha * op(A) @ op(B) + beta * C + * + * \note Requires cuBLAS 13.2+ (CUDA 13.1+) and Blackwell (SM100) or newer GPU architecture. + * Will error at runtime if compiled with an older cuBLAS version or run on + * a pre-Blackwell GPU. + * + * Performs batched GEMM on a collection of matrices with potentially different shapes. + * All tensors in the group must have compatible dimensions for matrix multiplication. + * Uses NVTEGroupedTensor to efficiently handle collections of tensors with contiguous + * memory layout and shape metadata. + * + * \param[in] A Input grouped tensor A. + * \param[in] transa Whether to transpose A matrices. + * \param[in] B Input grouped tensor B. + * \param[in] transb Whether to transpose B matrices. + * \param[in] C Input grouped tensor C (can be NULL for beta=0). + * \param[out] D Output grouped tensor D. + * \param[in] alpha Scale multipliers for A @ B (NVTETensor with num_tensors elements). + * \param[in] beta Scale multipliers for C (NVTETensor with num_tensors elements). + * \param[in] workspace_setup Workspace tensor for pointer array setup. + * \param[in] workspace_cublas Workspace tensor for cuBLAS operations. + * \param[in] config Additional configuration (can be NULL for defaults). + * \param[in] stream CUDA stream for the operation. + * + * Requirements: + * - cuBLAS 13.2+ (CUDA 13.1+) + * - Blackwell (SM100) or newer GPU architecture + * - A, B, C (if provided), D must have the same num_tensors + * - For each i: D[i] = alpha[i] * op(A[i]) @ op(B[i]) + beta[i] * C[i] + * - Shape compatibility: if transa=false, transb=false: + * - A[i]: (M[i], K[i]), B[i]: (K[i], N[i]), D[i]: (M[i], N[i]) + */ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedTensor B, int transb, const NVTEGroupedTensor C, NVTEGroupedTensor D, const NVTETensor alpha, const NVTETensor beta, NVTETensor workspace_setup, @@ -398,8 +398,19 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT * instead of NVTEGroupedTensor. This enables discrete per-expert weights as inputA * for Grouped GEMM. * - * \param[in] A_list List of A tensors (length = num_tensors). + * \param[in] A_list List of A tensors (length = num_a_tensors). * \param[in] num_a_tensors Number of tensors in A_list. + * \param[in] transa Whether to transpose A matrices. + * \param[in] B Input grouped tensor B. + * \param[in] transb Whether to transpose B matrices. + * \param[in] C Input grouped tensor C (can be NULL for beta=0). + * \param[out] D Output grouped tensor D. + * \param[in] alpha Scale multipliers for A @ B (NVTETensor with num_tensors elements). + * \param[in] beta Scale multipliers for C (NVTETensor with num_tensors elements). + * \param[in] workspace_setup Workspace tensor for pointer array setup. + * \param[in] workspace_cublas Workspace tensor for cuBLAS operations. + * \param[in] config Additional configuration (can be NULL for defaults). + * \param[in] stream CUDA stream for the operation. */ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors, int transa, const NVTEGroupedTensor B, int transb, @@ -415,10 +426,20 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num * instead of NVTEGroupedTensor. This enables accumulation into non-contiguous * per-expert buffers (for wgrads). * -* \param[in] C_list Optional list of C tensors (length = num_tensors). +* \param[in] A Input grouped tensor A. +* \param[in] transa Whether to transpose A matrices. +* \param[in] B Input grouped tensor B. +* \param[in] transb Whether to transpose B matrices. +* \param[in] C_list Optional list of C tensors (length = num_c_tensors). * \param[in] num_c_tensors Number of tensors in C_list (Can be 0 if C is not provided). -* \param[out] D_list List of D tensors (length = num_tensors). +* \param[out] D_list List of D tensors (length = num_d_tensors). * \param[in] num_d_tensors Number of tensors in D_list. +* \param[in] alpha Scale multipliers for A @ B (NVTETensor with num_tensors elements). +* \param[in] beta Scale multipliers for C (NVTETensor with num_tensors elements). +* \param[in] workspace_setup Workspace tensor for pointer array setup. +* \param[in] workspace_cublas Workspace tensor for cuBLAS operations. +* \param[in] config Additional configuration (can be NULL for defaults). +* \param[in] stream CUDA stream for the operation. * \note All tensors in C_list and D_list must share the same dtype. */ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index 488f25915..045ae8889 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -282,7 +282,7 @@ void nvte_zero_tensor(const NVTETensor tensor, cudaStream_t stream); * * \warning Deprecated in favor of nvte_set_tensor_param_v2. * - * \param[in/out] tensor Tensor. + * \param[in,out] tensor Tensor. * \param[in] param_name The parameter to be set. * \param[in] param The value to be set. */ @@ -300,7 +300,7 @@ NVTEBasicTensor nvte_get_tensor_param(const NVTETensor tensor, NVTETensorParam p /*! \brief Set a tensor parameter. * - * \param[in/out] tensor Tensor. + * \param[in,out] tensor Tensor. * \param[in] param Tensor parameter type. * \param[in] buf Memory address to read parameter value. * \param[in] size_in_bytes Size of buf. @@ -406,7 +406,7 @@ void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config, /*! \brief Set an option in quantization config. * - * \param[in/out] config Quantization config. + * \param[in,out] config Quantization config. * \param[in] attr Option type. * \param[in] buf Memory address to read option value. * \param[in] size_in_bytes Size of buf. @@ -510,7 +510,7 @@ void nvte_destroy_grouped_tensor(NVTEGroupedTensor tensor); /* EXPERIMENTAL FEATURE AND SUBJECT TO CHANGE. */ /*! \brief Set a grouped tensor parameter. * - * \param[in/out] tensor Grouped tensor. + * \param[in,out] tensor Grouped tensor. * \param[in] param Grouped tensor parameter type. * \param[in] buf Memory address to read parameter value. * \param[in] size_in_bytes Size of buf. From 4eab389b9ceac999e9700005c0a96c176f93429a Mon Sep 17 00:00:00 2001 From: Phuong Nguyen Date: Tue, 12 May 2026 14:01:06 -0700 Subject: [PATCH 056/180] [JAX] Add wait per multi-proc cleanup in `L0_jax_distributed_unittest` (#2979) add wait per multi-proc test cleanup Signed-off-by: Phuong Nguyen --- examples/jax/collective_gemm/run_test_cgemm.sh | 1 + examples/jax/encoder/run_test_multiprocessing_encoder.sh | 1 + 2 files changed, 2 insertions(+) diff --git a/examples/jax/collective_gemm/run_test_cgemm.sh b/examples/jax/collective_gemm/run_test_cgemm.sh index 8340d2010..c0a095d6b 100644 --- a/examples/jax/collective_gemm/run_test_cgemm.sh +++ b/examples/jax/collective_gemm/run_test_cgemm.sh @@ -143,5 +143,6 @@ wait # Final cleanup (trap will also call cleanup on exit) cleanup +wait exit $HAS_FAILURE diff --git a/examples/jax/encoder/run_test_multiprocessing_encoder.sh b/examples/jax/encoder/run_test_multiprocessing_encoder.sh index 3c1f2ba1f..4242c77c1 100644 --- a/examples/jax/encoder/run_test_multiprocessing_encoder.sh +++ b/examples/jax/encoder/run_test_multiprocessing_encoder.sh @@ -98,5 +98,6 @@ wait # Final cleanup (trap will also call cleanup on exit) cleanup +wait exit $HAS_FAILURE From 472ae5519bcfacc0bab4c361f2dd6263754f792e Mon Sep 17 00:00:00 2001 From: vasunvidia <108759426+vasunvidia@users.noreply.github.com> Date: Tue, 12 May 2026 15:18:27 -0700 Subject: [PATCH 057/180] Avoid CPU offload wait_event for validation (#2793) * Avoid CPU offload wait_event for validation Signed-off-by: Vasudevan Rengasamy * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci Signed-off-by: Vasudevan Rengasamy --------- Signed-off-by: Vasudevan Rengasamy Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/pytorch/cpu_offload.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/transformer_engine/pytorch/cpu_offload.py b/transformer_engine/pytorch/cpu_offload.py index ed10909b8..c81c18e64 100644 --- a/transformer_engine/pytorch/cpu_offload.py +++ b/transformer_engine/pytorch/cpu_offload.py @@ -307,8 +307,9 @@ def start_offload(self): # needed to restore pre-offload state after reload. self.aux = aux - self.finish_offload_event = torch.cuda.Event() - self.finish_offload_event.record(self.offload_stream) + if len(self.fwd_gpu_tensor_group.tensor_list) > 0: + self.finish_offload_event = torch.cuda.Event() + self.finish_offload_event.record(self.offload_stream) def release_activation_forward_gpu_memory(self): """ @@ -319,13 +320,13 @@ def release_activation_forward_gpu_memory(self): func_name="release_activation_forward_gpu_memory", allowed_states=["offload_started"] ) self.state = "offload_finished" + if len(self.fwd_gpu_tensor_group.tensor_list) > 0: + torch.cuda.current_stream().wait_event(self.finish_offload_event) # type: ignore[arg-type] - torch.cuda.current_stream().wait_event(self.finish_offload_event) # type: ignore[arg-type] - - # GPU memory can be released safely after the offload. - # Notice that the memory needs to be kept alive when GPU->CPU copy is performed. - self.fwd_gpu_tensor_group = TensorGroup() - del self.finish_offload_event + # GPU memory can be released safely after the offload. + # Notice that the memory needs to be kept alive when GPU->CPU copy is performed. + self.fwd_gpu_tensor_group = TensorGroup() + del self.finish_offload_event def start_reload(self): """ From c3a1d30227911a66038592b1b18899d00ddb2686 Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Tue, 12 May 2026 17:01:50 -0700 Subject: [PATCH 058/180] [Core] Report CUDA versions when NVRTC compilation fails (#2842) * [NVRTC] Warn on CUDA version mismatch after compilation failure When NVRTC kernel compilation fails, detect whether the linked NVRTC library and the CUDA headers used for compilation are from different CUDA versions, and if so emit an actionable note to stderr pointing the user toward NVTE_CUDA_INCLUDE_DIR / CUDA_HOME / LD_LIBRARY_PATH. The header version is obtained by compiling a tiny probe program that embeds CUDA_VERSION (from cuda.h) into a static_assert failure message, so the macro is resolved by the actual preprocessor rather than by parsing header text. All probe failures are silent; the check is purely informational and never causes a premature error. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Tim Moon * Move CUDA header version check to CUDA runtime utils Still buggy, include_directory_version returns CUDA runtime version instead of header version. Signed-off-by: Tim Moon * [NVRTC] Fix CUDA header version detection The NVRTC probe approach was broken: NVRTC pre-defines CUDART_VERSION to its own version before processing any includes, so the probe always returned the NVRTC version regardless of the headers on the include path. Fix by reading cuda_runtime_api.h as text and parsing the "#define CUDART_VERSION " line directly. This is immune to NVRTC's internal macro management, and the format has been stable across all CUDA versions. Also decode raw CUDA version integers to "major.minor" strings in the error message for readability. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Tim Moon * [NVRTC] Add unit tests for CUDA header detection Test that the CUDA include directory is found and that its version matches the compile-time CUDART_VERSION. Also export transformer_engine::cuda::* symbols and tighten the rtc export pattern in the version script. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tweak version message Suggestion from @ptrendx Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Remove test Test required exposing CUDA utility functions externally, which is beyond the scope of this work. Signed-off-by: Tim Moon --------- Signed-off-by: Tim Moon Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Claude Sonnet 4.6 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Przemyslaw Tredak --- .../common/util/cuda_runtime.cpp | 44 +++++++++++++++++++ transformer_engine/common/util/cuda_runtime.h | 15 +++++++ transformer_engine/common/util/rtc.cpp | 37 +++++++++++++++- 3 files changed, 94 insertions(+), 2 deletions(-) diff --git a/transformer_engine/common/util/cuda_runtime.cpp b/transformer_engine/common/util/cuda_runtime.cpp index 4b43940a5..504d761bb 100644 --- a/transformer_engine/common/util/cuda_runtime.cpp +++ b/transformer_engine/common/util/cuda_runtime.cpp @@ -9,6 +9,7 @@ #include #include +#include #include #include "../common.h" @@ -202,6 +203,49 @@ const std::string &include_directory(bool required) { return path; } +int include_directory_version(bool required) { + // Header path + const auto &include_dir = cuda::include_directory(false); + if (include_dir.empty()) { + if (required) { + NVTE_ERROR( + "Could not detect version of CUDA Toolkit headers " + "(CUDA Toolkit headers not found)."); + } + return -1; + } + + // Parse CUDART_VERSION from cuda_runtime_api.h. + const auto header_path = std::filesystem::path(include_dir) / "cuda_runtime_api.h"; + std::ifstream header_file(header_path); + if (header_file.is_open()) { + const std::string define_prefix = "#define CUDART_VERSION "; + std::string line; + while (std::getline(header_file, line)) { + const auto pos = line.find(define_prefix); + if (pos == std::string::npos) { + continue; + } + try { + const int version = std::stoi(line.substr(pos + define_prefix.size())); + if (version > 0) { + return version; + } + } catch (...) { + continue; + } + } + } + + if (required) { + NVTE_ERROR( + "Could not detect version of CUDA Toolkit headers " + "(Could not parse CUDART_VERSION from ", + header_path.string(), ")."); + } + return -1; +} + int cudart_version() { auto get_version = []() -> int { int version; diff --git a/transformer_engine/common/util/cuda_runtime.h b/transformer_engine/common/util/cuda_runtime.h index f0aa23962..0f3559400 100644 --- a/transformer_engine/common/util/cuda_runtime.h +++ b/transformer_engine/common/util/cuda_runtime.h @@ -67,6 +67,21 @@ bool supports_multicast(int device_id = -1); */ const std::string &include_directory(bool required = false); +/* \brief Version number of CUDA Toolkit headers + * + * The headers are accessed at run-time and its CUDA version may + * differ from compile-time and from the CUDA Runtime. The header path + * can be configured by setting NVTE_CUDA_INCLUDE_DIR in the + * environment (default is to search in common install paths). + * + * \param[in] required Whether to throw exception if headers are not + * found or if version cannot be determined. + * + * \return CUDA version encoded as major * 1000 + minor * 10, or -1 if + * it could not be determined. + */ +int include_directory_version(bool required = false); + /* \brief CUDA Runtime version number at run-time * * Versions may differ between compile-time and run-time. diff --git a/transformer_engine/common/util/rtc.cpp b/transformer_engine/common/util/rtc.cpp index 7925fdcee..70024a202 100644 --- a/transformer_engine/common/util/rtc.cpp +++ b/transformer_engine/common/util/rtc.cpp @@ -12,6 +12,7 @@ #include "../common.h" #include "../util/cuda_driver.h" +#include "../util/cuda_runtime.h" #include "../util/string.h" #include "../util/system.h" @@ -175,14 +176,46 @@ void KernelManager::compile(const std::string& kernel_label, const std::string& const nvrtcResult compile_result = nvrtcCompileProgram(program, opts_ptrs.size(), opts_ptrs.data()); if (compile_result != NVRTC_SUCCESS) { - // Display log if compilation failed - std::string log = concat_strings("NVRTC compilation log for ", filename, ":\n"); + std::string log; + + // Decode CUDA version number to "major.minor" string + auto version_string = [](int v) -> std::string { + if (v < 0) { + return ""; + } + return concat_strings(v / 1000, ".", (v % 1000) / 10); + }; + + // Check CUDA versions + const int build_version = CUDA_VERSION; + int nvrtc_version = -1; + int nvrtc_version_major = 0, nvrtc_version_minor = 0; + if (nvrtcVersion(&nvrtc_version_major, &nvrtc_version_minor) == NVRTC_SUCCESS) { + nvrtc_version = nvrtc_version_major * 1000 + nvrtc_version_minor * 10; + } + const int header_version = cuda::include_directory_version(); + log += concat_strings("Compile-time CUDA version: ", version_string(build_version), "\n", + "Run-time NVRTC version: ", version_string(nvrtc_version), "\n", + "Run-time CUDA headers version: ", version_string(header_version), "\n"); + if (nvrtc_version != header_version) { + log += concat_strings( + "\nWarning: CUDA versions do not match between NVRTC and CUDA headers (", + cuda::include_directory(), + "). " + "Consider changing the CUDA header search path (by setting NVTE_CUDA_INCLUDE_DIR) " + "or the linked CUDA Runtime (by setting CUDA_HOME or LD_LIBRARY_PATH).\n\n"); + } + + // Get build log + log += concat_strings("NVRTC compilation log for ", filename, ":\n"); const size_t log_offset = log.size(); size_t log_size; NVTE_CHECK_NVRTC(nvrtcGetProgramLogSize(program, &log_size)); log.resize(log_offset + log_size); NVTE_CHECK_NVRTC(nvrtcGetProgramLog(program, &log[log_offset])); log.back() = '\n'; + + // Display log and throw error std::cerr << log; NVTE_CHECK_NVRTC(compile_result); } From 4631d97fdf32f721fcc9353bbc979237f1120c4e Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Tue, 12 May 2026 17:02:40 -0700 Subject: [PATCH 059/180] [pyTorch] Replace the make_empty implementation to use C++ implementation (#2666) * Replace the make_empty implementation to use C++ implementation for the known quantizers Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix lint Signed-off-by: Przemek Tredak * Handle the device passed as string Signed-off-by: Przemek Tredak * Fix Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fixes Signed-off-by: Przemek Tredak * Replace the make_empty implementation to use C++ implementation for the known quantizers Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix lint Signed-off-by: Przemek Tredak * Handle the device passed as string Signed-off-by: Przemek Tredak * Fix Signed-off-by: Przemek Tredak * Fixes Signed-off-by: Przemek Tredak * Fix duplicate create_empty_quantized_tensor after merge The merge with main introduced duplicate function definition, declaration, and pybind registration for create_empty_quantized_tensor. Remove the duplicates. Signed-off-by: Przemek Tredak * Fix device index resolution in create_tensor Change the device parameter from at::Device with default torch::kCUDA to std::optional with default nullopt. When no device is specified, resolve to the current CUDA device via c10::cuda::current_device(), ensuring the device always has a valid index. This fixes autograd engine assertions when tensors created without an explicit device are used in backward passes. Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Guard make_empty for custom quantizers without C++ converter Custom quantizers that set self.custom = True and don't override make_empty() will now get a clear NotImplementedError instead of hitting an opaque C++ NVTE_ERROR("Unexpected type for quantizer"). Signed-off-by: Przemek Tredak * Fix the device from the passed data case Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Przemek Tredak Signed-off-by: vthumbe1503 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: vthumbe1503 --- transformer_engine/pytorch/csrc/common.h | 43 ++++--- transformer_engine/pytorch/csrc/extensions.h | 5 +- .../pytorch/csrc/extensions/cast.cpp | 8 ++ .../pytorch/csrc/extensions/pybind.cpp | 4 + transformer_engine/pytorch/csrc/quantizer.cpp | 105 ++++++++++++------ .../pytorch/quantized_tensor.py | 31 +++++- .../pytorch/tensor/float8_blockwise_tensor.py | 56 ---------- .../pytorch/tensor/float8_tensor.py | 85 -------------- .../pytorch/tensor/mxfp8_tensor.py | 64 ----------- .../pytorch/tensor/nvfp4_tensor.py | 93 ---------------- 10 files changed, 142 insertions(+), 352 deletions(-) diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 35a459351..94350da1e 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -102,8 +102,9 @@ class Quantizer { virtual void set_quantization_params(TensorWrapper* tensor) const = 0; /*! @brief Construct a tensor with uninitialized data */ - virtual std::pair create_tensor(const std::vector& shape, - DType dtype) const = 0; + virtual std::pair create_tensor( + const std::vector& shape, DType dtype, + std::optional device = std::nullopt, bool pin_memory = false) const = 0; /*! @brief Construct a grouped tensor with uninitialized data */ virtual std::pair create_grouped_tensor( @@ -144,8 +145,9 @@ class NoneQuantizer : public Quantizer { void set_quantization_params(TensorWrapper* tensor) const override {} - std::pair create_tensor(const std::vector& shape, - DType dtype) const override; + std::pair create_tensor( + const std::vector& shape, DType dtype, + std::optional device = std::nullopt, bool pin_memory = false) const override; std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, @@ -174,8 +176,9 @@ class Float8Quantizer : public Quantizer { void set_quantization_params(TensorWrapper* tensor) const override; - std::pair create_tensor(const std::vector& shape, - DType dtype) const override; + std::pair create_tensor( + const std::vector& shape, DType dtype, + std::optional device = std::nullopt, bool pin_memory = false) const override; std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, @@ -183,10 +186,10 @@ class Float8Quantizer : public Quantizer { size_t logical_last_dim) const override; /*! @brief Construct a tensor with pre-initialized data */ - std::pair create_tensor(const std::vector& shape, DType dtype, - std::optional data, - std::optional transpose, - std::optional scale_inv) const; + std::pair create_tensor( + const std::vector& shape, DType dtype, std::optional data, + std::optional transpose, std::optional scale_inv, + std::optional device = std::nullopt, bool pin_memory = false) const; std::pair convert_and_update_tensor(py::object shape) const override; @@ -208,8 +211,9 @@ class Float8CurrentScalingQuantizer : public Quantizer { void set_quantization_params(TensorWrapper* tensor) const override; - std::pair create_tensor(const std::vector& shape, - DType dtype) const override; + std::pair create_tensor( + const std::vector& shape, DType dtype, + std::optional device = std::nullopt, bool pin_memory = false) const override; std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, @@ -270,8 +274,9 @@ class Float8BlockQuantizer : public Quantizer { // Create a python Float8BlockQuantized tensor and C++ wrapper // for the tensor. Should set quantized data, scales for rowwise // and optionally columnwise usage. - std::pair create_tensor(const std::vector& shape, - DType dtype) const override; + std::pair create_tensor( + const std::vector& shape, DType dtype, + std::optional device = std::nullopt, bool pin_memory = false) const override; std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, @@ -294,8 +299,9 @@ class MXFP8Quantizer : public Quantizer { void set_quantization_params(TensorWrapper* tensor) const override; - std::pair create_tensor(const std::vector& shape, - DType dtype) const override; + std::pair create_tensor( + const std::vector& shape, DType dtype, + std::optional device = std::nullopt, bool pin_memory = false) const override; std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, @@ -333,8 +339,9 @@ class NVFP4Quantizer : public Quantizer { void set_quantization_params(TensorWrapper* tensor) const override; - std::pair create_tensor(const std::vector& shape, - DType dtype) const override; + std::pair create_tensor( + const std::vector& shape, DType dtype, + std::optional device = std::nullopt, bool pin_memory = false) const override; std::pair create_grouped_tensor( size_t num_tensors, const std::vector& logical_shape, DType dtype, diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 9b10a9c5a..8082ff07e 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -320,9 +320,12 @@ std::vector bulk_allocate(const std::vector> &sh std::optional> alignments = std::nullopt); /*************************************************************************************************** - * Cast + * Quantize **************************************************************************************************/ +py::object create_empty_quantized_tensor(py::handle quantizer, const std::vector &shape, + at::ScalarType dtype, at::Device device, bool pin_memory); + py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::object &output, std::optional noop_flag); diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 3ada2459c..2b38339d6 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -65,6 +65,14 @@ py::object quantize(const at::Tensor &tensor, py::handle quantizer, const py::ob return output_py; } +py::object create_empty_quantized_tensor(py::handle quantizer, const std::vector &shape, + at::ScalarType dtype, at::Device device, bool pin_memory) { + auto quantizer_cpp = convert_quantizer(quantizer); + auto te_dtype = GetTransformerEngineDType(dtype); + auto [_, output_py] = quantizer_cpp->create_tensor(shape, te_dtype, device, pin_memory); + return output_py; +} + namespace { // helper functions for NVFP4 grouped quantization (cuda graph safe with shapes stored in device without D2H copy) diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index a813f3119..a4571c64e 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -139,6 +139,10 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("output") = py::none(), py::arg("noop") = py::none()); m.def("dequantize", &transformer_engine::pytorch::dequantize, "Dequantize", py::arg("input"), py::arg("otype")); + m.def("create_empty_quantized_tensor", + &transformer_engine::pytorch::create_empty_quantized_tensor, + "Create an empty quantized tensor", py::arg("quantizer"), py::arg("shape"), + py::arg("dtype"), py::arg("device"), py::arg("pin_memory")); m.def("group_quantize", transformer_engine::pytorch::group_quantize, py::arg("tensor"), py::arg("quantizer"), py::arg("num_tensors"), py::arg("first_dims")); m.def("group_dequantize", transformer_engine::pytorch::group_dequantize, diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 82dfe4d22..7045995dd 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -16,6 +16,29 @@ namespace transformer_engine::pytorch { namespace { +/*! @brief Resolve an optional device to a concrete CUDA device + * + * If no device is provided, uses the current CUDA device. + */ +at::Device resolve_device(std::optional device, + const std::optional& data = std::nullopt) { + if (device.has_value() && data.has_value()) { + // Ensure that they are the same + const auto provided_device = *device; + const auto data_device = data->device(); + NVTE_CHECK(provided_device == data_device, + "Provided device and the device of the provided data tensor are not the same."); + return provided_device; + } + if (device.has_value()) { + return *device; + } + if (data.has_value()) { + return data->device(); + } + return at::Device(torch::kCUDA, c10::cuda::current_device()); +} + /*! @brief Transposed tensor shape * * The tensor is interpreted as a 2D matrix by flattening all but the @@ -129,10 +152,13 @@ Float8Quantizer::Float8Quantizer(const py::handle& quantizer) : Quantizer(quanti this->dtype = type; } -std::pair NoneQuantizer::create_tensor(const std::vector& shape, - DType dtype) const { +std::pair NoneQuantizer::create_tensor( + const std::vector& shape, DType dtype, std::optional device_opt, + bool pin_memory) const { + const auto device = resolve_device(device_opt); const std::vector shape_int64(shape.begin(), shape.end()); - const auto opts = at::TensorOptions().dtype(GetATenDType(dtype)).device(torch::kCUDA); + const auto opts = + at::TensorOptions().dtype(GetATenDType(dtype)).device(device).pinned_memory(pin_memory); return create_tensor(shape, dtype, at::empty(shape_int64, opts)); } @@ -240,22 +266,29 @@ void Float8Quantizer::set_quantization_params(TensorWrapper* tensor) const { } std::pair Float8Quantizer::create_tensor( - const std::vector& shape, DType dtype) const { - const auto opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + const std::vector& shape, DType dtype, std::optional device_opt, + bool pin_memory) const { + const auto device = resolve_device(device_opt); + const auto opts = + at::TensorOptions().dtype(torch::kFloat32).device(device).pinned_memory(pin_memory); at::Tensor scale_inv = at::empty(std::vector{1}, opts); - return create_tensor(shape, dtype, std::nullopt, std::nullopt, std::move(scale_inv)); + return create_tensor(shape, dtype, std::nullopt, std::nullopt, std::move(scale_inv), device, + pin_memory); } std::pair Float8Quantizer::create_tensor( const std::vector& shape, DType dtype, std::optional data, - std::optional transpose, std::optional scale_inv) const { + std::optional transpose, std::optional scale_inv, + std::optional device_opt, bool pin_memory) const { + const auto device = resolve_device(device_opt, data); using namespace pybind11::literals; int is_non_tn_fp8_gemm_supported = nvte_is_non_tn_fp8_gemm_supported(); // Initialize data tensor const bool with_data = rowwise_usage || is_non_tn_fp8_gemm_supported; if (with_data && !data) { const std::vector shape_int64(shape.begin(), shape.end()); - const auto opts = at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA); + const auto opts = + at::TensorOptions().dtype(torch::kUInt8).device(device).pinned_memory(pin_memory); data = at::empty(shape_int64, opts); } else if (!with_data && data) { data.reset(); @@ -266,7 +299,8 @@ std::pair Float8Quantizer::create_tensor( const bool with_transpose = columnwise_usage && !is_non_tn_fp8_gemm_supported; if (with_transpose && !transpose) { const auto transpose_shape = make_transpose_shape(shape); - const auto opts = at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA); + const auto opts = + at::TensorOptions().dtype(torch::kUInt8).device(device).pinned_memory(pin_memory); transpose = at::empty(transpose_shape, opts); } else if (!with_transpose && transpose) { transpose.reset(); @@ -277,10 +311,6 @@ std::pair Float8Quantizer::create_tensor( scale_inv = at::reciprocal(scale); } py::object scale_inv_py = py::cast(*scale_inv); - at::Device device = - with_data ? data->device() - : (with_transpose ? transpose->device() - : at::Device(torch::kCUDA, c10::cuda::current_device())); // Construct Python FP8 tensor py::object out_py; if (internal) { @@ -555,7 +585,9 @@ Float8CurrentScalingQuantizer::Float8CurrentScalingQuantizer(const py::handle& q void Float8CurrentScalingQuantizer::set_quantization_params(TensorWrapper* tensor) const {} std::pair Float8CurrentScalingQuantizer::create_tensor( - const std::vector& shape, DType dtype) const { + const std::vector& shape, DType dtype, std::optional device_opt, + bool pin_memory) const { + const auto device = resolve_device(device_opt); using namespace pybind11::literals; // Initialize data tensor @@ -564,7 +596,8 @@ std::pair Float8CurrentScalingQuantizer::create_tenso const bool with_data = rowwise_usage || is_non_tn_fp8_gemm_supported; if (with_data) { const std::vector shape_int64(shape.begin(), shape.end()); - const auto opts = at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA); + const auto opts = + at::TensorOptions().dtype(torch::kUInt8).device(device).pinned_memory(pin_memory); data_tensor = at::empty(shape_int64, opts); } @@ -573,20 +606,18 @@ std::pair Float8CurrentScalingQuantizer::create_tenso const bool with_transpose = columnwise_usage && !is_non_tn_fp8_gemm_supported; if (with_transpose) { const auto transpose_shape = make_transpose_shape(shape); - const auto opts = at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA); + const auto opts = + at::TensorOptions().dtype(torch::kUInt8).device(device).pinned_memory(pin_memory); transpose_tensor = at::empty(transpose_shape, opts); } // Initialize scale-inverse tensor at::Tensor scale_inv_tensor; { const std::vector scale_inv_shape = {1}; - const auto opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + const auto opts = + at::TensorOptions().dtype(torch::kFloat32).device(device).pinned_memory(pin_memory); scale_inv_tensor = at::empty(scale_inv_shape, opts); } - at::Device device = - with_data ? data_tensor.device() - : (with_transpose ? transpose_tensor.device() - : at::Device(torch::kCUDA, c10::cuda::current_device())); // Construct Python FP8 tensor py::object out_py; py::object scale_inv_py = py::cast(scale_inv_tensor); @@ -924,7 +955,9 @@ Float8BlockQuantizer::Float8BlockQuantizer(const py::handle& quantizer) : Quanti void Float8BlockQuantizer::set_quantization_params(TensorWrapper* tensor) const {} std::pair Float8BlockQuantizer::create_tensor( - const std::vector& shape, DType dtype) const { + const std::vector& shape, DType dtype, std::optional device_opt, + bool pin_memory) const { + const auto device = resolve_device(device_opt); using namespace pybind11::literals; std::vector torch_shape; for (auto s : shape) { @@ -935,8 +968,8 @@ std::pair Float8BlockQuantizer::create_tensor( at::TensorOptions opts; at::TensorOptions scale_opts; at::Tensor data_rowwise, data_colwise, scale_inv_rowwise, scale_inv_colwise; - opts = opts.dtype(torch::kUInt8).device(torch::kCUDA); - scale_opts = scale_opts.dtype(torch::kFloat32).device(torch::kCUDA); + opts = opts.dtype(torch::kUInt8).device(device).pinned_memory(pin_memory); + scale_opts = scale_opts.dtype(torch::kFloat32).device(device).pinned_memory(pin_memory); if (rowwise_usage) { data_rowwise = at::empty(torch_shape, opts); @@ -1015,6 +1048,7 @@ std::pair Float8BlockQuantizer::create_tensor( kwargs["fp8_dtype"] = py::cast(this->dtype); kwargs["quantizer"] = this->quantizer; kwargs["is_2D_scaled"] = py::cast(block_scaling_dim == 2); + kwargs["device"] = py::cast(device); py::tuple args(0); PyObject* result = PyObject_Call(reinterpret_cast(Float8BlockwiseQTensorPythonClass), @@ -1334,8 +1368,10 @@ MXFP8Quantizer::MXFP8Quantizer(const py::handle& quantizer) : Quantizer(quantize void MXFP8Quantizer::set_quantization_params(TensorWrapper* tensor) const {} -std::pair MXFP8Quantizer::create_tensor(const std::vector& shape, - DType dtype) const { +std::pair MXFP8Quantizer::create_tensor( + const std::vector& shape, DType dtype, std::optional device_opt, + bool pin_memory) const { + const auto device = resolve_device(device_opt); using namespace pybind11::literals; // Scaling factor format @@ -1353,7 +1389,8 @@ std::pair MXFP8Quantizer::create_tensor(const std::ve // Allocate tensors at::Tensor rowwise_data_tensor, rowwise_scale_inv_tensor; at::Tensor columnwise_data_tensor, columnwise_scale_inv_tensor; - const auto uint8_tensor_opts = at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA); + const auto uint8_tensor_opts = + at::TensorOptions().dtype(torch::kUInt8).device(device).pinned_memory(pin_memory); if (rowwise_usage) { const std::vector scale_inv_shape_int64(rowwise_scale_inv_shape.begin(), rowwise_scale_inv_shape.end()); @@ -1413,6 +1450,7 @@ std::pair MXFP8Quantizer::create_tensor(const std::ve kwargs["fp8_dtype"] = py::cast(this->dtype); kwargs["quantizer"] = this->quantizer; kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); + kwargs["device"] = py::cast(device); py::tuple args(0); PyObject* result = PyObject_Call(reinterpret_cast(MXFP8TensorPythonClass), @@ -1722,8 +1760,10 @@ void NVFP4Quantizer::set_quantization_params(TensorWrapper* tensor) const { columnwise_data.shape); } -std::pair NVFP4Quantizer::create_tensor(const std::vector& shape, - DType dtype) const { +std::pair NVFP4Quantizer::create_tensor( + const std::vector& shape, DType dtype, std::optional device_opt, + bool pin_memory) const { + const auto device = resolve_device(device_opt); using namespace pybind11::literals; // Scaling factor format @@ -1749,8 +1789,10 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve // Allocate tensors at::Tensor rowwise_data_tensor, rowwise_scale_inv_tensor, amax_rowwise; at::Tensor columnwise_data_tensor, columnwise_scale_inv_tensor, amax_columnwise; - const auto bit8_tensor_opts = at::TensorOptions().dtype(torch::kUInt8).device(torch::kCUDA); - const auto bit32_tensor_opts = at::TensorOptions().dtype(torch::kFloat32).device(torch::kCUDA); + const auto bit8_tensor_opts = + at::TensorOptions().dtype(torch::kUInt8).device(device).pinned_memory(pin_memory); + const auto bit32_tensor_opts = + at::TensorOptions().dtype(torch::kFloat32).device(device).pinned_memory(pin_memory); if (rowwise_usage) { const std::vector scale_inv_shape_int64(rowwise_scale_inv_shape.begin(), rowwise_scale_inv_shape.end()); @@ -1831,6 +1873,7 @@ std::pair NVFP4Quantizer::create_tensor(const std::ve kwargs["fp4_dtype"] = py::cast(this->dtype); kwargs["quantizer"] = this->quantizer; kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); + kwargs["device"] = py::cast(device); kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); py::tuple args(0); PyObject* result = PyObject_Call(reinterpret_cast(NVFP4TensorPythonClass), diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index a7722f777..7163e2b17 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -13,6 +13,8 @@ import torch from torch.utils._pytree import tree_map +import transformer_engine_torch as tex + from transformer_engine.common.recipe import Recipe from transformer_engine.pytorch.tensor._quantization_helpers import ( _QuantizeFunc, @@ -311,13 +313,34 @@ def make_empty( shape: Iterable[int], *, dtype: torch.dtype = torch.float32, - device: Optional[torch.device] = None, + device: Optional[Union[torch.device, str]] = None, + requires_grad: bool = False, + pin_memory: bool = False, ) -> QuantizedTensor: """Construct quantized tensor with uninitialized data""" - raise NotImplementedError( - f"{self.__class__.__name__} class does not implement make_empty function, " - "required for construction of unintialized quantized tensor" + + # Guard for custom quantizers that don't have a registered C++ converter. + # Without this, they would hit an opaque C++ NVTE_ERROR. + if getattr(self, "custom", False): + raise NotImplementedError( + f"{self.__class__.__name__} class does not implement make_empty function, " + "required for construction of uninitialized quantized tensor" + ) + + if device is None: + device = torch.device("cuda") + # Handle the device passed as string + device = torch.device(device) + result = tex.create_empty_quantized_tensor( + self, + list(shape), + dtype, + device, + pin_memory, ) + if requires_grad: + result.requires_grad_(True) + return result def calibrate(self, tensor: torch.Tensor) -> None: """Calibrate quantizer state diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index 914397b9b..d0296902a 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -202,62 +202,6 @@ def is_quantizable(self, inp: torch.Tensor) -> bool: return False return True - def make_empty( - self, - shape: Iterable[int], - *, - dtype: torch.dtype = torch.float32, - device: Optional[torch.device] = None, - requires_grad: bool = False, - pin_memory: bool = False, - ) -> Float8BlockwiseQTensor: - """Construct quantized tensor with uninitialized data""" - - tensor_kwargs = { - "device": torch.device("cuda") if device is None else device, - "pin_memory": pin_memory, - } - - # Allocate buffers for row-scaled data - rowwise_data = None - rowwise_scale_inv = None - if self.rowwise_usage: - rowwise_data = torch.empty(shape, dtype=torch.uint8, **tensor_kwargs) - rowwise_scale_inv = torch.empty( - self.get_scale_shape(shape, columnwise=False), - dtype=torch.float32, - **tensor_kwargs, - ) - - # Allocate buffers for column-scaled data - columnwise_data = None - columnwise_scale_inv = None - if self.columnwise_usage: - columnwise_data = torch.empty( - self.get_columnwise_shape(shape), - dtype=torch.uint8, - **tensor_kwargs, - ) - columnwise_scale_inv = torch.empty( - self.get_scale_shape(shape, columnwise=True), - dtype=torch.float32, - **tensor_kwargs, - ) - - # Construct FP8 tensor - return Float8BlockwiseQTensor( - shape=shape, - dtype=dtype, - fp8_dtype=self.dtype, - rowwise_data=rowwise_data, - rowwise_scale_inv=rowwise_scale_inv, - columnwise_data=columnwise_data, - columnwise_scale_inv=columnwise_scale_inv, - quantizer=self, - is_2D_scaled=self.block_scaling_dim == 2, - requires_grad=requires_grad, - ) - def calibrate(self, tensor: torch.Tensor) -> None: # NOTE: This interface is specific to requirements like delayed scaling # where state from an estimator influences distribution parameters. diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index ed6091c85..c4c5934f9 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -112,49 +112,6 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: """Quantize tensor implementation""" return tex.quantize(tensor, self) - def make_empty( - self, - shape: Iterable[int], - *, - dtype: torch.dtype = torch.float32, - device: Optional[torch.device] = None, - requires_grad: bool = False, - pin_memory: bool = False, - ) -> Float8Tensor: - - # Canonicalize tensor attributes - if device is None: - device = torch.device("cuda") - - # Allocate FP8 data - data = None - if self.rowwise_usage: - data = torch.empty(shape, dtype=torch.uint8, device=device, pin_memory=pin_memory) - - # Allocate FP8 data transpose if needed - data_transpose = None - if self.columnwise_usage: - transpose_shape = [shape[-1]] + list(shape[:-1]) - data_transpose = torch.empty( - transpose_shape, - dtype=torch.uint8, - device=device, - pin_memory=pin_memory, - ) - - # Construct FP8 tensor - return Float8Tensor( - shape=shape, - dtype=dtype, - data=data, - fp8_scale_inv=torch.empty(1, dtype=torch.float32, device=device, pin_memory=pin_memory), - fp8_dtype=self.dtype, - requires_grad=requires_grad, - data_transpose=data_transpose, - quantizer=self, - device=device, - ) - def calibrate(self, tensor: torch.Tensor) -> None: amin, amax = tensor.aminmax() self.amax.copy_(torch.max(-amin, amax)) @@ -337,48 +294,6 @@ def quantize_impl(self, tensor: torch.Tensor) -> QuantizedTensor: """Quantize tensor implementation""" return tex.quantize(tensor, self) - def make_empty( - self, - shape: Iterable[int], - *, - dtype: torch.dtype = torch.float32, - device: Optional[torch.device] = None, - requires_grad: bool = False, - pin_memory: bool = False, - ) -> Float8Tensor: - - # Canonicalize tensor attributes - if device is None: - device = torch.device("cuda") - - # Allocate FP8 data - data = None - if self.rowwise_usage: - data = torch.empty(shape, dtype=torch.uint8, device=device, pin_memory=pin_memory) - - # Allocate FP8 data transpose if needed - data_transpose = None - if self.columnwise_usage: - transpose_shape = [shape[-1]] + list(shape[:-1]) - data_transpose = torch.empty( - transpose_shape, - dtype=torch.uint8, - device=device, - pin_memory=pin_memory, - ) - # Construct FP8 tensor - return Float8Tensor( - shape=shape, - dtype=dtype, - data=data, - fp8_scale_inv=torch.empty(1, dtype=torch.float32, device=device, pin_memory=pin_memory), - fp8_dtype=self.dtype, - requires_grad=requires_grad, - data_transpose=data_transpose, - quantizer=self, - device=device, - ) - def calibrate(self, tensor: torch.Tensor) -> None: # current scaling don't need to calibrate return diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 5cab519c7..134f8b5a6 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -96,70 +96,6 @@ def is_quantizable(self, inp: torch.Tensor) -> bool: return False return True - def make_empty( - self, - shape: Iterable[int], - *, - dtype: torch.dtype = torch.float32, - device: Optional[torch.device] = None, - requires_grad: bool = False, - pin_memory: bool = False, - ) -> MXFP8Tensor: - - # Canonicalize tensor attributes - if device is None: - device = torch.device("cuda") - - assert ( - shape[-1] % MXFP8_BLOCK_SCALING_SIZE == 0 - and math.prod(shape[:-1]) % MXFP8_BLOCK_SCALING_SIZE == 0 - ), ( - f"Incorrect shape {shape} for MXFP8. Tensor dims must be divisible by" - f" {MXFP8_BLOCK_SCALING_SIZE}" - ) - - # Allocate FP8 data - data = None - scale_inv = None - if self.rowwise_usage: - data = torch.empty(shape, dtype=torch.uint8, device=device, pin_memory=pin_memory) - scale_inv = torch.empty( - round_up_to_nearest_multiple(math.prod(shape[:-1]), 128), - round_up_to_nearest_multiple(shape[-1] // MXFP8_BLOCK_SCALING_SIZE, 4), - dtype=torch.uint8, - device=device, - pin_memory=pin_memory, - ) - - # Allocate FP8 data transpose if needed - columnwise_data = None - columnwise_scale_inv = None - if self.columnwise_usage: - columnwise_data = torch.empty( - shape, dtype=torch.uint8, device=device, pin_memory=pin_memory - ) - columnwise_scale_inv = torch.empty( - round_up_to_nearest_multiple(math.prod(shape[:-1]) // MXFP8_BLOCK_SCALING_SIZE, 4), - round_up_to_nearest_multiple(shape[-1], 128), - dtype=torch.uint8, - device=device, - pin_memory=pin_memory, - ) - - # Construct FP8 tensor - return MXFP8Tensor( - shape=shape, - dtype=dtype, - fp8_dtype=self.dtype, - rowwise_data=data, - rowwise_scale_inv=scale_inv, - columnwise_data=columnwise_data, - columnwise_scale_inv=columnwise_scale_inv, - quantizer=self, - requires_grad=requires_grad, - with_gemm_swizzled_scales=self.optimize_for_gemm, - ) - def calibrate(self, tensor: torch.Tensor) -> None: # TODO(ksivamani): No calibration needed for mxfp8? pass diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 285a7f030..df7a2b4bd 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -297,99 +297,6 @@ def convert_shape_for_fp4(shape: Iterable[int]) -> Tuple[int, ...]: shape[-1] = shape[-1] // 2 return tuple(shape) - def make_empty( - self, - shape: Iterable[int], - *, - dtype: torch.dtype = torch.float32, - device: Optional[torch.device] = None, - pin_memory: bool = False, - requires_grad: bool = False, - ) -> NVFP4Tensor: - - # Canonicalize tensor attributes - if device is None: - device = torch.device("cuda") - - assert shape[-1] % NVFP4_BLOCK_SCALING_SIZE == 0, ( - f"Incorrect shape {shape} for NVFP4. Tensor dims must be divisible by" - f" {NVFP4_BLOCK_SCALING_SIZE}" - ) - - flat_first_dim = math.prod(shape[:-1]) - assert flat_first_dim % NVFP4_BLOCK_SCALING_SIZE == 0, ( - f"Incorrect shape {shape} for NVFP4. Tensor dims must be divisible by" - f" {NVFP4_BLOCK_SCALING_SIZE}" - ) - if self.row_scaled_nvfp4: - if not self.rowwise_usage: - raise ValueError("Row-scaled NVFP4 quantization requires rowwise usage.") - if self.columnwise_usage: - raise ValueError("Row-scaled NVFP4 quantization does not support columnwise usage.") - - # Allocate FP4 data - data = None - scale_inv = None - amax_rowwise = None - if self.rowwise_usage: - data = torch.empty( - self.convert_shape_for_fp4(shape), - dtype=torch.uint8, - device=device, - pin_memory=pin_memory, - ) - scale_shape = self.get_scale_shape(shape, columnwise=False) - scale_inv = torch.empty( - scale_shape, dtype=torch.uint8, device=device, pin_memory=pin_memory - ) - # Allocate global amax metadata. Row-scaled NVFP4 stores one value per row. - amax_rows = flat_first_dim if self.row_scaled_nvfp4 else 1 - amax_rowwise = torch.zeros( - amax_rows, dtype=torch.float32, device=device, pin_memory=pin_memory - ) - - # Allocate FP8 data transpose if needed - columnwise_data = None - columnwise_scale_inv = None - amax_columnwise = None - if self.columnwise_usage: - # enforce 2D shape to avoid [S, B, H] shape and B and be 1 - # and the transposed shape is [H, S, B], so divide last dim by 2 gives zero - shape_2d = tuple([flat_first_dim, shape[-1]]) - columnwise_data = torch.empty( - self.convert_shape_for_fp4(self.get_columnwise_shape(shape_2d)), - dtype=torch.uint8, - device=device, - pin_memory=pin_memory, - ) - columnwise_scale_shape = self.get_scale_shape(shape, columnwise=True) - columnwise_scale_inv = torch.empty( - columnwise_scale_shape, - dtype=torch.uint8, - device=device, - pin_memory=pin_memory, - ) - amax_columnwise = torch.zeros( - 1, dtype=torch.float32, device=device, pin_memory=pin_memory - ) - - # Construct FP8 tensor - return NVFP4Tensor( - shape=shape, - dtype=dtype, - rowwise_data=data, - rowwise_scale_inv=scale_inv, - columnwise_data=columnwise_data, - columnwise_scale_inv=columnwise_scale_inv, - amax_rowwise=amax_rowwise, - amax_columnwise=amax_columnwise, - fp4_dtype=self.dtype, - quantizer=self, - requires_grad=requires_grad, - with_gemm_swizzled_scales=False, - row_scaled_nvfp4=self.row_scaled_nvfp4, - ) - def calibrate(self, tensor: torch.Tensor) -> None: pass # Calibration is no-op From 76c2a9e90275a6856039fbba9fcb09bc98c48605 Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Tue, 12 May 2026 17:04:08 -0700 Subject: [PATCH 060/180] Added the CODEOWNERS file (#2980) Signed-off-by: Przemek Tredak --- CODEOWNERS | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 CODEOWNERS diff --git a/CODEOWNERS b/CODEOWNERS new file mode 100644 index 000000000..3087832fa --- /dev/null +++ b/CODEOWNERS @@ -0,0 +1,24 @@ +# IMPORTANT: +# This file is ONLY used to subscribe for notifications for PRs +# related to a specific file path. Approvals from people in this +# file are not required for merges. + +# C API +/transformer_engine/common/include/ @ptrendx + +# TE/JAX +/transformer_engine/jax/ @jberchtold-nvidia + +# TE/PyTorch +/transformer_engine/pytorch/ @ksivaman + +# te.ops API +/transformer_engine/pytorch/ops/ @timmoon10 + +# Quantization kernels +/transformer_engine/common/cast/ @Oleg-Goncharov + +# Attention +/transformer_engine/pytorch/attention/ @cyanguwa +/transformer_engine/common/fused_attn/ @cyanguwa +/transformer_engine/jax/cpp_extensions/attention.py @KshitijLakhani From 4322c0ab1c37b845ec7fe71f7711713f5acf2f2e Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Thu, 14 May 2026 17:46:43 -0400 Subject: [PATCH 061/180] Remove `epel-release` package from wheel Dockerfiles (#2987) Remove epel-release package from wheel Dockerfiles Signed-off-by: Kirthi Shankar Sivamani --- build_tools/wheel_utils/Dockerfile.aarch | 1 - build_tools/wheel_utils/Dockerfile.x86 | 3 +-- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/build_tools/wheel_utils/Dockerfile.aarch b/build_tools/wheel_utils/Dockerfile.aarch index 8c5b81d92..c040dadcd 100644 --- a/build_tools/wheel_utils/Dockerfile.aarch +++ b/build_tools/wheel_utils/Dockerfile.aarch @@ -23,7 +23,6 @@ ENV CUDA_MAJOR=${CUDA_MAJOR} # Cuda toolkit, cudnn, driver. RUN dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/sbsa/cuda-rhel8.repo -RUN dnf -y install epel-release RUN dnf -y install cuda-compiler-${CUDA_MAJOR}-${CUDA_MINOR}.aarch64 \ cuda-libraries-${CUDA_MAJOR}-${CUDA_MINOR}.aarch64 \ cuda-libraries-devel-${CUDA_MAJOR}-${CUDA_MINOR}.aarch64 diff --git a/build_tools/wheel_utils/Dockerfile.x86 b/build_tools/wheel_utils/Dockerfile.x86 index b77920250..2728b6b7c 100644 --- a/build_tools/wheel_utils/Dockerfile.x86 +++ b/build_tools/wheel_utils/Dockerfile.x86 @@ -23,7 +23,6 @@ ENV CUDA_MAJOR=${CUDA_MAJOR} # Cuda toolkit, cudnn, driver. RUN dnf config-manager --add-repo https://developer.download.nvidia.com/compute/cuda/repos/rhel8/x86_64/cuda-rhel8.repo -RUN dnf -y install epel-release RUN dnf -y install cuda-compiler-${CUDA_MAJOR}-${CUDA_MINOR}.x86_64 \ cuda-libraries-${CUDA_MAJOR}-${CUDA_MINOR}.x86_64 \ cuda-libraries-devel-${CUDA_MAJOR}-${CUDA_MINOR}.x86_64 @@ -44,4 +43,4 @@ ENV CUDA_PATH=/usr/local/cuda ENV CUDADIR=/usr/local/cuda ENV NVTE_RELEASE_BUILD=1 -CMD ["/bin/bash", "-c", "bash /TransformerEngine/build_tools/wheel_utils/build_wheels.sh manylinux_2_28_x86_64 $BUILD_METAPACKAGE $BUILD_COMMON $BUILD_PYTORCH $BUILD_JAX $CUDA_MAJOR"] \ No newline at end of file +CMD ["/bin/bash", "-c", "bash /TransformerEngine/build_tools/wheel_utils/build_wheels.sh manylinux_2_28_x86_64 $BUILD_METAPACKAGE $BUILD_COMMON $BUILD_PYTORCH $BUILD_JAX $CUDA_MAJOR"] From c40398c4996c1d5f8a1e5a1b12db2978d973ca8d Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Thu, 14 May 2026 16:22:40 -0700 Subject: [PATCH 062/180] [JAX] Size autotuned Triton grids per config (#2975) * [JAX] Size autotuned Triton grids per config (3x perm-kernel speedup) The autotuned path in triton_call_lowering compiled all BLOCK_SIZE configs but dispatched every one with the same fixed grid sized for the smallest BLOCK_SIZE, so larger configs over-launched by the BLOCK_SIZE ratio. Make grid accept a callable(meta)->tuple evaluated per config, matching the jax-triton API. Update _permute_kernel, _unpermute_kernel, and _sort_chunks_by_map_kernel lowerings. Measured 22.6ms -> 7.4ms (3.06x) on GB200 for sort_chunks at 524k tokens, hidden=4096, fp32. * [JAX] Triton wrapper defaults match jax-triton (3.25ms speedup) num_warps default 32->4 and num_stages 1->3 in triton_call_lowering match Triton's own triton.Config defaults. Non-autotuned kernels (e.g. _make_chunk_sort_map_kernel) were running with 1024 threads/block, an 8x kernel slowdown. Also: tuple/callable grid assertion + comment trims. Signed-off-by: tdophung --- .../jax/triton_extensions/permutation.py | 25 +++-- .../jax/triton_extensions/utils.py | 93 +++++++++++++------ 2 files changed, 84 insertions(+), 34 deletions(-) diff --git a/transformer_engine/jax/triton_extensions/permutation.py b/transformer_engine/jax/triton_extensions/permutation.py index 98c54e52b..22f983f07 100644 --- a/transformer_engine/jax/triton_extensions/permutation.py +++ b/transformer_engine/jax/triton_extensions/permutation.py @@ -589,10 +589,13 @@ def lowering( probs_stride_token = 0 probs_stride_expert = 0 - # Grid function equivalent: (num_tokens, cdiv(hidden_size, BLOCK_SIZE)) - # Use minimum BLOCK_SIZE from autotune configs to ensure grid covers all elements + # We use BLOCK_SIZE in the grid calculation to ensure the grid is the + # proper size. If the grid size is an overestimate it can significantly + # hurt performance. + def grid(meta): + return (num_tokens, triton.cdiv(hidden_size, meta["BLOCK_SIZE"])) + block_size = _get_min_block_size(_permute_kernel) - grid = (num_tokens, triton.cdiv(hidden_size, block_size)) # Use input_output_aliases to alias pre-zeroed buffers to outputs. # This ensures padding positions contain zeros since the kernel only writes valid positions. @@ -997,9 +1000,13 @@ def lowering( unpermuted_probs_stride_token = num_experts unpermuted_probs_stride_expert = 1 - # Grid - use minimum BLOCK_SIZE from autotune configs + # We use BLOCK_SIZE in the grid calculation to ensure the grid is the + # proper size. If the grid size is an overestimate it can significantly + # hurt performance. + def grid(meta): + return (num_tokens, triton.cdiv(hidden_size, meta["BLOCK_SIZE"])) + block_size = _get_min_block_size(_unpermute_kernel) - grid = (num_tokens, triton.cdiv(hidden_size, block_size)) return triton_call_lowering( ctx, @@ -1720,9 +1727,13 @@ def lowering( probs_stride_token = 1 permuted_probs_stride_token = 1 - # Grid - use minimum BLOCK_SIZE from autotune configs + # We use BLOCK_SIZE in the grid calculation to ensure the grid is the + # proper size. If the grid size is an overestimate it can significantly + # hurt performance. + def grid(meta): + return (num_tokens, triton.cdiv(hidden_size, meta["BLOCK_SIZE"])) + block_size = _get_min_block_size(_sort_chunks_by_map_kernel) - grid = (num_tokens, triton.cdiv(hidden_size, block_size)) # Declare input_output_aliases so XLA knows output slot 0 is claimed by # input 3 (output_buf). This prevents XLA from implicitly aliasing any diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 2a86321c3..332bc6ddb 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -390,7 +390,16 @@ def triton_call_lowering( ctx: MLIR lowering context kernel_fn: Triton kernel function *array_args: Input arrays (from ctx) - grid: Grid dimensions (int or tuple) + grid: Grid dimensions. May be either: + - an int or tuple (fixed grid for every config), or + - a callable ``meta -> int|tuple`` (evaluated per autotune config). + + Use the callable form for autotuned kernels whose grid depends on + ``BLOCK_SIZE`` (or any other autotuned constexpr); otherwise the + launch grid will not match the autotuner-selected config and the + kernel will either over-launch (waste) or under-cover. ``meta`` is + the merged dict ``{**constexprs, **config.kwargs}`` for the chosen + config — the same convention as jax-triton's ``triton_call``. input_output_aliases: Mapping of input to output aliases constexprs: Compile-time constants for the kernel. This includes both tl.constexpr arguments AND scalar runtime arguments (like @@ -404,13 +413,12 @@ def triton_call_lowering( def lowering(ctx, x, *, block_size): from ..triton_extensions import triton_call_lowering n = ctx.avals_in[0].size + + def grid(meta): + return (triton.cdiv(n, meta["BLOCK_SIZE"]),) + return triton_call_lowering( - ctx, my_kernel, x, - grid=(triton.cdiv(n, block_size),), - constexprs={ - "n_elements": n, # scalar arg (not tl.constexpr in kernel) - "BLOCK_SIZE": block_size, # tl.constexpr arg - }, + ctx, my_kernel, x, grid=grid, constexprs={"n_elements": n}, ) """ # Get compute capability using gpu_triton @@ -431,22 +439,39 @@ def lowering(ctx, x, *, block_size): tensor_arg_names = [n for n in arg_names if n not in constexpr_names] signature = {n: get_triton_dtype(a) for n, a in zip(tensor_arg_names, all_avals)} - # Normalize grid to 3D - if isinstance(grid, int): - grid_tuple = (grid, 1, 1) - elif len(grid) == 1: - grid_tuple = (grid[0], 1, 1) - elif len(grid) == 2: - grid_tuple = (grid[0], grid[1], 1) - else: - grid_tuple = grid[:3] + assert callable(grid) or isinstance(grid, tuple), ( + "Argument 'grid' must be a tuple or a callable but received: " + f"type={type(grid)}, value={grid}" + ) - # Default values for the kernel + # Normalize grid to 3D. When `grid` is a callable, defer evaluation until + # we know the per-config meta (so each autotune config gets its own grid, + # matching jax-triton's behavior). + def _normalize_grid(grid_tuple): + if isinstance(grid_tuple, int): + return (grid_tuple, 1, 1) + if len(grid_tuple) == 1: + return (grid_tuple[0], 1, 1) + if len(grid_tuple) == 2: + return (grid_tuple[0], grid_tuple[1], 1) + return tuple(grid_tuple[:3]) + + grid_callable = grid if callable(grid) else None + if grid_callable is None: + grid_tuple = _normalize_grid(grid) + else: + grid_tuple = None # evaluated per-config below + + # Default kernel launch parameters. These apply to non-autotuned kernels + # and as a fallback when an autotuned config doesn't specify them. Values + # match Triton's own `triton.Config` defaults (num_warps=4, num_stages=3, + # num_ctas=1) and jax-triton's `get_or_create_triton_kernel`. Using a + # larger default (e.g. num_warps=32) over-provisions threads per block, + # which slashes SM occupancy on non-autotuned kernels — measured as an 8× + # slowdown on `_make_chunk_sort_map_kernel` vs jax-triton. actual_kernel_fn = kernel_fn - num_warps = 32 - num_stages = ( - 1 # TODO(Phuong): consider if it is beneficial to expose num_warps, num_stages, num_ctas - ) + num_warps = 4 + num_stages = 3 num_ctas = 1 kernel_constexprs = constexprs if constexprs is not None else {} @@ -510,11 +535,18 @@ def lowering(ctx, x, *, block_size): for _ in list(ctx.avals_in) + list(ctx.avals_out): config_params.append(gpu_triton.create_array_parameter(0, 16)) + # Per-config grid: evaluate `grid(meta)` if grid is a callable so + # the launch shape matches this config's BLOCK_SIZE (etc.). + if grid_callable is not None: + config_grid = _normalize_grid(grid_callable(config_constexprs)) + else: + config_grid = grid_tuple + config_call = gpu_triton.TritonKernelCall( config_kernel, - grid_tuple[0], - grid_tuple[1], - grid_tuple[2], + config_grid[0], + config_grid[1], + config_grid[2], config_params, ) @@ -571,11 +603,18 @@ def lowering(ctx, x, *, block_size): for _ in list(ctx.avals_in) + list(ctx.avals_out): kernel_params.append(gpu_triton.create_array_parameter(0, 16)) + # Non-autotuned dispatch: evaluate `grid(meta)` once with the merged + # constexprs (which already reflect the single config we'll launch). + if grid_callable is not None: + single_grid = _normalize_grid(grid_callable(kernel_constexprs)) + else: + single_grid = grid_tuple + kernel_call = gpu_triton.TritonKernelCall( kernel, - grid_tuple[0], - grid_tuple[1], - grid_tuple[2], + single_grid[0], + single_grid[1], + single_grid[2], kernel_params, ) From eca05d3b554e36d99106c368f45df2cc350ddebc Mon Sep 17 00:00:00 2001 From: Arpit Jain <3242828+arpitjain099@users.noreply.github.com> Date: Fri, 15 May 2026 08:53:54 +0900 Subject: [PATCH 063/180] ci: declare contents:read on Lint workflow (#2989) The Lint workflow runs cpplint and pylint against the checked-out tree. No cache, no GitHub API write. `permissions: contents: read` captures that and matches the per-job permissions blocks already used in deploy_nightly_docs.yml (pages:write + id-token:write) and upload-ci-logs.yml (statuses:write). build.yml is left out because it pulls mozilla-actions/sccache-action (which writes to the Actions cache) and easimon/maximize-build-space. A drive-by permissions block there would need actions:write for the sccache save path, which deserves a separate look. Signed-off-by: Arpit Jain --- .github/workflows/lint.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 1d2fb272f..016d2079d 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -11,6 +11,8 @@ concurrency: # Group by workflow name + PR number (for PRs) or ref (for branch/tag pushes) group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} cancel-in-progress: true +permissions: + contents: read jobs: pytorch_cpplint: name: 'PyTorch C++' From 583d2d12b1dd7515c4b1f261f1c3501a52fc6b1b Mon Sep 17 00:00:00 2001 From: Przemek Tredak Date: Mon, 18 May 2026 09:52:13 -0700 Subject: [PATCH 064/180] Changed VERSION to 2.17.0.dev0 Signed-off-by: Przemek Tredak --- build_tools/VERSION.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/build_tools/VERSION.txt b/build_tools/VERSION.txt index 36334f690..73198da8a 100644 --- a/build_tools/VERSION.txt +++ b/build_tools/VERSION.txt @@ -1 +1 @@ -2.16.0.dev0 +2.17.0.dev0 From ca50bbf9ba9194465bf704fa7f7a711c33c5985b Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Tue, 19 May 2026 13:35:40 -0400 Subject: [PATCH 065/180] Add license to framework sdist builds (#3002) Signed-off-by: ksivamani --- transformer_engine/jax/setup.py | 8 ++++++++ transformer_engine/pytorch/setup.py | 8 ++++++++ 2 files changed, 16 insertions(+) diff --git a/transformer_engine/jax/setup.py b/transformer_engine/jax/setup.py index 2d2524282..678062df9 100644 --- a/transformer_engine/jax/setup.py +++ b/transformer_engine/jax/setup.py @@ -42,6 +42,11 @@ shutil.rmtree(build_tools_copy) shutil.copytree(build_tools_dir, build_tools_copy) +license_src = current_file_path.parent.parent / "LICENSE" +license_dst = current_file_path / "LICENSE" +if license_src.is_file(): + shutil.copyfile(license_src, license_dst) + from build_tools.build_ext import get_build_ext from build_tools.utils import copy_common_headers, min_python_version_str @@ -131,7 +136,10 @@ def get_cuda_major_version() -> int: python_requires=f">={min_python_version_str()}", install_requires=install_requires, tests_require=test_requirements(), + license_files=("LICENSE",), ) if any(x in sys.argv for x in (".", "sdist", "bdist_wheel")): shutil.rmtree(common_headers_dir) shutil.rmtree("build_tools") + if license_dst.is_file(): + license_dst.unlink() diff --git a/transformer_engine/pytorch/setup.py b/transformer_engine/pytorch/setup.py index 99f6a99ef..593a3169d 100644 --- a/transformer_engine/pytorch/setup.py +++ b/transformer_engine/pytorch/setup.py @@ -43,6 +43,11 @@ shutil.rmtree(build_tools_copy) shutil.copytree(build_tools_dir, build_tools_copy) +license_src = current_file_path.parent.parent / "LICENSE" +license_dst = current_file_path / "LICENSE" +if license_src.is_file(): + shutil.copyfile(license_src, license_dst) + from build_tools.build_ext import get_build_ext from build_tools.utils import copy_common_headers, min_python_version_str @@ -177,7 +182,10 @@ def run(self): python_requires=f">={min_python_version_str()}", install_requires=install_requires, tests_require=test_requirements(), + license_files=("LICENSE",), ) if any(x in sys.argv for x in (".", "sdist", "bdist_wheel")): shutil.rmtree(common_headers_dir) shutil.rmtree("build_tools") + if license_dst.is_file(): + license_dst.unlink() From b629e6e54cb3197927c0799c5aeca7537adebe68 Mon Sep 17 00:00:00 2001 From: Shaurya Singh Date: Tue, 19 May 2026 11:02:55 -0700 Subject: [PATCH 066/180] docs: fix comm GEMM overlap README typos (#3010) Signed-off-by: LeSingh1 --- examples/pytorch/comm_gemm_overlap/README.md | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/examples/pytorch/comm_gemm_overlap/README.md b/examples/pytorch/comm_gemm_overlap/README.md index fc8458844..b7ecb2d06 100644 --- a/examples/pytorch/comm_gemm_overlap/README.md +++ b/examples/pytorch/comm_gemm_overlap/README.md @@ -6,7 +6,7 @@ - `CUDA_DEVICE_MAX_CONNECTIONS=1` must be enabled in the environment. - For best performance, point-to-point communication via _CUDA Multicast_ needs CUDA Toolkit 12.0+ and CUDA driver 535+ on devices with compute capability 9.0 or newer. -- Devices older than compute capability 9.0 require `UB_SKIPMC=1` in the environment in order fall +- Devices older than compute capability 9.0 require `UB_SKIPMC=1` in the environment in order to fall back on a less performant implementation based on CUDA Inter-Process Communication (IPC) handles. ## Examples @@ -22,7 +22,7 @@ $ torchrun --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) te_layer_with_ov # [rank0:node0] |-- Created tensor-parallel group: [0, 1, 2, 3, 4, 5, 6, 7] # !!! [UB] Create UbufP2PCommOverlap Communicator # UB_TIMEOUT is set to 110 sec, 217800000000 cycles, freq: 1980000khz -# MC initialized succesfully, window size = 549755813888 +# MC initialized successfully, window size = 549755813888 # !!! [UBP2P] Register UBuf 1 # !!! [UBP2P] Register UBuf 2 # !!! [UBP2P] Register UBuf 3 @@ -66,7 +66,7 @@ $ torchrun --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) te_layer_with_ov ``` ### Single node, mixed data- and tensor-parallel LayerNormMLP: -Uses `torch.nn.parallel.DistributedDataParallel` for replicatin the model across 2 tensor-parallel +Uses `torch.nn.parallel.DistributedDataParallel` for replicating the model across 2 tensor-parallel groups in a single node. ```bash @@ -81,7 +81,7 @@ $ torchrun --nnodes=1 --nproc-per-node=$(nvidia-smi -L | wc -l) te_layer_with_ov # [rank2:node0] |-- Created data-parallel group: [2, 6] # !!! [UB] Create UbufP2PCommOverlap Communicator # UB_TIMEOUT is set to 110 sec, 217800000000 cycles, freq: 1980000khz -# MC initialized succesfully, window size = 549755813888 +# MC initialized successfully, window size = 549755813888 # !!! [UBP2P] Register UBuf 1 # !!! [UBP2P] Register UBuf 2 # !!! [UBP2P] Register UBuf 3 From 50ac303f788ef44a0f8b9fc06e1cc3847c7e3f42 Mon Sep 17 00:00:00 2001 From: Kirthi Shankar Sivamani Date: Wed, 20 May 2026 13:55:59 -0400 Subject: [PATCH 067/180] Update `cudnn-frontend` to 1.23.0 (#3003) * Update cudnn-FE to 1.23.0 Signed-off-by: ksivamani * Point to correct commit Signed-off-by: ksivamani --------- Signed-off-by: ksivamani --- 3rdparty/cudnn-frontend | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend index 97f6cb3b8..fb682ce76 160000 --- a/3rdparty/cudnn-frontend +++ b/3rdparty/cudnn-frontend @@ -1 +1 @@ -Subproject commit 97f6cb3b88cacff507cca1280db5650a457d92b3 +Subproject commit fb682ce761a2705e40f9b5d528737a3e0eb33cec From a12f7aade84b84efb8cab8a9220a6a0141314754 Mon Sep 17 00:00:00 2001 From: francesco-bertolotti Date: Wed, 20 May 2026 22:51:45 +0200 Subject: [PATCH 068/180] mnnvl guard (#3013) * guarding nvmlGpuFabricInfo_v2 Signed-off-by: Francesco Bertolotti * precision errors Signed-off-by: Francesco Bertolotti * reverting NVIDIA_TF32_OVERRIDE=0 Signed-off-by: Francesco Bertolotti --------- Signed-off-by: Francesco Bertolotti --- .../common/comm_gemm_overlap/userbuffers/userbuffers-host.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers-host.cpp b/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers-host.cpp index 1dcde51d4..c8d5977fb 100644 --- a/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers-host.cpp +++ b/transformer_engine/common/comm_gemm_overlap/userbuffers/userbuffers-host.cpp @@ -92,7 +92,7 @@ int stringCmp(const void *a, const void *b) { return strcmp((const char *)a, (co } while (0); bool has_mnnvl_fabric(int device_id) { -#if CUDA_VERSION < 12040 +#if !defined(nvmlGpuFabricInfo_v2) if (getenv("NVTE_UBDEBUG")) { printf( "TransformerEngine does not support multi-node NVLINK " From aab7bc947bb09a72a9b068f3205c28dcba6e7064 Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Wed, 20 May 2026 16:34:40 -0700 Subject: [PATCH 069/180] Add GitHub actions to automatically mark community contributions (#3007) * Add GitHub actions to automatically mark community contributions Signed-off-by: Przemek Tredak * Fix Signed-off-by: Przemek Tredak * Use exact commit hash Signed-off-by: Przemek Tredak * Remove unnecessary indentation Signed-off-by: Przemek Tredak --------- Signed-off-by: Przemek Tredak --- .github/workflows/community_label.yml | 65 +++++++++++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 .github/workflows/community_label.yml diff --git a/.github/workflows/community_label.yml b/.github/workflows/community_label.yml new file mode 100644 index 000000000..52afa2477 --- /dev/null +++ b/.github/workflows/community_label.yml @@ -0,0 +1,65 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +# A workflow to automatically label the contributions as community/org +name: Label community contributions + +on: + pull_request_target: + types: [opened, reopened, ready_for_review, synchronize] + +permissions: + contents: read + issues: write + +jobs: + label: + runs-on: ubuntu-latest + steps: + - uses: actions/github-script@3a2844b7e9c422d3c10d287c895573f7108da1b3 + with: + script: | + const pr = context.payload.pull_request; + const user = pr.user.login; + const association = pr.author_association; + + const communityLabel = "community-contribution"; + const orgLabel = "org-contribution"; + + let targetLabel = null; + + const isOrgMember = + association === "MEMBER" || association === "OWNER"; + + if (!isOrgMember) { + targetLabel = communityLabel; + } else { + let permission = "none"; + + try { + const res = await github.rest.repos.getCollaboratorPermissionLevel({ + owner: context.repo.owner, + repo: context.repo.repo, + username: user, + }); + permission = res.data.permission; + } catch (e) { + if (e.status !== 404) throw e; + } + + const isCore = permission === "write" || permission === "admin"; + + if (!isCore) { + targetLabel = orgLabel; + } + } + + if (targetLabel) { + await github.rest.issues.addLabels({ + owner: context.repo.owner, + repo: context.repo.repo, + issue_number: pr.number, + labels: [targetLabel], + }); + } From a01430023cf494e3f865dd37c69e79c4400f06d1 Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Wed, 20 May 2026 17:25:38 -0700 Subject: [PATCH 070/180] Split grouped quantize/activations and dbias for faster compilation on multicore machines (#2983) * Split grouped and dbias CUDA wrappers Signed-off-by: Przemek Tredak * Fix Signed-off-by: Przemek Tredak * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Przemek Tredak Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- transformer_engine/common/CMakeLists.txt | 23 ++++- transformer_engine/common/activation/gelu.cu | 99 ------------------- .../common/activation/gelu_dbias.cu | 34 +++++++ .../common/activation/gelu_grouped.cu | 53 ++++++++++ .../common/activation/gelu_grouped_dbias.cu | 36 +++++++ transformer_engine/common/activation/relu.cu | 99 ------------------- .../common/activation/relu_dbias.cu | 34 +++++++ .../common/activation/relu_grouped.cu | 53 ++++++++++ .../common/activation/relu_grouped_dbias.cu | 36 +++++++ .../common/activation/swiglu.cu | 49 --------- .../common/activation/swiglu_dbias.cu | 21 ++++ .../common/activation/swiglu_grouped.cu | 30 ++++++ .../common/activation/swiglu_grouped_dbias.cu | 22 +++++ transformer_engine/common/cast/cast.cu | 59 ----------- transformer_engine/common/cast/cast_dbias.cu | 24 +++++ .../common/cast/cast_grouped.cu | 47 +++++++++ .../common/cast/cast_grouped_dbias.cu | 24 +++++ .../common/cast/core/common.cuh | 1 + 18 files changed, 437 insertions(+), 307 deletions(-) create mode 100644 transformer_engine/common/activation/gelu_dbias.cu create mode 100644 transformer_engine/common/activation/gelu_grouped.cu create mode 100644 transformer_engine/common/activation/gelu_grouped_dbias.cu create mode 100644 transformer_engine/common/activation/relu_dbias.cu create mode 100644 transformer_engine/common/activation/relu_grouped.cu create mode 100644 transformer_engine/common/activation/relu_grouped_dbias.cu create mode 100644 transformer_engine/common/activation/swiglu_dbias.cu create mode 100644 transformer_engine/common/activation/swiglu_grouped.cu create mode 100644 transformer_engine/common/activation/swiglu_grouped_dbias.cu create mode 100644 transformer_engine/common/cast/cast_dbias.cu create mode 100644 transformer_engine/common/cast/cast_grouped.cu create mode 100644 transformer_engine/common/cast/cast_grouped_dbias.cu diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt index 030023d94..06d85b6d8 100644 --- a/transformer_engine/common/CMakeLists.txt +++ b/transformer_engine/common/CMakeLists.txt @@ -212,10 +212,22 @@ list(APPEND transformer_engine_cuda_sources list(APPEND transformer_engine_cuda_arch_specific_sources fused_attn/flash_attn.cu activation/gelu.cu + activation/gelu_dbias.cu + activation/gelu_grouped.cu + activation/gelu_grouped_dbias.cu activation/glu.cu activation/relu.cu + activation/relu_dbias.cu + activation/relu_grouped.cu + activation/relu_grouped_dbias.cu activation/swiglu.cu + activation/swiglu_dbias.cu + activation/swiglu_grouped.cu + activation/swiglu_grouped_dbias.cu cast/cast.cu + cast/cast_dbias.cu + cast/cast_grouped.cu + cast/cast_grouped_dbias.cu gemm/cutlass_grouped_gemm.cu hadamard_transform/group_hadamard_transform.cu hadamard_transform/graph_safe_group_hadamard_transform.cu @@ -447,9 +459,18 @@ list(APPEND nvte_sources_with_fast_math fused_softmax/scaled_masked_softmax.cu option(NVTE_BUILD_ACTIVATION_WITH_FAST_MATH "Compile activation kernels with --use_fast_math option" OFF) if (NVTE_BUILD_ACTIVATION_WITH_FAST_MATH) list(APPEND nvte_sources_with_fast_math activation/gelu.cu + activation/gelu_dbias.cu + activation/gelu_grouped.cu + activation/gelu_grouped_dbias.cu activation/glu.cu activation/relu.cu - activation/swiglu.cu) + activation/relu_dbias.cu + activation/relu_grouped.cu + activation/relu_grouped_dbias.cu + activation/swiglu.cu + activation/swiglu_dbias.cu + activation/swiglu_grouped.cu + activation/swiglu_grouped_dbias.cu) endif() foreach(cuda_source IN LISTS nvte_sources_with_fast_math) diff --git a/transformer_engine/common/activation/gelu.cu b/transformer_engine/common/activation/gelu.cu index ea864813b..6bd63672c 100644 --- a/transformer_engine/common/activation/gelu.cu +++ b/transformer_engine/common/activation/gelu.cu @@ -13,14 +13,6 @@ void nvte_gelu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { act_fn>(input, output, stream); } -void nvte_group_gelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_gelu); - using namespace transformer_engine; - constexpr bool IS_ACT = true; - dispatch::group_quantize_fwd_helper>(input, output, nullptr, - stream); -} - void nvte_dgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_dgelu); @@ -28,47 +20,6 @@ void nvte_dgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output dact_fn>(grad, input, output, stream); } -void nvte_group_dgelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, - NVTEGroupedTensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_dgelu); - using namespace transformer_engine; - NVTEGroupedTensor dbias = nullptr; - NVTETensor workspace = nullptr; - - constexpr bool IS_DBIAS = false; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - grad, input, output, dbias, workspace, nullptr, stream); -} - -void nvte_quantize_dbias_dgelu(const NVTETensor input, const NVTETensor activation_input, - NVTETensor output, NVTETensor dbias, NVTETensor workspace, - cudaStream_t stream) { - NVTE_API_CALL(nvte_quantize_dbias_dgelu); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = true; - - dispatch::quantize_bwd_helper>( - input, activation_input, output, dbias, workspace, nullptr, stream); -} - -void nvte_group_quantize_dbias_dgelu(const NVTEGroupedTensor input, - const NVTEGroupedTensor activation_input, - NVTEGroupedTensor output, NVTEGroupedTensor dbias, - NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_quantize_dbias_dgelu); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - input, activation_input, output, dbias, workspace, nullptr, stream); -} - void nvte_geglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_geglu); using namespace transformer_engine; @@ -90,15 +41,6 @@ void nvte_qgelu(const NVTETensor input, NVTETensor output, cudaStream_t stream) act_fn>(input, output, stream); } -void nvte_group_qgelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, - cudaStream_t stream) { - NVTE_API_CALL(nvte_group_qgelu); - using namespace transformer_engine; - constexpr bool IS_ACT = true; - dispatch::group_quantize_fwd_helper>(input, output, nullptr, - stream); -} - void nvte_dqgelu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_dqgelu); @@ -106,47 +48,6 @@ void nvte_dqgelu(const NVTETensor grad, const NVTETensor input, NVTETensor outpu dact_fn>(grad, input, output, stream); } -void nvte_group_dqgelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, - NVTEGroupedTensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_dqgelu); - using namespace transformer_engine; - NVTEGroupedTensor dbias = nullptr; - NVTETensor workspace = nullptr; - - constexpr bool IS_DBIAS = false; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - grad, input, output, dbias, workspace, nullptr, stream); -} - -void nvte_quantize_dbias_dqgelu(const NVTETensor input, const NVTETensor activation_input, - NVTETensor output, NVTETensor dbias, NVTETensor workspace, - cudaStream_t stream) { - NVTE_API_CALL(nvte_quantize_dbias_dqgelu); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = true; - - dispatch::quantize_bwd_helper>( - input, activation_input, output, dbias, workspace, nullptr, stream); -} - -void nvte_group_quantize_dbias_dqgelu(const NVTEGroupedTensor input, - const NVTEGroupedTensor activation_input, - NVTEGroupedTensor output, NVTEGroupedTensor dbias, - NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_quantize_dbias_dqgelu); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - input, activation_input, output, dbias, workspace, nullptr, stream); -} - void nvte_qgeglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_qgeglu); using namespace transformer_engine; diff --git a/transformer_engine/common/activation/gelu_dbias.cu b/transformer_engine/common/activation/gelu_dbias.cu new file mode 100644 index 000000000..4eaa9e355 --- /dev/null +++ b/transformer_engine/common/activation/gelu_dbias.cu @@ -0,0 +1,34 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_quantize_dbias_dgelu(const NVTETensor input, const NVTETensor activation_input, + NVTETensor output, NVTETensor dbias, NVTETensor workspace, + cudaStream_t stream) { + NVTE_API_CALL(nvte_quantize_dbias_dgelu); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = true; + + dispatch::quantize_bwd_helper>( + input, activation_input, output, dbias, workspace, nullptr, stream); +} + +void nvte_quantize_dbias_dqgelu(const NVTETensor input, const NVTETensor activation_input, + NVTETensor output, NVTETensor dbias, NVTETensor workspace, + cudaStream_t stream) { + NVTE_API_CALL(nvte_quantize_dbias_dqgelu); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = true; + + dispatch::quantize_bwd_helper>( + input, activation_input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/activation/gelu_grouped.cu b/transformer_engine/common/activation/gelu_grouped.cu new file mode 100644 index 000000000..c3267356f --- /dev/null +++ b/transformer_engine/common/activation/gelu_grouped.cu @@ -0,0 +1,53 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_gelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_gelu); + using namespace transformer_engine; + constexpr bool IS_ACT = true; + dispatch::group_quantize_fwd_helper>(input, output, nullptr, + stream); +} + +void nvte_group_dgelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_dgelu); + using namespace transformer_engine; + NVTEGroupedTensor dbias = nullptr; + NVTETensor workspace = nullptr; + + constexpr bool IS_DBIAS = false; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + grad, input, output, dbias, workspace, nullptr, stream); +} + +void nvte_group_qgelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_group_qgelu); + using namespace transformer_engine; + constexpr bool IS_ACT = true; + dispatch::group_quantize_fwd_helper>(input, output, nullptr, + stream); +} + +void nvte_group_dqgelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_dqgelu); + using namespace transformer_engine; + NVTEGroupedTensor dbias = nullptr; + NVTETensor workspace = nullptr; + + constexpr bool IS_DBIAS = false; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + grad, input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/activation/gelu_grouped_dbias.cu b/transformer_engine/common/activation/gelu_grouped_dbias.cu new file mode 100644 index 000000000..e8b549f69 --- /dev/null +++ b/transformer_engine/common/activation/gelu_grouped_dbias.cu @@ -0,0 +1,36 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_quantize_dbias_dgelu(const NVTEGroupedTensor input, + const NVTEGroupedTensor activation_input, + NVTEGroupedTensor output, NVTEGroupedTensor dbias, + NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_quantize_dbias_dgelu); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + input, activation_input, output, dbias, workspace, nullptr, stream); +} + +void nvte_group_quantize_dbias_dqgelu(const NVTEGroupedTensor input, + const NVTEGroupedTensor activation_input, + NVTEGroupedTensor output, NVTEGroupedTensor dbias, + NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_quantize_dbias_dqgelu); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + input, activation_input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/activation/relu.cu b/transformer_engine/common/activation/relu.cu index fc9122b7e..57222262f 100644 --- a/transformer_engine/common/activation/relu.cu +++ b/transformer_engine/common/activation/relu.cu @@ -13,14 +13,6 @@ void nvte_relu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { act_fn>(input, output, stream); } -void nvte_group_relu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_relu); - using namespace transformer_engine; - constexpr bool IS_ACT = true; - dispatch::group_quantize_fwd_helper>(input, output, nullptr, - stream); -} - void nvte_drelu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_drelu); @@ -28,47 +20,6 @@ void nvte_drelu(const NVTETensor grad, const NVTETensor input, NVTETensor output dact_fn>(grad, input, output, stream); } -void nvte_group_drelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, - NVTEGroupedTensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_drelu); - using namespace transformer_engine; - NVTEGroupedTensor dbias = nullptr; - NVTETensor workspace = nullptr; - - constexpr bool IS_DBIAS = false; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - grad, input, output, dbias, workspace, nullptr, stream); -} - -void nvte_quantize_dbias_drelu(const NVTETensor input, const NVTETensor activation_input, - NVTETensor output, NVTETensor dbias, NVTETensor workspace, - cudaStream_t stream) { - NVTE_API_CALL(nvte_quantize_dbias_drelu); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = true; - - dispatch::quantize_bwd_helper>( - input, activation_input, output, dbias, workspace, nullptr, stream); -} - -void nvte_group_quantize_dbias_drelu(const NVTEGroupedTensor input, - const NVTEGroupedTensor activation_input, - NVTEGroupedTensor output, NVTEGroupedTensor dbias, - NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_quantize_dbias_drelu); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - input, activation_input, output, dbias, workspace, nullptr, stream); -} - void nvte_reglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_reglu); using namespace transformer_engine; @@ -90,15 +41,6 @@ void nvte_srelu(const NVTETensor input, NVTETensor output, cudaStream_t stream) act_fn>(input, output, stream); } -void nvte_group_srelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, - cudaStream_t stream) { - NVTE_API_CALL(nvte_group_srelu); - using namespace transformer_engine; - constexpr bool IS_ACT = true; - dispatch::group_quantize_fwd_helper>(input, output, nullptr, - stream); -} - void nvte_dsrelu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_dsrelu); @@ -106,47 +48,6 @@ void nvte_dsrelu(const NVTETensor grad, const NVTETensor input, NVTETensor outpu dact_fn>(grad, input, output, stream); } -void nvte_group_dsrelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, - NVTEGroupedTensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_dsrelu); - using namespace transformer_engine; - NVTEGroupedTensor dbias = nullptr; - NVTETensor workspace = nullptr; - - constexpr bool IS_DBIAS = false; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - grad, input, output, dbias, workspace, nullptr, stream); -} - -void nvte_quantize_dbias_dsrelu(const NVTETensor input, const NVTETensor activation_input, - NVTETensor output, NVTETensor dbias, NVTETensor workspace, - cudaStream_t stream) { - NVTE_API_CALL(nvte_quantize_dbias_dsrelu); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = true; - - dispatch::quantize_bwd_helper>( - input, activation_input, output, dbias, workspace, nullptr, stream); -} - -void nvte_group_quantize_dbias_dsrelu(const NVTEGroupedTensor input, - const NVTEGroupedTensor activation_input, - NVTEGroupedTensor output, NVTEGroupedTensor dbias, - NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_quantize_dbias_dsrelu); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - input, activation_input, output, dbias, workspace, nullptr, stream); -} - void nvte_sreglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_sreglu); using namespace transformer_engine; diff --git a/transformer_engine/common/activation/relu_dbias.cu b/transformer_engine/common/activation/relu_dbias.cu new file mode 100644 index 000000000..bd14dc6c9 --- /dev/null +++ b/transformer_engine/common/activation/relu_dbias.cu @@ -0,0 +1,34 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_quantize_dbias_drelu(const NVTETensor input, const NVTETensor activation_input, + NVTETensor output, NVTETensor dbias, NVTETensor workspace, + cudaStream_t stream) { + NVTE_API_CALL(nvte_quantize_dbias_drelu); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = true; + + dispatch::quantize_bwd_helper>( + input, activation_input, output, dbias, workspace, nullptr, stream); +} + +void nvte_quantize_dbias_dsrelu(const NVTETensor input, const NVTETensor activation_input, + NVTETensor output, NVTETensor dbias, NVTETensor workspace, + cudaStream_t stream) { + NVTE_API_CALL(nvte_quantize_dbias_dsrelu); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = true; + + dispatch::quantize_bwd_helper>( + input, activation_input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/activation/relu_grouped.cu b/transformer_engine/common/activation/relu_grouped.cu new file mode 100644 index 000000000..93ce6b82f --- /dev/null +++ b/transformer_engine/common/activation/relu_grouped.cu @@ -0,0 +1,53 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_relu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_relu); + using namespace transformer_engine; + constexpr bool IS_ACT = true; + dispatch::group_quantize_fwd_helper>(input, output, nullptr, + stream); +} + +void nvte_group_drelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_drelu); + using namespace transformer_engine; + NVTEGroupedTensor dbias = nullptr; + NVTETensor workspace = nullptr; + + constexpr bool IS_DBIAS = false; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + grad, input, output, dbias, workspace, nullptr, stream); +} + +void nvte_group_srelu(const NVTEGroupedTensor input, NVTEGroupedTensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_group_srelu); + using namespace transformer_engine; + constexpr bool IS_ACT = true; + dispatch::group_quantize_fwd_helper>(input, output, nullptr, + stream); +} + +void nvte_group_dsrelu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_dsrelu); + using namespace transformer_engine; + NVTEGroupedTensor dbias = nullptr; + NVTETensor workspace = nullptr; + + constexpr bool IS_DBIAS = false; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + grad, input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/activation/relu_grouped_dbias.cu b/transformer_engine/common/activation/relu_grouped_dbias.cu new file mode 100644 index 000000000..2b9dcd35d --- /dev/null +++ b/transformer_engine/common/activation/relu_grouped_dbias.cu @@ -0,0 +1,36 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_quantize_dbias_drelu(const NVTEGroupedTensor input, + const NVTEGroupedTensor activation_input, + NVTEGroupedTensor output, NVTEGroupedTensor dbias, + NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_quantize_dbias_drelu); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + input, activation_input, output, dbias, workspace, nullptr, stream); +} + +void nvte_group_quantize_dbias_dsrelu(const NVTEGroupedTensor input, + const NVTEGroupedTensor activation_input, + NVTEGroupedTensor output, NVTEGroupedTensor dbias, + NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_quantize_dbias_dsrelu); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + input, activation_input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/activation/swiglu.cu b/transformer_engine/common/activation/swiglu.cu index 12478af4c..0b5b6069b 100644 --- a/transformer_engine/common/activation/swiglu.cu +++ b/transformer_engine/common/activation/swiglu.cu @@ -13,14 +13,6 @@ void nvte_silu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { act_fn>(input, output, stream); } -void nvte_group_silu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_silu); - using namespace transformer_engine; - constexpr bool IS_ACT = true; - dispatch::group_quantize_fwd_helper>(input, output, nullptr, - stream); -} - void nvte_dsilu(const NVTETensor grad, const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_dsilu); @@ -28,47 +20,6 @@ void nvte_dsilu(const NVTETensor grad, const NVTETensor input, NVTETensor output dact_fn>(grad, input, output, stream); } -void nvte_group_dsilu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, - NVTEGroupedTensor output, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_dsilu); - using namespace transformer_engine; - NVTEGroupedTensor dbias = nullptr; - NVTETensor workspace = nullptr; - - constexpr bool IS_DBIAS = false; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - grad, input, output, dbias, workspace, nullptr, stream); -} - -void nvte_quantize_dbias_dsilu(const NVTETensor input, const NVTETensor activation_input, - NVTETensor output, NVTETensor dbias, NVTETensor workspace, - cudaStream_t stream) { - NVTE_API_CALL(nvte_quantize_dbias_dsilu); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = true; - - dispatch::quantize_bwd_helper>( - input, activation_input, output, dbias, workspace, nullptr, stream); -} - -void nvte_group_quantize_dbias_dsilu(const NVTEGroupedTensor input, - const NVTEGroupedTensor activation_input, - NVTEGroupedTensor output, NVTEGroupedTensor dbias, - NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_quantize_dbias_dsilu); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = true; - - dispatch::group_quantize_bwd_helper>( - input, activation_input, output, dbias, workspace, nullptr, stream); -} - void nvte_swiglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_swiglu); using namespace transformer_engine; diff --git a/transformer_engine/common/activation/swiglu_dbias.cu b/transformer_engine/common/activation/swiglu_dbias.cu new file mode 100644 index 000000000..0e532acc5 --- /dev/null +++ b/transformer_engine/common/activation/swiglu_dbias.cu @@ -0,0 +1,21 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_quantize_dbias_dsilu(const NVTETensor input, const NVTETensor activation_input, + NVTETensor output, NVTETensor dbias, NVTETensor workspace, + cudaStream_t stream) { + NVTE_API_CALL(nvte_quantize_dbias_dsilu); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = true; + + dispatch::quantize_bwd_helper>( + input, activation_input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/activation/swiglu_grouped.cu b/transformer_engine/common/activation/swiglu_grouped.cu new file mode 100644 index 000000000..160ab6628 --- /dev/null +++ b/transformer_engine/common/activation/swiglu_grouped.cu @@ -0,0 +1,30 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_silu(const NVTEGroupedTensor input, NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_silu); + using namespace transformer_engine; + constexpr bool IS_ACT = true; + dispatch::group_quantize_fwd_helper>(input, output, nullptr, + stream); +} + +void nvte_group_dsilu(const NVTEGroupedTensor grad, const NVTEGroupedTensor input, + NVTEGroupedTensor output, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_dsilu); + using namespace transformer_engine; + NVTEGroupedTensor dbias = nullptr; + NVTETensor workspace = nullptr; + + constexpr bool IS_DBIAS = false; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + grad, input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/activation/swiglu_grouped_dbias.cu b/transformer_engine/common/activation/swiglu_grouped_dbias.cu new file mode 100644 index 000000000..83d15e802 --- /dev/null +++ b/transformer_engine/common/activation/swiglu_grouped_dbias.cu @@ -0,0 +1,22 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include "../util/math.h" +#include "./activation_template.h" + +void nvte_group_quantize_dbias_dsilu(const NVTEGroupedTensor input, + const NVTEGroupedTensor activation_input, + NVTEGroupedTensor output, NVTEGroupedTensor dbias, + NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_quantize_dbias_dsilu); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = true; + + dispatch::group_quantize_bwd_helper>( + input, activation_input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/cast/cast.cu b/transformer_engine/common/cast/cast.cu index 61cfacd33..1e3c04573 100644 --- a/transformer_engine/common/cast/cast.cu +++ b/transformer_engine/common/cast/cast.cu @@ -26,15 +26,6 @@ void nvte_quantize(const NVTETensor input, NVTETensor output, cudaStream_t strea dispatch::quantize_fwd_helper(input, output, nullptr, stream); } -void nvte_group_quantize(const NVTEGroupedTensor input, NVTEGroupedTensor output, - const NVTEQuantizationConfig quant_config, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_quantize); - using namespace transformer_engine; - - constexpr bool IS_ACT = false; - dispatch::group_quantize_fwd_helper(input, output, quant_config, stream); -} - void nvte_quantize_noop(const NVTETensor input, NVTETensor output, NVTETensor noop, cudaStream_t stream) { NVTE_API_CALL(nvte_quantize_noop); @@ -56,32 +47,6 @@ void nvte_quantize_v2(const NVTETensor input, NVTETensor output, dispatch::quantize_fwd_helper(input, output, quant_config, stream); } -void nvte_quantize_dbias(const NVTETensor input, NVTETensor output, NVTETensor dbias, - NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_quantize_dbias); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = false; - constexpr const NVTETensor activation_input = nullptr; - - dispatch::quantize_bwd_helper( - input, activation_input, output, dbias, workspace, nullptr, stream); -} - -void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor output, - NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream) { - NVTE_API_CALL(nvte_group_quantize_dbias); - using namespace transformer_engine; - - constexpr bool IS_DBIAS = true; - constexpr bool IS_DACT = false; - constexpr const NVTEGroupedTensor activation_input = nullptr; - - dispatch::group_quantize_bwd_helper( - input, activation_input, output, dbias, workspace, nullptr, stream); -} - void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t stream) { NVTE_API_CALL(nvte_dequantize); using namespace transformer_engine; @@ -89,14 +54,6 @@ void nvte_dequantize(const NVTETensor input, NVTETensor output, cudaStream_t str stream); } -void nvte_group_dequantize(const NVTEGroupedTensor input, NVTEGroupedTensor output, - cudaStream_t stream) { - NVTE_API_CALL(nvte_group_dequantize); - using namespace transformer_engine; - dispatch::group_dequantize_helper(*convertNVTEGroupedTensorCheck(input), - convertNVTEGroupedTensorCheck(output), stream); -} - void nvte_multi_tensor_quantize(const NVTETensor *inputs, NVTETensor *outputs, const NVTEQuantizationConfig quant_configs, const size_t num_tensors, cudaStream_t stream) { @@ -130,19 +87,3 @@ void nvte_multi_tensor_quantize(const NVTETensor *inputs, NVTETensor *outputs, NVTE_CHECK_CUDA(cudaStreamWaitEvent(stream, detail::get_compute_stream_event(s))); } } - -// Group quantize assumes contiguous inputs and outputs in memory allocation -// Note: this API assumes knowing split sections from the host, if split information -// comes from D2H copy, it will break cuda graph capture -void nvte_group_nvfp4_quantize_with_amax(const NVTETensor input, NVTETensor *outputs, - const size_t *split_sections, const size_t num_tensors, - const NVTEQuantizationConfig quant_config, - cudaStream_t stream) { - NVTE_API_CALL(nvte_group_nvfp4_quantize_with_amax); - using namespace transformer_engine; - - constexpr bool IS_ACT = false; - - dispatch::group_quantize_fwd_host_aware_helper( - input, outputs, split_sections, num_tensors, quant_config, stream); -} diff --git a/transformer_engine/common/cast/cast_dbias.cu b/transformer_engine/common/cast/cast_dbias.cu new file mode 100644 index 000000000..480e8ca74 --- /dev/null +++ b/transformer_engine/common/cast/cast_dbias.cu @@ -0,0 +1,24 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include + +#include "../common.h" +#include "dispatch/quantize.cuh" + +void nvte_quantize_dbias(const NVTETensor input, NVTETensor output, NVTETensor dbias, + NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_quantize_dbias); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = false; + constexpr const NVTETensor activation_input = nullptr; + + dispatch::quantize_bwd_helper( + input, activation_input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/cast/cast_grouped.cu b/transformer_engine/common/cast/cast_grouped.cu new file mode 100644 index 000000000..853634c81 --- /dev/null +++ b/transformer_engine/common/cast/cast_grouped.cu @@ -0,0 +1,47 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include +#include +#include + +#include "../common.h" +#include "dispatch/dequantize.cuh" +#include "dispatch/quantize.cuh" + +void nvte_group_quantize(const NVTEGroupedTensor input, NVTEGroupedTensor output, + const NVTEQuantizationConfig quant_config, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_quantize); + using namespace transformer_engine; + + constexpr bool IS_ACT = false; + dispatch::group_quantize_fwd_helper(input, output, quant_config, stream); +} + +void nvte_group_dequantize(const NVTEGroupedTensor input, NVTEGroupedTensor output, + cudaStream_t stream) { + NVTE_API_CALL(nvte_group_dequantize); + using namespace transformer_engine; + dispatch::group_dequantize_helper(*convertNVTEGroupedTensorCheck(input), + convertNVTEGroupedTensorCheck(output), stream); +} + +// Group quantize assumes contiguous inputs and outputs in memory allocation. +// Note: this API assumes knowing split sections from the host. If split information +// comes from D2H copy, it will break cuda graph capture. +void nvte_group_nvfp4_quantize_with_amax(const NVTETensor input, NVTETensor *outputs, + const size_t *split_sections, const size_t num_tensors, + const NVTEQuantizationConfig quant_config, + cudaStream_t stream) { + NVTE_API_CALL(nvte_group_nvfp4_quantize_with_amax); + using namespace transformer_engine; + + constexpr bool IS_ACT = false; + + dispatch::group_quantize_fwd_host_aware_helper( + input, outputs, split_sections, num_tensors, quant_config, stream); +} diff --git a/transformer_engine/common/cast/cast_grouped_dbias.cu b/transformer_engine/common/cast/cast_grouped_dbias.cu new file mode 100644 index 000000000..5290255a0 --- /dev/null +++ b/transformer_engine/common/cast/cast_grouped_dbias.cu @@ -0,0 +1,24 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +#include +#include + +#include "../common.h" +#include "dispatch/quantize.cuh" + +void nvte_group_quantize_dbias(const NVTEGroupedTensor input, NVTEGroupedTensor output, + NVTEGroupedTensor dbias, NVTETensor workspace, cudaStream_t stream) { + NVTE_API_CALL(nvte_group_quantize_dbias); + using namespace transformer_engine; + + constexpr bool IS_DBIAS = true; + constexpr bool IS_DACT = false; + constexpr const NVTEGroupedTensor activation_input = nullptr; + + dispatch::group_quantize_bwd_helper( + input, activation_input, output, dbias, workspace, nullptr, stream); +} diff --git a/transformer_engine/common/cast/core/common.cuh b/transformer_engine/common/cast/core/common.cuh index 90e57a6fe..3e6eb55b7 100644 --- a/transformer_engine/common/cast/core/common.cuh +++ b/transformer_engine/common/cast/core/common.cuh @@ -17,6 +17,7 @@ #include #include "../../common.h" +#include "../../util/ptx.cuh" #include "../../utils.cuh" namespace transformer_engine { From 8c0f1d242ed1b6eb84de1b1ce576662b8d606dc1 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Thu, 21 May 2026 10:06:32 -0700 Subject: [PATCH 071/180] [JAX] Improve JAX tutorial documentation (#2976) Signed-off-by: Jeremy Berchtold Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Teddy Do --- docs/examples/jax/attention.rst | 11 + docs/examples/jax/collective_gemm.rst | 11 + docs/examples/jax/dense.out | 21 + docs/examples/jax/dense.py | 180 +++++++ docs/examples/jax/dense.rst | 166 +++++++ docs/examples/jax/expert_parallelism.rst | 11 + .../{ => jax}/quickstart_jax_utils.py | 52 ++ docs/examples/jax/test_dense.py | 87 ++++ docs/examples/te_jax_integration.ipynb | 462 ------------------ docs/examples/te_jax_integration.rst | 95 ++++ docs/index.rst | 2 +- qa/L0_jax_distributed_unittest/test.sh | 5 + qa/L0_jax_unittest/test.sh | 5 + 13 files changed, 645 insertions(+), 463 deletions(-) create mode 100644 docs/examples/jax/attention.rst create mode 100644 docs/examples/jax/collective_gemm.rst create mode 100644 docs/examples/jax/dense.out create mode 100644 docs/examples/jax/dense.py create mode 100644 docs/examples/jax/dense.rst create mode 100644 docs/examples/jax/expert_parallelism.rst rename docs/examples/{ => jax}/quickstart_jax_utils.py (64%) create mode 100644 docs/examples/jax/test_dense.py delete mode 100644 docs/examples/te_jax_integration.ipynb create mode 100644 docs/examples/te_jax_integration.rst diff --git a/docs/examples/jax/attention.rst b/docs/examples/jax/attention.rst new file mode 100644 index 000000000..c9f84da63 --- /dev/null +++ b/docs/examples/jax/attention.rst @@ -0,0 +1,11 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +JAX: Attention with TransformerEngine +===================================== + +**TODO — Coming soon.** + +`← Back to the JAX integration overview <../te_jax_integration.html>`_ diff --git a/docs/examples/jax/collective_gemm.rst b/docs/examples/jax/collective_gemm.rst new file mode 100644 index 000000000..05b39ea01 --- /dev/null +++ b/docs/examples/jax/collective_gemm.rst @@ -0,0 +1,11 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +JAX: Collective GEMMs with TransformerEngine +============================================= + +**TODO — Coming soon.** + +`← Back to the JAX integration overview <../te_jax_integration.html>`_ diff --git a/docs/examples/jax/dense.out b/docs/examples/jax/dense.out new file mode 100644 index 000000000..22b93ff04 --- /dev/null +++ b/docs/examples/jax/dense.out @@ -0,0 +1,21 @@ +# Numbers below are illustrative (captured on a GB200). Regenerate with: +# python3 docs/examples/jax/dense.py > dense.out + +# SINGLE_GPU_OUTPUT_START +Variable collections: ['params'] +{'params': {'Dense_0': {'kernel': ((8192, 32768), dtype('float32'))}}} + +bf16 baseline: +Mean time: 18.056 ms + +TE MXFP8BlockScaling: +Mean time: 11.260 ms +# SINGLE_GPU_OUTPUT_END + +# MULTI_GPU_OUTPUT_START +bf16 DP=2/TP=2: +Mean time: 5.516 ms + +TE MXFP8BlockScaling DP=2/TP=2: +Mean time: 3.712 ms +# MULTI_GPU_OUTPUT_END diff --git a/docs/examples/jax/dense.py b/docs/examples/jax/dense.py new file mode 100644 index 000000000..9ddc5a9e8 --- /dev/null +++ b/docs/examples/jax/dense.py @@ -0,0 +1,180 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""JAX: Dense GEMMs with TransformerEngine. + +Companion source for ``dense.rst``. Code blocks between ``# DENSE_*_START`` / +``# DENSE_*_END`` markers are pulled into the RST via ``literalinclude``. + +Run as a script to exercise the example end-to-end: + + python docs/examples/jax/dense.py + +Pytest tests live in ``test_dense.py``; the multi-GPU section auto-skips when +fewer than 4 GPUs are visible. +""" + +# DENSE_IMPORTS_START +import jax +import jax.numpy as jnp +from flax import linen as nn + +import quickstart_jax_utils as utils + +# DENSE_IMPORTS_END + + +# DENSE_BASELINE_MODEL_START +class FlaxDenseBlock(nn.Module): + """One linear layer. ``dot_general_cls`` lets us swap the GEMM impl.""" + + features: int + dtype: jnp.dtype = jnp.bfloat16 + dot_general_cls: callable = lambda: None + + @nn.compact + def __call__(self, x): + return nn.Dense( + features=self.features, + use_bias=False, + dtype=self.dtype, + dot_general=self.dot_general_cls(), + )(x) + + +# DENSE_BASELINE_MODEL_END + + +# DENSE_INPUTS_SETUP_START +batch, seq, hidden, out_features = 8, 2048, 8192, 32768 +dtype = jnp.bfloat16 + +key = jax.random.PRNGKey(0) +k_init, k_x, k_dy = jax.random.split(key, 3) +x = jax.random.normal(k_x, (batch, seq, hidden)).astype(dtype) +dy = jax.random.normal(k_dy, (batch, seq, out_features)).astype(dtype) + +baseline = FlaxDenseBlock(features=out_features) +baseline_vars = baseline.init(k_init, x) +# DENSE_INPUTS_SETUP_END + + +# DENSE_TE_SETUP_START +from transformer_engine.jax import flax as te_flax +from transformer_engine.common.recipe import MXFP8BlockScaling + +recipe = MXFP8BlockScaling() +te_dot_general_cls = te_flax.make_dot_general_cls(recipe) + +te_model = FlaxDenseBlock(features=out_features, dot_general_cls=te_dot_general_cls) +te_vars = te_model.init(k_init, x) + +print("Variable collections:", list(te_vars.keys())) +print(jax.tree_util.tree_map(lambda a: (a.shape, a.dtype), te_vars)) +# DENSE_TE_SETUP_END + + +# DENSE_SINGLE_GPU_BENCH_START +def run_single_gpu_bench(): + print("bf16 baseline:") + utils.speedometer( + model_apply_fn=baseline.apply, + variables=baseline_vars, + input=x, + output_grad=dy, + ) + + print(f"\nTE {type(recipe).__name__}:") + utils.speedometer( + model_apply_fn=te_model.apply, + variables=te_vars, + input=x, + output_grad=dy, + ) + + +# DENSE_SINGLE_GPU_BENCH_END + + +# DENSE_MULTI_GPU_MESH_SETUP_START +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P +from jax.experimental import mesh_utils +from transformer_engine.jax.sharding import MeshResource, global_shard_guard + + +def build_dp_tp_mesh(): + # 2x2 mesh: DP on one axis, TP on the other. + devices = mesh_utils.create_device_mesh((2, 2)) + mesh = Mesh(devices, axis_names=("dp", "tp")) + + # Tell TE which mesh axis is which. This is a *global* setting, established + # outside JIT, so TE's GEMM primitives can plan comms accordingly. + mesh_resource = MeshResource(dp_resource="dp", tp_resource="tp") + return mesh, mesh_resource + + +# DENSE_MULTI_GPU_MESH_SETUP_END + + +# DENSE_MULTI_GPU_SHARD_SETUP_START +def shard_variables(mesh, variables_dict): + kernel_sharding = NamedSharding(mesh, P(None, "tp")) + + def _shard(variables): + params = variables["params"] + sharded = jax.device_put(params["Dense_0"]["kernel"], kernel_sharding) + return { + **variables, + "params": { + **params, + "Dense_0": {**params["Dense_0"], "kernel": sharded}, + }, + } + + input_sharding = NamedSharding(mesh, P("dp", None, None)) + output_grad_sharding = NamedSharding(mesh, P("dp", None, "tp")) + + return { + "x": jax.device_put(x, input_sharding), + "dy": jax.device_put(dy, output_grad_sharding), + **{name: _shard(vars_) for name, vars_ in variables_dict.items()}, + } + + +# DENSE_MULTI_GPU_SHARD_SETUP_END + + +# DENSE_MULTI_GPU_BENCH_START +def run_multi_gpu_bench(): + mesh, mesh_resource = build_dp_tp_mesh() + sharded = shard_variables(mesh, {"baseline": baseline_vars, "te": te_vars}) + + with jax.set_mesh(mesh), global_shard_guard(mesh_resource): + print("bf16 DP=2/TP=2:") + utils.speedometer( + model_apply_fn=baseline.apply, + variables=sharded["baseline"], + input=sharded["x"], + output_grad=sharded["dy"], + ) + + print(f"\nTE {type(recipe).__name__} DP=2/TP=2:") + utils.speedometer( + model_apply_fn=te_model.apply, + variables=sharded["te"], + input=sharded["x"], + output_grad=sharded["dy"], + ) + + +# DENSE_MULTI_GPU_BENCH_END + + +if __name__ == "__main__": + run_single_gpu_bench() + if len(jax.devices()) >= 4: + print() + run_multi_gpu_bench() + else: + print("\n[skipped multi-GPU section: <4 devices visible]") diff --git a/docs/examples/jax/dense.rst b/docs/examples/jax/dense.rst new file mode 100644 index 000000000..2087d49c7 --- /dev/null +++ b/docs/examples/jax/dense.rst @@ -0,0 +1,166 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +JAX: Dense GEMMs with TransformerEngine +======================================= + +This document walks through replacing a plain ``flax.linen.Dense``'s GEMM with +TransformerEngine's quantized GEMM. + +**Recipe.** We use ``MXFP8BlockScaling`` in this tutorial. ``MXFP8BlockScaling`` and +``NVFP4BlockScaling`` require a Blackwell-class GPU; on Hopper, swap in +``DelayedScaling`` or ``Float8CurrentScaling``. For more information on recipes, see this :ref:`recipe overview `. + +`← Back to the JAX integration overview <../te_jax_integration.html>`_ + +1. Baseline: a plain Flax Dense block +------------------------------------- + +We isolate the optimization to a single linear layer so it's clear what's +changing. ``dot_general_cls`` is exposed as a constructor argument so we can swap +in TE later without touching the model definition. + +.. literalinclude:: dense.py + :language: python + :start-after: # DENSE_BASELINE_MODEL_START + :end-before: # DENSE_BASELINE_MODEL_END + +.. literalinclude:: dense.py + :language: python + :start-after: # DENSE_INPUTS_SETUP_START + :end-before: # DENSE_INPUTS_SETUP_END + + +2. Quantized Dense via ``make_dot_general_cls`` +----------------------------------------------- + +TE exposes a helper, ``te_flax.make_dot_general_cls(recipe)``, that returns a Flax +module class you pass directly to ``nn.Dense(..., dot_general=...)``. + +With this API, TE doesn't create the ``kernel`` params; it only wraps the GEMM. +All your initialization, sharding annotations, and optimizer state stay where +they were. + +.. literalinclude:: dense.py + :language: python + :start-after: # DENSE_TE_SETUP_START + :end-before: # DENSE_TE_SETUP_END + +If using ``DelayedScaling``, see [#delayedscaling]_. + + +3. Single-GPU performance +------------------------- + +``speedometer`` runs a JIT-compiled forward+backward loop with warmup, on the +same input for both models. + +.. literalinclude:: dense.py + :language: python + :start-after: # DENSE_SINGLE_GPU_BENCH_START + :end-before: # DENSE_SINGLE_GPU_BENCH_END + +.. raw:: html + +
+ Output: +
+ +.. container:: program-output + + .. literalinclude:: dense.out + :language: text + :start-after: # SINGLE_GPU_OUTPUT_START + :end-before: # SINGLE_GPU_OUTPUT_END + +On a single GB200, that's roughly **1.6× faster** for the fwd+bwd of one large +Dense — and the only code change was passing ``dot_general=te_dot_general_cls()`` +into ``nn.Dense``. + +The speedup depends on shape: large GEMMs benefit most. Very small GEMMs may +not benefit at all because the cast + scale overhead can dominate. + +.. warning:: + + **Remat / activation checkpointing.** If your training loop uses + ``jax.checkpoint_policies.checkpoint_dots`` (or any policy that matches + ``jax.lax.dot_general``), swap it for + ``transformer_engine.jax.checkpoint_policies.checkpoint_dots_and_te_gemms``. + Otherwise TE's quantized GEMM primitives won't be checkpointed correctly + and your performance comparison will not be accurate. + + +4. Multi-GPU: DP=2 / TP=2 on a single Dense +------------------------------------------- + +**Prerequisite:** this section requires four GPUs. + +Keeping the same ``FlaxDenseBlock`` from the rest of the document, we run it on +a 2×2 mesh with **data parallelism** on one axis and **tensor parallelism** +(column-parallel: shard the kernel's output dim) on the other. + +Two pieces wire this up: + +1. A ``jax.sharding.Mesh`` you build once at module scope (outside JIT). +2. TE's ``MeshResource``, set globally via ``global_shard_guard``, which tells + TE which mesh axes are DP and TP. + +.. literalinclude:: dense.py + :language: python + :start-after: # DENSE_MULTI_GPU_MESH_SETUP_START + :end-before: # DENSE_MULTI_GPU_MESH_SETUP_END + +**Sharding plan:** + +.. csv-table:: + :header: "Tensor", "Shape", "PartitionSpec" + :widths: 30, 40, 30 + + "Kernel (column-parallel)", "``(hidden, out_features)``", "``P(None, 'tp')``" + "Input activations", "``(batch, seq, hidden)``", "``P('dp', None, None)``" + "Gradient on output", "``(batch, seq, out_features)``", "``P('dp', None, 'tp')``" + +.. literalinclude:: dense.py + :language: python + :start-after: # DENSE_MULTI_GPU_SHARD_SETUP_START + :end-before: # DENSE_MULTI_GPU_SHARD_SETUP_END + +.. literalinclude:: dense.py + :language: python + :start-after: # DENSE_MULTI_GPU_BENCH_START + :end-before: # DENSE_MULTI_GPU_BENCH_END + +.. raw:: html + +
+ Output: +
+ +.. container:: program-output + + .. literalinclude:: dense.out + :language: text + :start-after: # MULTI_GPU_OUTPUT_START + :end-before: # MULTI_GPU_OUTPUT_END + + +Next steps +---------- + +* `Collective GEMM `_: further speedups by communicating between devices inside the GEMM. +* `← Hub <../te_jax_integration.html>`_ + +.. rubric:: Footnotes + +.. [#delayedscaling] **DelayedScaling state.** Most recipes are stateless — scaling factors are computed from each + tensor as it flows through the GEMM, so there is nothing to persist across steps. However, if you swap in + ``DelayedScaling`` instead, ``init`` will produce a second variable collection, + ``_overwrite_with_gradient``, holding ``kernel_amax_history``, ``kernel_scale``, + ``x_amax_history``, ``x_scale``, etc. These are **not** model parameters — they are Flax + variables that TE updates each step to compute per-tensor scales from a rolling amax window. + If you use ``DelayedScaling``, you must thread the *entire* ``var_collect`` through your + training loop (not just ``params``) so the history persists across steps, otherwise training + accuracy will be impacted. ``MXFP8BlockScaling``, ``NVFP4BlockScaling``, and + ``Float8CurrentScaling`` do not require this. diff --git a/docs/examples/jax/expert_parallelism.rst b/docs/examples/jax/expert_parallelism.rst new file mode 100644 index 000000000..5e94e1d29 --- /dev/null +++ b/docs/examples/jax/expert_parallelism.rst @@ -0,0 +1,11 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +JAX: Expert Parallelism with TransformerEngine +============================================== + +**TODO — Coming soon.** + +`← Back to the JAX integration overview <../te_jax_integration.html>`_ diff --git a/docs/examples/quickstart_jax_utils.py b/docs/examples/jax/quickstart_jax_utils.py similarity index 64% rename from docs/examples/quickstart_jax_utils.py rename to docs/examples/jax/quickstart_jax_utils.py index 0c5ec5295..6547a5ff1 100644 --- a/docs/examples/quickstart_jax_utils.py +++ b/docs/examples/jax/quickstart_jax_utils.py @@ -4,6 +4,7 @@ import jax import jax.numpy as jnp +import numpy as np import time from typing import Callable, Any, Dict, Optional, Tuple @@ -99,3 +100,54 @@ def _split_step_rngs( new_rngs[name] = new_key step_rngs[name] = step_key return new_rngs, step_rngs + + +def compare_fwd_bwd( + ref_apply_fn: Callable, + ref_variables: Any, + test_apply_fn: Callable, + test_variables: Any, + *, + input: jnp.ndarray, + output_grad: jnp.ndarray, + rtol: float = 1e-5, + atol: float = 1e-8, + rtol_dW: Optional[float] = None, + atol_dW: Optional[float] = None, +) -> None: + """Compare forward outputs and VJP gradients between two models. + + Runs ``y, vjp_fn = jax.vjp(apply_fn, variables, input)`` for each model, + then applies ``vjp_fn(output_grad)`` to get gradients wrt both the + parameters (``dW``) and the input (``dx``). Calls + ``numpy.testing.assert_allclose`` on each tensor (``y``, ``dx``, and every + leaf of ``dW``). ``rtol_dW`` / ``atol_dW`` override ``rtol`` / ``atol`` + for the params-grad comparison. + """ + rtol_dW = rtol if rtol_dW is None else rtol_dW + atol_dW = atol if atol_dW is None else atol_dW + + def _run(apply_fn: Callable) -> Callable: + @jax.jit + def go(variables, inp, dy): + y, vjp_fn = jax.vjp(apply_fn, variables, inp) + dvars, dx = vjp_fn(dy.astype(y.dtype)) + return y, dvars["params"], dx + + return go + + y_ref, dW_ref, dx_ref = _run(ref_apply_fn)(ref_variables, input, output_grad) + y_test, dW_test, dx_test = _run(test_apply_fn)(test_variables, input, output_grad) + + np.testing.assert_allclose( + y_test, y_ref, rtol=rtol, atol=atol, err_msg="forward output (y) mismatch" + ) + np.testing.assert_allclose( + dx_test, dx_ref, rtol=rtol, atol=atol, err_msg="input grad (dx) mismatch" + ) + for ref_leaf, test_leaf in zip( + jax.tree_util.tree_leaves(dW_ref), jax.tree_util.tree_leaves(dW_test) + ): + np.testing.assert_allclose( + test_leaf, ref_leaf, rtol=rtol_dW, atol=atol_dW, err_msg="params grad (dW) mismatch" + ) diff --git a/docs/examples/jax/test_dense.py b/docs/examples/jax/test_dense.py new file mode 100644 index 000000000..4bedd9d40 --- /dev/null +++ b/docs/examples/jax/test_dense.py @@ -0,0 +1,87 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Pytest entry points for ``dense.py``. + +These run the same code shown in ``dense.py`` and add numeric / smoke +assertions so CI catches regressions. + +Run with: + + pytest -v docs/examples/jax/test_dense.py + +The multi-GPU section auto-skips when fewer than 4 GPUs are visible. +""" + +import jax +import jax.numpy as jnp +import pytest + +import quickstart_jax_utils as utils +from transformer_engine.jax.quantize import is_scaling_mode_supported, ScalingMode + +# Imports from ``dense`` are intentionally deferred into each test body. dense.py +# runs ``te_vars = te_model.init(k_init, x)`` at module scope, which raises on +# devices without MXFP8 support (Hopper or older). A top-level import would fire +# that before pytest can apply the @requires_mxfp8 skip marks. + +_mxfp8_supported, _mxfp8_reason = is_scaling_mode_supported(ScalingMode.MXFP8_1D_SCALING) +requires_mxfp8 = pytest.mark.skipif( + not _mxfp8_supported, reason=f"MXFP8 not supported on this device: {_mxfp8_reason}" +) + + +@requires_mxfp8 +def test_baseline_runs(): + from dense import baseline, baseline_vars, batch, dtype, out_features, seq, x + + out = baseline.apply(baseline_vars, x) + assert out.shape == (batch, seq, out_features) + assert out.dtype == dtype + + +@requires_mxfp8 +def test_te_dense_runs(): + from dense import batch, out_features, seq, te_model, te_vars, x + + out = te_model.apply(te_vars, x) + assert out.shape == (batch, seq, out_features) + + +@requires_mxfp8 +def test_te_matches_baseline(): + """TE quantized Dense should match the bf16 baseline within MXFP8 tolerance.""" + from dense import baseline, baseline_vars, batch, dy, seq, te_model, te_vars, x + + fp8_rel_noise = float(jnp.finfo(jnp.float8_e4m3fn).eps) + atol_fwd = 10.0 * fp8_rel_noise + atol_dw = atol_fwd * jnp.sqrt(batch * seq).item() + + utils.compare_fwd_bwd( + baseline.apply, + baseline_vars, + te_model.apply, + te_vars, + input=x, + output_grad=dy, + rtol=fp8_rel_noise, + atol=atol_fwd, + rtol_dW=fp8_rel_noise, + atol_dW=atol_dw, + ) + + +@requires_mxfp8 +def test_single_gpu_benchmark(): + from dense import run_single_gpu_bench + + run_single_gpu_bench() + + +@requires_mxfp8 +@pytest.mark.skipif(len(jax.devices()) < 4, reason="needs 4 GPUs for DP=2/TP=2") +def test_multi_gpu_benchmark(): + from dense import run_multi_gpu_bench + + run_multi_gpu_bench() diff --git a/docs/examples/te_jax_integration.ipynb b/docs/examples/te_jax_integration.ipynb deleted file mode 100644 index 66d16ed52..000000000 --- a/docs/examples/te_jax_integration.ipynb +++ /dev/null @@ -1,462 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "id": "962d87bb", - "metadata": {}, - "source": [ - "\n", - "\n", - "# JAX: Integrating TE into an existing framework\n", - "\n", - "This tutorial will cover how to integrate TransformerEngine into an existing JAX model framework, such as [MaxText's TE integration](https://github.com/AI-Hypercomputer/maxtext/blob/ed517cf80d9aa81f76e236c5516dacebfe39e96d/src/MaxText/layers/quantizations.py#L753) or your own model framework. \n" - ] - }, - { - "cell_type": "markdown", - "id": "b36876bb", - "metadata": {}, - "source": [ - "Let's start with a standard JAX+Flax Transformer layer" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "id": "d5284a38", - "metadata": {}, - "outputs": [], - "source": [ - "import jax\n", - "import jax.numpy as jnp\n", - "from flax import linen as nn\n", - "import quickstart_jax_utils as utils\n", - "from typing import Optional" - ] - }, - { - "cell_type": "code", - "execution_count": 11, - "id": "a4d1cfdc", - "metadata": {}, - "outputs": [], - "source": [ - "class FlaxMLP(nn.Module):\n", - " \"\"\"Feed-forward network in Transformer layer\n", - " Built with plain Flax modules.\n", - " \"\"\"\n", - " hidden_size: int\n", - " ffn_hidden_size: int\n", - " dot_general_cls: callable = lambda: None\n", - "\n", - " @nn.compact\n", - " def __call__(self, x: jnp.ndarray) -> jnp.ndarray:\n", - " x = nn.Dense(features=self.ffn_hidden_size, use_bias=True, dot_general=self.dot_general_cls())(x)\n", - " x = nn.gelu(x, approximate=True) # equivalent to tanh approximation\n", - " x = nn.Dense(features=self.hidden_size, use_bias=True, dot_general=self.dot_general_cls())(x)\n", - " return x\n", - "\n", - "class FlaxTransformerLayer(nn.Module):\n", - " \"\"\"Basic Transformer layer using plain Flax modules\"\"\"\n", - " hidden_size: int\n", - " ffn_hidden_size: int\n", - " num_attention_heads: int\n", - " layernorm_eps: float = 1e-5\n", - " attention_dropout: float = 0.1\n", - " dot_general_cls: callable = lambda: None\n", - " \n", - " def setup(self):\n", - " self.kv_channels = self.hidden_size // self.num_attention_heads\n", - "\n", - " @nn.compact\n", - " def __call__(\n", - " self, \n", - " x: jnp.ndarray, \n", - " attention_mask: Optional[jnp.ndarray] = None,\n", - " deterministic: bool = False\n", - " ) -> jnp.ndarray:\n", - " # Create causal mask if not provided\n", - " if attention_mask is None:\n", - " attention_mask = nn.make_causal_mask(x[..., 0], dtype=jnp.bool_)\n", - " \n", - " res = x\n", - " x = nn.LayerNorm(epsilon=self.layernorm_eps)(x)\n", - " \n", - " # Fused QKV projection\n", - " qkv = nn.Dense(features=3 * self.hidden_size, use_bias=True, dot_general=self.dot_general_cls())(x)\n", - " qkv = qkv.reshape(qkv.shape[0], qkv.shape[1], self.num_attention_heads, 3 * self.kv_channels)\n", - " q, k, v = jnp.split(qkv, 3, axis=3)\n", - " \n", - " # q, k, v now have shape [batch, seq_len, num_heads, kv_channels]\n", - " # which is the correct format for dot_product_attention\n", - " \n", - " # Apply dot product attention\n", - " # Note: dot_product_attention expects mask to be broadcastable to \n", - " # [batch, num_heads, q_length, kv_length], but attention_mask from \n", - " # nn.make_causal_mask has shape [batch, 1, seq_len, seq_len]\n", - " \n", - " # Generate dropout RNG key when needed (not deterministic and dropout_rate > 0)\n", - " dropout_rng = None\n", - " if not deterministic and self.attention_dropout > 0:\n", - " dropout_rng = self.make_rng('dropout')\n", - " \n", - " # See quickstart_jax.ipynb for details on using TE's faster fused attention\n", - " x = nn.dot_product_attention(\n", - " query=q,\n", - " key=k,\n", - " value=v,\n", - " mask=attention_mask,\n", - " dropout_rng=dropout_rng,\n", - " dropout_rate=self.attention_dropout,\n", - " deterministic=deterministic,\n", - " broadcast_dropout=True,\n", - " )\n", - " \n", - " # Reshape output from [batch, seq_len, num_heads, kv_channels] to [batch, seq_len, hidden_size]\n", - " x = x.reshape(x.shape[0], x.shape[1], self.hidden_size)\n", - "\n", - " # Output projection\n", - " x = nn.Dense(features=self.hidden_size, use_bias=True, dot_general=self.dot_general_cls())(x)\n", - " \n", - " x = res + x\n", - " \n", - " # Second residual connection\n", - " res = x\n", - " x = nn.LayerNorm(epsilon=self.layernorm_eps)(x)\n", - " \n", - " # MLP\n", - " mlp = FlaxMLP(\n", - " hidden_size=self.hidden_size,\n", - " ffn_hidden_size=self.ffn_hidden_size,\n", - " dot_general_cls=self.dot_general_cls,\n", - " )\n", - " x = mlp(x)\n", - " \n", - " return x + res\n" - ] - }, - { - "cell_type": "markdown", - "id": "db16bf70", - "metadata": {}, - "source": [ - "We've exposed `dot_general_cls` here so we can test out different GEMM implementations later. By default, Flax's `nn.Dense` will use JAX's GEMM `jax.lax.dot_general` when `dot_general` is `None`." - ] - }, - { - "cell_type": "markdown", - "id": "fbc3510b", - "metadata": {}, - "source": [ - "## Testing Performance\n", - "\n", - "Now let's test the performance of our FlaxTransformerLayer:\n" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "id": "8b44649d", - "metadata": {}, - "outputs": [], - "source": [ - "# Layer configuration\n", - "hidden_size = 4096\n", - "sequence_length = 2048\n", - "batch_size = 4\n", - "ffn_hidden_size = 16384\n", - "num_attention_heads = 32\n", - "dtype = jnp.bfloat16\n", - "\n", - "# Synthetic data\n", - "key, dropout_key = jax.random.split(jax.random.PRNGKey(42))\n", - "x = jax.random.normal(key, (batch_size, sequence_length, hidden_size)).astype(dtype)\n", - "dy = jax.random.normal(key, (batch_size, sequence_length, hidden_size)).astype(dtype)\n" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "id": "e44ed26d", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Pure Flax FlaxTransformerLayer initialized successfully!\n", - "Parameter shapes: {'params': {'Dense_0': {'bias': (12288,), 'kernel': (4096, 12288)}, 'Dense_1': {'bias': (4096,), 'kernel': (4096, 4096)}, 'FlaxMLP_0': {'Dense_0': {'bias': (16384,), 'kernel': (4096, 16384)}, 'Dense_1': {'bias': (4096,), 'kernel': (16384, 4096)}}, 'LayerNorm_0': {'bias': (4096,), 'scale': (4096,)}, 'LayerNorm_1': {'bias': (4096,), 'scale': (4096,)}}}\n" - ] - } - ], - "source": [ - "# Initialize the FlaxTransformerLayer\n", - "flax_transformer = FlaxTransformerLayer(\n", - " hidden_size=hidden_size,\n", - " ffn_hidden_size=ffn_hidden_size,\n", - " num_attention_heads=num_attention_heads,\n", - ")\n", - "\n", - "# Initialize parameters\n", - "params = flax_transformer.init(key, x, attention_mask=None, deterministic=False)\n", - "\n", - "print(\"Pure Flax FlaxTransformerLayer initialized successfully!\")\n", - "print(f\"Parameter shapes: {jax.tree_util.tree_map(lambda x: x.shape, params)}\")\n" - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "id": "de91af7a", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Input shape: (4, 2048, 4096)\n", - "Output shape: (4, 2048, 4096)\n", - "Output dtype: float32\n", - "Forward pass completed successfully!\n" - ] - } - ], - "source": [ - "# Example usage of forward pass\n", - "y = flax_transformer.apply(params, x, attention_mask=None, deterministic=True)\n", - "print(f\"Input shape: {x.shape}\")\n", - "print(f\"Output shape: {y.shape}\")\n", - "print(f\"Output dtype: {y.dtype}\")\n", - "print(\"Forward pass completed successfully!\")\n" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "id": "037bc8d9", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Mean time: 18.83516788482666 ms\n" - ] - } - ], - "source": [ - "import importlib\n", - "import quickstart_jax_utils\n", - "importlib.reload(quickstart_jax_utils)\n", - "\n", - "utils.speedometer(\n", - " model_apply_fn=flax_transformer.apply,\n", - " variables=params,\n", - " input=x,\n", - " output_grad=dy,\n", - " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", - " rngs={\"dropout\": dropout_key},\n", - ")" - ] - }, - { - "cell_type": "markdown", - "id": "5e9310c9", - "metadata": {}, - "source": [ - "## Transformer Engine" - ] - }, - { - "cell_type": "markdown", - "id": "1f8e213e", - "metadata": {}, - "source": [ - "TransformerEngine/JAX is currently using Flax Linen. However, it is easily compatible with Flax NNX or Haiku.\n", - "* [Use Flax NNX and Linen together](https://flax.readthedocs.io/en/latest/guides/bridge_guide.html)\n", - "* [Haiku and Flax interop](https://dm-haiku.readthedocs.io/en/latest/notebooks/flax.html)\n", - "\n", - "Additionally, with the tutorial below, no model parameters need to be managed by TransformerEngine. You can keep all your existing model parameters, initialization, and sharding the same. The only change required is to call TE's dot_general_cls instead of the default Dense dot_general implementation. TE's dot_general_cls is a small module that performs a quantized dense VJP and stores some small recipe-specific state." - ] - }, - { - "cell_type": "markdown", - "id": "4477d4e9", - "metadata": {}, - "source": [ - "Now we'll select a recipe. `DelayedScaling` and `CurrentScaling` use per-tensor scaling and are supported on Hopper and Blackwell. `MXFP8BlockScaling` and `NVFP4BlockScaling` use block scaling or a combination of both per-tensor and block scaling and are supported on Blackwell.\n", - "\n", - "If you would like to customize the recipe further, various options can be changed by passing args to the recipe's constructor." - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "id": "5ddf41e7", - "metadata": {}, - "outputs": [], - "source": [ - "from transformer_engine.common.recipe import DelayedScaling, Float8CurrentScaling, MXFP8BlockScaling, NVFP4BlockScaling\n", - "from transformer_engine.jax import flax as te_flax \n", - "\n", - "# Choose a quantization recipe. This can be modified to any of the recipes imported above.\n", - "quantization_recipe = DelayedScaling()\n", - "\n", - "te_dot_general_cls = te_flax.make_dot_general_cls(quantization_recipe)\n", - "\n", - "rngs = {'dropout': dropout_key}\n", - "if isinstance(quantization_recipe, NVFP4BlockScaling):\n", - " # The NVFP4 recipe requires a Flax RNG for stochastic rounding\n", - " rngs['sr_rng'] = jax.random.PRNGKey(0)\n" - ] - }, - { - "cell_type": "markdown", - "id": "c8769655", - "metadata": {}, - "source": [ - "Now using this quantized dense in our model is as simple as passing in `dot_general_fn=te_dot_general`. Let's try it out!\n", - "\n", - "
\n", - "\n", - "Important: Remat Policy\n", - "\n", - "TE's quantization uses specialized TE quantized GEMM primitives. If you are using any built-in JAX checkpoint policies that look for JAX GEMMs (dots), such as `jax.checkpoint_policies.checkpoint_dots`, please replace the policy with `transformer_engine.jax.checkpoint_policies.checkpoint_dots_and_te_gemms` or similar policies to ensure TE's quantized GEMM primitives are checkpointed correctly.\n", - "\n", - "If this is not performed, TE GEMMs will be rematerialized introducing an incorrect performance comparison.\n", - "\n", - "
" - ] - }, - { - "cell_type": "code", - "execution_count": 12, - "id": "8407d2ea", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Pure Flax FlaxTransformerLayer initialized successfully!\n", - "Parameter shapes: {'Dense_0': {'bias': (12288,), 'kernel': (4096, 12288)}, 'Dense_1': {'bias': (4096,), 'kernel': (4096, 4096)}, 'FlaxMLP_0': {'Dense_0': {'bias': (16384,), 'kernel': (4096, 16384)}, 'Dense_1': {'bias': (4096,), 'kernel': (16384, 4096)}}, 'LayerNorm_0': {'bias': (4096,), 'scale': (4096,)}, 'LayerNorm_1': {'bias': (4096,), 'scale': (4096,)}}\n", - "Additional state: {'_overwrite_with_gradient': {'FlaxMLP_0': {'TEWrapper_dot_general_0': {'grad_amax_history': (1024,), 'grad_scale': (1,), 'kernel_amax_history': (1024,), 'kernel_scale': (1,), 'x_amax_history': (1024,), 'x_scale': (1,)}, 'TEWrapper_dot_general_1': {'grad_amax_history': (1024,), 'grad_scale': (1,), 'kernel_amax_history': (1024,), 'kernel_scale': (1,), 'x_amax_history': (1024,), 'x_scale': (1,)}}, 'TEWrapper_dot_general_0': {'grad_amax_history': (1024,), 'grad_scale': (1,), 'kernel_amax_history': (1024,), 'kernel_scale': (1,), 'x_amax_history': (1024,), 'x_scale': (1,)}, 'TEWrapper_dot_general_1': {'grad_amax_history': (1024,), 'grad_scale': (1,), 'kernel_amax_history': (1024,), 'kernel_scale': (1,), 'x_amax_history': (1024,), 'x_scale': (1,)}}}\n" - ] - } - ], - "source": [ - "# Initialize the FlaxTransformerLayer\n", - "flax_transformer = FlaxTransformerLayer(\n", - " hidden_size=hidden_size,\n", - " ffn_hidden_size=ffn_hidden_size,\n", - " num_attention_heads=num_attention_heads,\n", - " dot_general_cls=te_dot_general_cls,\n", - ")\n", - "\n", - "# Initialize parameters\n", - "var_collect = flax_transformer.init(key, x, attention_mask=None, deterministic=False)\n", - "\n", - "print(\"Pure Flax FlaxTransformerLayer initialized successfully!\")\n", - "print(f\"Parameter shapes: {jax.tree_util.tree_map(lambda x: x.shape, var_collect['params'])}\")\n", - "print(f\"Additional state: {jax.tree_util.tree_map(lambda x: x.shape, {k: v for k, v in var_collect.items() if k != 'params'})}\")" - ] - }, - { - "cell_type": "markdown", - "id": "abe27237", - "metadata": {}, - "source": [ - "If using a recipe that stores additional state, such as `DelayedScaling`, you'll see this additional state stored as Flax variables. It is important to maintain and pass the whole state of Flax variables `var_collect` across training steps, not just the model params, for proper usage of stateful recipes like `DelayedScaling`.\n", - "\n", - "For example, above inside `Additional state: ` you'll see the `amax_history` of each quantization which is used to compute the per-tensor scale in the `DelayedScaling` recipe." - ] - }, - { - "cell_type": "markdown", - "id": "5ab72935", - "metadata": {}, - "source": [ - "The reason we need `te_dot_general_cls` as a Flax module instead of a module-less function like `jax.lax.dot_general` is for some quantization recipes to track internal state separate from model parameters.\n", - "\n", - "Flax modules can manage 3 things:\n", - "1. Model parameters/weights, e.g. your Dense \"kernel\", \"bias\", etc.\n", - "2. RNGs for dropout, stochastic rounding, etc.\n", - "3. Flax variables. These are additional state variables that are used across training steps but are distinct from model params in that you don't take gradients or optimize them. Currently, we only use this for DelayedScaling's amax_history state\n", - "\n", - "With the simplest quantization integration shown in this tutorial, we want users to keep their existing model param setup so they don't need to worry about preserving the sharding, init distribution, etc.. So we don't need point 1 since we don't do model param creation in this codepath with dot_general_cls, but we still do need `te_dot_general_cls()` to produce a Flax module since we potentially need to do points 2 or 3 which need to be in a Flax module." - ] - }, - { - "cell_type": "code", - "execution_count": 13, - "id": "3b6b344b", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Input shape: (4, 2048, 4096)\n", - "Output shape: (4, 2048, 4096)\n", - "Output dtype: float32\n", - "Forward pass completed successfully!\n" - ] - } - ], - "source": [ - "# Example usage of forward pass\n", - "y = flax_transformer.apply(var_collect, x, attention_mask=None, deterministic=True, rngs=rngs)\n", - "print(f\"Input shape: {x.shape}\")\n", - "print(f\"Output shape: {y.shape}\")\n", - "print(f\"Output dtype: {y.dtype}\")\n", - "print(\"Forward pass completed successfully!\")\n" - ] - }, - { - "cell_type": "markdown", - "id": "d178f247", - "metadata": {}, - "source": [ - "Now let's measure the performance!" - ] - }, - { - "cell_type": "code", - "execution_count": 14, - "id": "5cc6c2a7", - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Mean time: 10.553865432739258 ms\n" - ] - } - ], - "source": [ - "import importlib\n", - "import quickstart_jax_utils\n", - "importlib.reload(quickstart_jax_utils)\n", - "\n", - "utils.speedometer(\n", - " model_apply_fn=flax_transformer.apply,\n", - " variables=var_collect,\n", - " input=x,\n", - " output_grad=dy,\n", - " forward_kwargs={\"attention_mask\": None, \"deterministic\": False},\n", - " rngs=rngs,\n", - ")" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": "Python 3 (ipykernel)", - "language": "python", - "name": "python3" - } - }, - "nbformat": 4, - "nbformat_minor": 5 -} diff --git a/docs/examples/te_jax_integration.rst b/docs/examples/te_jax_integration.rst new file mode 100644 index 000000000..a15a10e0b --- /dev/null +++ b/docs/examples/te_jax_integration.rst @@ -0,0 +1,95 @@ +.. + Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + See LICENSE for license information. + +JAX: Integrating TransformerEngine into an existing framework +============================================================= + +This is the landing page for a series of focused documents on bringing +TransformerEngine into a JAX+Flax codebase one optimization at a time. Each +linked page isolates a single feature so you can see exactly what changes are +required and what are the performance benefits. + +Pick a topic +------------ + +.. list-table:: + :header-rows: 1 + :widths: 25, 15, 60 + + * - Document + - Status + - Covers + * - `Dense GEMMs `_ + - **Available** + - ``nn.Dense`` → quantized GEMM; single-GPU speedup; multi-GPU speedup; + * - `Collective GEMMs `_ + - *Coming soon* + - + * - `Attention `_ + - *Coming soon* + - + * - `Expert Parallelism `_ + - *Coming soon* + - + + +Quantization recipes at a glance +-------------------------------- + +TE exposes its quantization choices as **recipes**. Please see +`Low-precision Training +`_ +for a more detailed description of each recipe. + +.. _jax_recipe_table_overview: +.. list-table:: + :header-rows: 1 + :widths: 25, 15, 30, 30 + + * - Recipe + - Hardware + - State + - Description + * - ``MXFP8BlockScaling`` + - Blackwell+ + - none + - Block-scaled FP8 (32-element blocks) + * - ``NVFP4BlockScaling`` + - Blackwell+ + - requires a Flax RNG ``sr_rng`` + - FP4 with 2D block scaling and stochastic rounding + * - ``DelayedScaling`` + - Hopper+ + - amax history (Flax variables) + - Per-tensor FP8 with amax history + * - ``Float8CurrentScaling`` + - Hopper+ + - none + - Per-tensor FP8 without an amax history + +Import them from ``transformer_engine.common.recipe``. + + +Conventions used across these documents +--------------------------------------- + +* **Framework.** Flax Linen. (TE/JAX uses Linen; see + `Flax NNX/Linen interop + `_ and + `Haiku/Flax interop + `_ if you're on + a different stack.) +* **Baseline dtype.** bf16 for inputs and parameters. +* **Benchmarking.** ``quickstart_jax_utils.speedometer`` runs a JIT-compiled + fwd+bwd loop with warmup + + +.. toctree:: + :hidden: + + jax/dense + jax/collective_gemm + jax/attention + jax/expert_parallelism diff --git a/docs/index.rst b/docs/index.rst index 738955367..53c4b0e37 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -57,7 +57,7 @@ Transformer Engine documentation examples/te_llama/tutorial_accelerate_hf_llama_with_te.ipynb examples/te_gemma/tutorial_generation_gemma_with_te.ipynb examples/onnx/onnx_export.ipynb - examples/te_jax_integration.ipynb + examples/te_jax_integration.rst examples/op_fuser/op_fuser.rst .. toctree:: diff --git a/qa/L0_jax_distributed_unittest/test.sh b/qa/L0_jax_distributed_unittest/test.sh index 3f2581660..c62d7a4ba 100644 --- a/qa/L0_jax_distributed_unittest/test.sh +++ b/qa/L0_jax_distributed_unittest/test.sh @@ -37,6 +37,11 @@ wait TE_PATH=$TE_PATH bash $TE_PATH/examples/jax/collective_gemm/run_test_cgemm.sh || test_fail "run_test_cgemm.sh" wait +# Exercise the multi-GPU tutorial in docs/examples/jax (needs >= 4 GPUs; +# auto-skips otherwise). +CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_docs_examples_jax_distributed.xml -k multi_gpu $TE_PATH/docs/examples/jax/ || test_fail "docs/examples/jax (multi-GPU)" +wait + if [ $RET -ne 0 ]; then echo "Error: some sub-tests failed: $FAILED_CASES" exit 1 diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh index 3453e35d2..e4bcdc4e5 100644 --- a/qa/L0_jax_unittest/test.sh +++ b/qa/L0_jax_unittest/test.sh @@ -42,6 +42,11 @@ python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/py export XLA_FLAGS="${XLA_FLAGS} --xla_gpu_deterministic_ops" NVTE_JAX_CUSTOM_CALLS="false" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_test_single_gpu_encoder_without_custom_call.xml $TE_PATH/examples/jax/encoder/test_single_gpu_encoder.py || test_fail "test_single_gpu_encoder.py without custom calls" +# Exercise the docs/examples/jax tutorials. The multi-GPU tests are +# skipped at runtime when fewer than 4 devices are visible, so this is safe on +# single-GPU runners. +python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_docs_examples_jax.xml $TE_PATH/docs/examples/jax/ || test_fail "docs/examples/jax" + if [ $RET -ne 0 ]; then echo "Error: some sub-tests failed: $FAILED_CASES" exit 1 From d95b34c8516fa11e3b126d7bb93082a694ce52be Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Thu, 21 May 2026 12:02:36 -0700 Subject: [PATCH 072/180] Fix the permissions in the automatic labeler (#3029) Fix the permissions Signed-off-by: Przemek Tredak --- .github/workflows/community_label.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/community_label.yml b/.github/workflows/community_label.yml index 52afa2477..c09debb87 100644 --- a/.github/workflows/community_label.yml +++ b/.github/workflows/community_label.yml @@ -12,6 +12,7 @@ on: permissions: contents: read issues: write + pull-requests: write jobs: label: From 82776bc04ac20fb070f5c5863b3bdea221ccaa21 Mon Sep 17 00:00:00 2001 From: Muu Date: Fri, 22 May 2026 03:05:05 +0800 Subject: [PATCH 073/180] refactor(distributed): deduplicate TE module class lookups with caching (#2992) - Extract common get_te_classes() with @lru_cache for reuse - Refactor has_te_modules() and _is_te_module() to use tuple isinstance check - Remove duplicated import lists across multiple functions Signed-off-by: Muu --- transformer_engine/pytorch/distributed.py | 48 +++++++++-------------- 1 file changed, 18 insertions(+), 30 deletions(-) diff --git a/transformer_engine/pytorch/distributed.py b/transformer_engine/pytorch/distributed.py index a0d4ac353..670eecaa5 100644 --- a/transformer_engine/pytorch/distributed.py +++ b/transformer_engine/pytorch/distributed.py @@ -613,18 +613,21 @@ def get_activation_recompute_contexts(): return forward_ctx, recompute_ctx -def has_te_modules(network): +@lru_cache +def get_te_classes(): """ - Check if there are any Transformer Engine modules in the network. + Return all Transformer Engine modules. """ from .module import LayerNorm, RMSNorm from .module.base import TransformerEngineBaseModule + from .attention.dot_product_attention.dot_product_attention import ( + DotProductAttention, + ) from .attention.dot_product_attention.backends import UnfusedDotProductAttention - from .attention.dot_product_attention.dot_product_attention import DotProductAttention from .attention.multi_head_attention import MultiheadAttention from .transformer import TransformerLayer - te_classes_list = [ + return ( LayerNorm, RMSNorm, TransformerEngineBaseModule, @@ -632,12 +635,17 @@ def has_te_modules(network): DotProductAttention, MultiheadAttention, TransformerLayer, - ] + ) + +def has_te_modules(network): + """ + Check if there are any Transformer Engine modules in the network. + """ + te_classes = get_te_classes() if isinstance(network, torch.nn.Module): - for module in network.modules(): - if any(isinstance(module, te_class) for te_class in te_classes_list): - return True + if any(isinstance(module, te_classes) for module in network.modules()): + return True return False # Cannot check for TE modules inside a custom class/callable that's not a torch.nn.Module, @@ -2040,28 +2048,8 @@ def _is_te_module(module): Check if given module is a Transformer Engine module that requires the TE checkpoint implementation for activation recompute. """ - from .module import LayerNorm, RMSNorm - from .module.base import TransformerEngineBaseModule - from .attention.dot_product_attention.dot_product_attention import DotProductAttention - from .attention.dot_product_attention.backends import UnfusedDotProductAttention - from .attention.multi_head_attention import MultiheadAttention - from .transformer import TransformerLayer - - te_classes_list = [ - LayerNorm, - RMSNorm, - TransformerEngineBaseModule, - UnfusedDotProductAttention, - DotProductAttention, - MultiheadAttention, - TransformerLayer, - ] - is_te_module = False - for te_class in te_classes_list: - if isinstance(module, te_class): - is_te_module = True - break - return is_te_module + te_classes = get_te_classes() + return isinstance(module, te_classes) def prepare_te_modules_for_fsdp(fsdp_root: torch.nn.Module) -> None: From 390eac83b532d9c2776a05d02c00e7b006387b79 Mon Sep 17 00:00:00 2001 From: Przemyslaw Tredak Date: Thu, 21 May 2026 15:51:04 -0700 Subject: [PATCH 074/180] Fixes to the community labeling GitHub Action (#3030) * Debugging Signed-off-by: Przemek Tredak * More debugging Signed-off-by: Przemek Tredak * Distinguish the users based on the write permissions rather than relying on the member field, which could be set to private Signed-off-by: Przemek Tredak * Fix Signed-off-by: Przemek Tredak --------- Signed-off-by: Przemek Tredak --- .github/workflows/community_label.yml | 35 ++++++++++++--------------- 1 file changed, 16 insertions(+), 19 deletions(-) diff --git a/.github/workflows/community_label.yml b/.github/workflows/community_label.yml index c09debb87..c0d31d2a4 100644 --- a/.github/workflows/community_label.yml +++ b/.github/workflows/community_label.yml @@ -33,30 +33,27 @@ jobs: const isOrgMember = association === "MEMBER" || association === "OWNER"; + let permission = "none"; + + try { + const res = await github.rest.repos.getCollaboratorPermissionLevel({ + owner: context.repo.owner, + repo: context.repo.repo, + username: user, + }); + permission = res.data.permission; + } catch (e) { + if (e.status !== 404) throw e; + } + + const isCore = permission === "write" || permission === "admin"; if (!isOrgMember) { targetLabel = communityLabel; } else { - let permission = "none"; - - try { - const res = await github.rest.repos.getCollaboratorPermissionLevel({ - owner: context.repo.owner, - repo: context.repo.repo, - username: user, - }); - permission = res.data.permission; - } catch (e) { - if (e.status !== 404) throw e; - } - - const isCore = permission === "write" || permission === "admin"; - - if (!isCore) { - targetLabel = orgLabel; - } + targetLabel = orgLabel; } - if (targetLabel) { + if (!isCore) { await github.rest.issues.addLabels({ owner: context.repo.owner, repo: context.repo.repo, From 1bd99646ddff278f81480aea331602ee3fe7e962 Mon Sep 17 00:00:00 2001 From: sraman-rgb Date: Thu, 21 May 2026 18:51:04 -0500 Subject: [PATCH 075/180] GGEMM+srelu kernels for MxFP8 Nemotron (#2981) * Add MXFP8 grouped MLP SReLU fusion Signed-off-by: sraman-rgb * Address grouped MLP fused op review comments Signed-off-by: Siddhartha Raman S * Avoid quantizing ScaledSReLU backward in basic op Signed-off-by: Siddhartha Raman S * Wire ScaledSReLU recompute in grouped MLP Signed-off-by: Siddhartha Raman S * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address grouped MLP ScaledSReLU review comments Signed-off-by: Siddhartha Raman S * Gate ScaledSReLU recompute support Signed-off-by: Siddhartha Raman S * Use version check for dSReLU reuse arg Signed-off-by: Siddhartha Raman S * Reuse forward dSReLU recompute decision Signed-off-by: Siddhartha Raman S * Reject activation recompute without grouped MLP fusion Signed-off-by: Siddhartha Raman S * Rename grouped MLP activation recompute flag Signed-off-by: Siddhartha Raman S --------- Signed-off-by: sraman-rgb Signed-off-by: Siddhartha Raman S Signed-off-by: Siddhartha Raman S Co-authored-by: Siddhartha Raman S Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: vthumbe1503 --- tests/pytorch/test_fusible_ops.py | 142 +++++++++++--- transformer_engine/pytorch/ops/_common.py | 84 ++++++--- .../pytorch/ops/basic/__init__.py | 1 + .../pytorch/ops/basic/activation.py | 119 +++++++++++- .../pytorch/ops/basic/swiglu.py | 32 +++- .../pytorch/ops/fused/__init__.py | 6 +- .../pytorch/ops/fused/backward_grouped_mlp.py | 173 +++++++++++++---- .../pytorch/ops/fused/forward_grouped_mlp.py | 176 +++++++++++++----- 8 files changed, 591 insertions(+), 142 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 7691582f9..3a3aa8be9 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -22,6 +22,7 @@ import transformer_engine.pytorch.ops as te_ops from transformer_engine.pytorch.ops._common import ( _cudnn_frontend_version_supported, + is_glu_activation, ) from transformer_engine.pytorch.ops.fused import ( @@ -2480,6 +2481,59 @@ def test_scaled_swiglu( assert_close_grads(x_test, x_ref, **tols) assert_close_grads(scales_test, scales_ref, **tols) + @pytest.mark.parametrize("in_shape", ((71, 192), (5, 7, 128))) + @pytest.mark.parametrize("input_requires_grad", (False, True)) + @pytest.mark.parametrize("scales_requires_grad", (False, True)) + def test_scaled_srelu( + self, + *, + in_shape: Iterable[int], + dtype: torch.dtype = torch.float32, + device: torch.device = "cuda", + input_requires_grad: bool, + scales_requires_grad: bool, + ) -> None: + """SReLU with post-scale""" + + # Random data + x_ref, x_test = make_reference_and_test_tensors( + in_shape, + test_dtype=dtype, + test_device=device, + requires_grad=input_requires_grad, + ) + scales_ref, scales_test = make_reference_and_test_tensors( + in_shape[:-1], + test_dtype=dtype, + test_device=device, + requires_grad=scales_requires_grad, + ) + dy_ref, dy_test = make_reference_and_test_tensors( + in_shape, + test_dtype=dtype, + test_device=device, + requires_grad=False, + ) + + # Plain PyTorch implementation + y = torch.nn.functional.relu(x_ref).square() + y_ref = scales_ref.unsqueeze(-1) * y + if input_requires_grad or scales_requires_grad: + y_ref.backward(dy_ref) + + # Implementation with fusible operation + op = te_ops.ScaledSReLU() + y_test = op(x_test, scales_test) + if input_requires_grad or scales_requires_grad: + y_test.backward(dy_test) + + # Check results + tols = dtype_tols(dtype) + y_test = y_test.to(dtype=torch.float64, device="cpu") + assert_close(y_test, y_ref, **tols) + assert_close_grads(x_test, x_ref, **tols) + assert_close_grads(scales_test, scales_ref, **tols) + def test_interleaved_scaled_swiglu(self): """SwiGLU with post-scale and block interleaved input format""" self.test_scaled_swiglu( @@ -2489,6 +2543,15 @@ def test_interleaved_scaled_swiglu(self): scales_requires_grad=True, ) + @pytest.mark.parametrize( + "op_cls", + (te_ops.ScaledSwiGLU, te_ops.ScaledSReLU, te_ops.ScaledClampedQGeGLU), + ) + def test_scaled_activation_recompute_in_mlp_config(self, op_cls) -> None: + """Scaled activations expose a per-op recompute knob.""" + assert op_cls().activation_recompute_in_mlp is False + assert op_cls(activation_recompute_in_mlp=True).activation_recompute_in_mlp is True + @pytest.mark.parametrize("in_shape", ((71, 192), (5, 7, 128))) @pytest.mark.parametrize("input_requires_grad", (False, True)) @pytest.mark.parametrize("scales_requires_grad", (False, True)) @@ -3570,7 +3633,9 @@ def test_layernorm_mlp( @pytest.mark.parametrize("glu_interleave_size", (None, 32)) @pytest.mark.parametrize("delay_wgrad_compute", (False, True)) @pytest.mark.parametrize("hidden_size", (128, 256)) - @pytest.mark.parametrize("activation", ("scaled_swiglu", "scaled_clamped_qgeglu")) + @pytest.mark.parametrize( + "activation", ("scaled_swiglu", "scaled_clamped_qgeglu", "scaled_srelu") + ) def test_grouped_mlp( self, *, @@ -3588,7 +3653,7 @@ def test_grouped_mlp( delay_wgrad_compute: bool, activation: str, ) -> None: - """GroupedLinear + ScaledSwiGLU / ScaledClampedQGeGLU + GroupedLinear""" + """GroupedLinear + scaled activation + GroupedLinear""" # Split sizes split_sizes = [split_alignment * (i) for i in range(group_size)] @@ -3601,6 +3666,15 @@ def test_grouped_mlp( # Skip invalid configurations with_quantization = quantization is not None + if activation == "scaled_swiglu": + scaled_act = te_ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) + elif activation == "scaled_clamped_qgeglu": + scaled_act = te_ops.ScaledClampedQGeGLU(glu_interleave_size=glu_interleave_size) + elif activation == "scaled_srelu": + scaled_act = te_ops.ScaledSReLU() + else: + raise ValueError(f"Unexpected grouped MLP activation ({activation})") + activation_is_glu = is_glu_activation(scaled_act) maybe_skip_quantization(quantization, dims=in_shape, device=device, dtype=dtype) if single_grouped_weight and quantization != "mxfp8": pytest.skip("single_grouped_weight is only supported for MXFP8 quantization") @@ -3608,9 +3682,14 @@ def test_grouped_mlp( pytest.skip("single_grouped_bias requires bias=True") if with_quantization and dtype not in (torch.bfloat16, torch.float16): pytest.skip("Quantized group GEMM is only supported with BF16/FP16") + if not activation_is_glu and quantization != "mxfp8": + pytest.skip("Scaled unary grouped MLP fusion is only supported with MXFP8") + if not activation_is_glu and glu_interleave_size is not None: + pytest.skip("Unary activations do not use GLU interleaving") if quantization == "nvfp4" and activation == "scaled_clamped_qgeglu" and bias: # TODO: ksivaman: Need to debug numerics for this case. pytest.skip("Bias/dbias not yet supported in NVFP4 fused grouped MLP with GeGLU") + fc1_out_features = 2 * hidden_size if activation_is_glu else hidden_size # Random data x_ref, x_test = make_reference_and_test_tensors( @@ -3641,7 +3720,7 @@ def test_grouped_mlp( fc2_bs_ref, fc2_bs_test = [], [] for _ in range(group_size): fc1_w_ref, fc1_w_test = make_reference_and_test_tensors( - (2 * hidden_size, hidden_size), + (fc1_out_features, hidden_size), min=-0.25, max=0.25, quantization=quantization, @@ -3660,7 +3739,7 @@ def test_grouped_mlp( fc2_b_ref, fc2_b_test = None, None if bias: fc1_b_ref, fc1_b_test = make_reference_and_test_tensors( - (2 * hidden_size,), + (fc1_out_features,), min=-0.5, max=0.5, test_dtype=dtype, @@ -3689,7 +3768,7 @@ def test_grouped_mlp( for group_idx in range(group_size): x = xs[group_idx] x = torch.nn.functional.linear(x, fc1_ws_ref[group_idx], bias=fc1_bs_ref[group_idx]) - if glu_interleave_size is not None: + if activation_is_glu and glu_interleave_size is not None: x = x.reshape( -1, 2 * hidden_size // (2 * glu_interleave_size), @@ -3698,15 +3777,20 @@ def test_grouped_mlp( ) x = x.transpose(1, 2) x = x.reshape(-1, 2 * hidden_size) - x1, x2 = x.chunk(2, dim=-1) if activation == "scaled_swiglu": + x1, x2 = x.chunk(2, dim=-1) x = torch.nn.functional.silu(x1) * x2 - else: + elif activation == "scaled_clamped_qgeglu": + x1, x2 = x.chunk(2, dim=-1) lim = torch.tensor(7.0, device=x1.device, dtype=x1.dtype) geglu_alpha = 1.702 x1c = torch.minimum(x1, lim) x2c = torch.clamp(x2, -lim, lim) x = (x2c + 1) * (x1c * torch.sigmoid(geglu_alpha * x1c)) + elif activation == "scaled_srelu": + x = torch.nn.functional.relu(x).square() + else: + raise ValueError(f"Unexpected grouped MLP activation ({activation})") x = x * probs[group_idx].unsqueeze(-1) x = torch.nn.functional.linear(x, fc2_ws_ref[group_idx]) if bias: @@ -3717,16 +3801,11 @@ def test_grouped_mlp( # Construct operations recipe = make_recipe(quantization) - scaled_act = ( - te_ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) - if activation == "scaled_swiglu" - else te_ops.ScaledClampedQGeGLU(glu_interleave_size=glu_interleave_size) - ) with te.quantized_model_init(enabled=with_quantization, recipe=recipe): fc1 = te_ops.GroupedLinear( group_size, hidden_size, - 2 * hidden_size, + fc1_out_features, bias=bias, device=device, dtype=dtype, @@ -3810,22 +3889,31 @@ def test_grouped_mlp( if ( quantization == "mxfp8" and dtype in (torch.bfloat16, torch.float16) - and glu_interleave_size == 32 + and ( + (not activation_is_glu and glu_interleave_size is None) + or (activation_is_glu and glu_interleave_size == 32) + ) and _cudnn_frontend_version_supported() ): - if te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported(): + if activation_is_glu: + forward_cls = te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU_MXFP8 + backward_cls = te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8 + else: + forward_cls = te_ops.fused.ForwardGroupedMLP_CuTeGEMMUnary_MXFP8 + backward_cls = te_ops.fused.BackwardGroupedMLP_CuTeGEMMDUnary_MXFP8 + if forward_cls.is_supported(): forward_ops = module._module_groups[0]._forward_ops assert len(forward_ops) == 1 assert isinstance( forward_ops[0][0], - te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8, + forward_cls, ) - if te_ops.fused.BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8.is_supported(): + if backward_cls is not None and backward_cls.is_supported(): backward_ops = module._module_groups[0]._backward_ops assert len(backward_ops) == 1 assert isinstance( backward_ops[0][0], - te_ops.fused.BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8, + backward_cls, ) # Loose tols for sanity checking @@ -3910,9 +3998,9 @@ def test_grouped_mlp_single_weight_numerics( ) -> None: """single_grouped_weight=True/False should match exactly for fused MXFP8 grouped MLP.""" - if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported(): + if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU_MXFP8.is_supported(): pytest.skip("MXFP8 fused grouped MLP forward is not supported on this system") - if not te_ops.fused.BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8.is_supported(): + if not te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8.is_supported(): pytest.skip("MXFP8 fused grouped MLP backward is not supported on this system") split_sizes = [split_alignment * (i + 1) for i in range(group_size)] @@ -4014,12 +4102,12 @@ def _run_case(single_grouped_weight: bool) -> tuple[torch.Tensor, ...]: assert len(forward_ops) == 1 assert isinstance( forward_ops[0][0], - te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8, + te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU_MXFP8, ) assert len(backward_ops) == 1 assert isinstance( backward_ops[0][0], - te_ops.fused.BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8, + te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8, ) if single_grouped_weight: @@ -4132,9 +4220,9 @@ def test_grouped_mlp_overwrite_main_grad( that read ``.grad`` don't see stale bytes from the cached dummy). """ - if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported(): + if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU_MXFP8.is_supported(): pytest.skip("MXFP8 fused grouped MLP forward is not supported on this system") - if not te_ops.fused.BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8.is_supported(): + if not te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8.is_supported(): pytest.skip("MXFP8 fused grouped MLP backward is not supported on this system") recipe = make_recipe("mxfp8") @@ -4266,7 +4354,7 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8( ) -> None: """Grouped MLP forward+backward should be CUDA graph capturable (MXFP8).""" - if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported(): + if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU_MXFP8.is_supported(): pytest.skip("MXFP8 fused grouped MLP is not supported on this system") if dtype not in (torch.bfloat16, torch.float16): pytest.skip("MXFP8 fused grouped MLP is only supported with BF16/FP16") @@ -4408,12 +4496,12 @@ def train_step( assert len(forward_ops) == 1 assert isinstance( forward_ops[0][0], - te_ops.fused.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8, + te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU_MXFP8, ) assert len(backward_ops) == 1 assert isinstance( backward_ops[0][0], - te_ops.fused.BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8, + te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8, ) fresh_x = torch.randn_like(static_x) diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 9325d87ae..c0474220e 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -21,16 +21,31 @@ from ..utils import canonicalize_dtype -@functools.lru_cache(maxsize=1) +@functools.lru_cache(maxsize=None) +def _cudnn_frontend_version_at_least(min_version: str) -> bool: + """Check cuDNN frontend package version.""" + try: + return PkgVersion(get_pkg_version("nvidia-cudnn-frontend")) >= PkgVersion(min_version) + except PackageNotFoundError: + return False + + def _cudnn_frontend_version_supported() -> bool: """Check cuDNN frontend is at least 1.23.0. All grouped MLP fused-kernel features require cuDNN frontend 1.23.0. """ - try: - return PkgVersion(get_pkg_version("nvidia-cudnn-frontend")) >= PkgVersion("1.23.0") - except PackageNotFoundError: - return False + return _cudnn_frontend_version_at_least("1.23.0") + + +def _cudnn_frontend_supports_grouped_gemm_srelu() -> bool: + """Check cuDNN frontend min version for grouped GEMM SReLU kernels.""" + return _cudnn_frontend_version_at_least("1.24.0") + + +def _nvidia_cudnn_frontend_supports_wgrad() -> bool: + """Check cuDNN FE min version for grouped GEMM wgrad kernel.""" + return _cudnn_frontend_version_supported() def is_quantized_tensor(tensor: torch.Tensor | QuantizedTensorStorage) -> bool: @@ -182,8 +197,21 @@ def get_dummy_wgrads_for_params( return out -def validate_grouped_mlp_dims(fc1, glu_op, fc2) -> None: - """Validate FC1 / scaled GLU / FC2 dimensions for fused grouped MLP.""" +def is_glu_activation(activation_op) -> bool: + """Whether an activation consumes a GLU-style doubled input.""" + from .basic import ( # pylint: disable=import-outside-toplevel + ScaledClampedQGeGLU, + ScaledSwiGLU, + ) + + return isinstance(activation_op, (ScaledSwiGLU, ScaledClampedQGeGLU)) + + +def validate_grouped_mlp_dims(fc1, activation_op, fc2) -> None: + """Validate FC1 / activation / FC2 dimensions for fused grouped MLP.""" + from .basic import ( # pylint: disable=import-outside-toplevel + ScaledSReLU, + ) if fc1.in_features % 64 != 0 or fc1.out_features % 64 != 0: raise ValueError( @@ -195,17 +223,24 @@ def validate_grouped_mlp_dims(fc1, glu_op, fc2) -> None: f"Unsupported dims for FC2 (num_groups={fc2.num_groups}, " f"in_features={fc2.in_features}, out_features={fc2.out_features})." ) - if fc1.out_features != 2 * fc2.in_features or fc1.num_groups != fc2.num_groups: + if is_glu_activation(activation_op): + expected_fc1_out_features = 2 * fc2.in_features + elif isinstance(activation_op, ScaledSReLU): + expected_fc1_out_features = fc2.in_features + else: + raise TypeError(f"Unsupported grouped MLP activation ({activation_op.__class__.__name__}).") + + if fc1.out_features != expected_fc1_out_features or fc1.num_groups != fc2.num_groups: raise ValueError( f"FC1 (num_groups={fc1.num_groups}, in_features={fc1.in_features}, " f"out_features={fc1.out_features}) " f"and FC2 (num_groups={fc2.num_groups}, in_features={fc2.in_features}, " f"out_features={fc2.out_features}) do not match." ) - if glu_op.glu_interleave_size != 32: + if is_glu_activation(activation_op) and activation_op.glu_interleave_size != 32: raise ValueError( "Fused kernel requires 32-wide GLU interleaving, " - f"but got glu_interleave_size={glu_op.glu_interleave_size}." + f"but got glu_interleave_size={activation_op.glu_interleave_size}." ) @@ -214,8 +249,9 @@ def fuse_grouped_mlp_ops( *, recipe, fused_op_cls, + activation_op_types=None, ): - """Sliding-window fusion for GroupedLinear + scaled GLU + GroupedLinear. + """Sliding-window fusion for GroupedLinear + activation + GroupedLinear. Parameters ---------- @@ -225,9 +261,7 @@ def fuse_grouped_mlp_ops( Quantization recipe. fused_op_cls : type Fused operation class with ``is_supported()`` classmethod and - constructor accepting ``fc1``, ``glu_op``, ``fc2`` keyword args. The - ``glu_op`` must be :class:`~transformer_engine.pytorch.ops.basic.swiglu.ScaledSwiGLU` - or :class:`~transformer_engine.pytorch.ops.basic.swiglu.ScaledClampedQGeGLU`. + constructor accepting ``fc1``, ``activation``, and ``fc2`` keyword args. Returns ------- @@ -244,6 +278,8 @@ def fuse_grouped_mlp_ops( return ops if recipe is None or not recipe.mxfp8(): return ops + if activation_op_types is None: + activation_op_types = (ScaledSwiGLU, ScaledClampedQGeGLU) out = [] window, ops = ops[:3], ops[3:] @@ -252,7 +288,7 @@ def fuse_grouped_mlp_ops( matches_pattern = True if not ( isinstance(window[0], GroupedLinear) - and isinstance(window[1], (ScaledSwiGLU, ScaledClampedQGeGLU)) + and isinstance(window[1], activation_op_types) and isinstance(window[2], GroupedLinear) ): matches_pattern = False @@ -260,22 +296,16 @@ def fuse_grouped_mlp_ops( abs(window[1]._clamped.alpha - 1.702) > 0.001 ): matches_pattern = False - elif window[0].num_groups != window[2].num_groups: - matches_pattern = False - elif ( - window[0].in_features % 64 != 0 - or window[0].out_features % 64 != 0 - or window[2].in_features % 64 != 0 - or window[2].out_features % 64 != 0 - ): - matches_pattern = False - elif window[1].glu_interleave_size != 32: - matches_pattern = False + else: + try: + validate_grouped_mlp_dims(window[0], window[1], window[2]) + except (TypeError, ValueError): + matches_pattern = False if matches_pattern: op = fused_op_cls( fc1=window[0], - swiglu=window[1], + activation=window[1], fc2=window[2], ) window = [op] diff --git a/transformer_engine/pytorch/ops/basic/__init__.py b/transformer_engine/pytorch/ops/basic/__init__.py index 45c938ede..6def36ffc 100644 --- a/transformer_engine/pytorch/ops/basic/__init__.py +++ b/transformer_engine/pytorch/ops/basic/__init__.py @@ -13,6 +13,7 @@ ReLU, ReGLU, SReLU, + ScaledSReLU, SReGLU, SiLU, ) diff --git a/transformer_engine/pytorch/ops/basic/activation.py b/transformer_engine/pytorch/ops/basic/activation.py index 13cb519c1..eacc36b36 100644 --- a/transformer_engine/pytorch/ops/basic/activation.py +++ b/transformer_engine/pytorch/ops/basic/activation.py @@ -6,7 +6,8 @@ from __future__ import annotations import abc -from typing import Optional +from collections.abc import Iterable +from typing import Any, Optional import torch @@ -26,6 +27,7 @@ "ReLU", "ReGLU", "SReLU", + "ScaledSReLU", "SReGLU", "SiLU", ] @@ -345,6 +347,121 @@ def _activation_backward_impl(self, *args, **kwargs) -> torch.Tensor: return tex.dsrelu(*args, **kwargs) +class ScaledSReLU(BasicOperation): + r"""Squared ReLU with per-row post-scaling. + + If the SReLU output has shape ``(d_1, ..., d_n)``, it is multiplied + with an extra input tensor of shape ``(d_1, ..., d_{n-1})``. + + Parameters + ---------- + activation_recompute_in_mlp : bool, default = ``False`` + Enable fused grouped MLP kernels to recompute activation outputs + during backward when supported instead of saving them. + """ + + num_extra_inputs: int = 1 + + def __init__(self, *, activation_recompute_in_mlp: bool = False) -> None: + super().__init__() + self.activation_recompute_in_mlp: bool = activation_recompute_in_mlp + + def op_forward(self, *args, **kwargs) -> None: + raise RuntimeError( + f"{self.__class__.__name__} operation has " + f"{self.num_extra_inputs} extra tensor inputs " + f"and {self.num_extra_outputs} extra tensor outputs. " + "It overrides `fuser_forward` instead of `op_forward`." + ) + + def op_backward(self, *args, **kwargs) -> None: + raise RuntimeError( + f"{self.__class__.__name__} operation has " + f"{self.num_extra_inputs} extra tensor inputs " + f"and {self.num_extra_outputs} extra tensor outputs. " + "It overrides `fuser_backward` instead of `op_backward`." + ) + + def fuser_forward( + self, + basic_op_ctxs: list[OperationContext], + input_: torch.Tensor, + *, + basic_op_extra_inputs: list[tuple[torch.Tensor, ...]], + prev_op_grad_output_quantizer: Optional[Quantizer], # pylint: disable=unused-argument + next_op_input_quantizer: Optional[Quantizer], # pylint: disable=unused-argument + basic_op_kwargs: list[dict[str, Any]], # pylint: disable=unused-argument + ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + if self.activation_recompute_in_mlp: + raise RuntimeError( + f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " + "fused grouped MLP path." + ) + + extra_input = basic_op_extra_inputs[0][0] + + if torch.is_autocast_enabled(): + dtype = torch.get_autocast_dtype("cuda") + elif isinstance(input_, torch.Tensor): + dtype = input_.dtype + else: + dtype = extra_input.dtype + + x = maybe_dequantize(input_.contiguous(), dtype) + scales = maybe_dequantize(extra_input, dtype) + y = tex.srelu(x, None) * scales.unsqueeze(-1) + + ctx = basic_op_ctxs[0] + if ctx.requires_grad: + if is_cpu_offload_enabled(): + mark_activation_offload(x) + ctx.input_requires_grad = True + ctx.extra_input_requires_grad = extra_input.requires_grad + ctx.dtype = dtype + ctx.save_for_backward(x, scales) + + return y, [()] + + def fuser_backward( + self, + basic_op_ctxs: list[OperationContext], + grad_output: torch.Tensor, + *, + basic_op_grad_extra_outputs: list[tuple[torch.Tensor, ...]], + ) -> tuple[ + torch.Tensor, + Iterable[Iterable[Optional[torch.Tensor]]], + Iterable[Iterable[Optional[torch.Tensor]]], + ]: + del basic_op_grad_extra_outputs + + if self.activation_recompute_in_mlp: + raise RuntimeError( + f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " + "fused grouped MLP path." + ) + + ctx = basic_op_ctxs[0] + x, scales = ctx.saved_tensors + x = maybe_dequantize(x.contiguous(), ctx.dtype) + scales = maybe_dequantize(scales, ctx.dtype) + grad_output = maybe_dequantize(grad_output.contiguous(), ctx.dtype) + + grad_input = None + if ctx.input_requires_grad: + grad_srelu_out = grad_output * scales.unsqueeze(-1) + grad_input = tex.dsrelu(grad_srelu_out, x, None) + + grad_extra_input = None + if ctx.extra_input_requires_grad: + srelu_out = tex.srelu(x, None) + grad_extra_input = torch.linalg.vecdot(srelu_out, grad_output) + + clear_tensor_data(ctx.saved_tensors[0]) + + return grad_input, [()], [(grad_extra_input,)] + + class SReGLU(_ActivationOperation): r"""Squared Rectified Gated Linear Unit diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index 9c0bc86bc..9267d9bbb 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -369,9 +369,15 @@ class _ScaledGLU(BasicOperation): num_extra_inputs: int = 1 - def __init__(self, glu_interleave_size: Optional[int] = None) -> None: + def __init__( + self, + glu_interleave_size: Optional[int] = None, + *, + activation_recompute_in_mlp: bool = False, + ) -> None: super().__init__() self.glu_interleave_size: Optional[int] = glu_interleave_size + self.activation_recompute_in_mlp: bool = activation_recompute_in_mlp def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor: raise NotImplementedError @@ -409,6 +415,12 @@ def fuser_forward( next_op_input_quantizer: Optional[Quantizer], basic_op_kwargs: list[dict[str, Any]], ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: + if self.activation_recompute_in_mlp: + raise RuntimeError( + f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " + "fused grouped MLP path." + ) + extra_input = basic_op_extra_inputs[0][0] # Determine compute dtype @@ -465,6 +477,12 @@ def fuser_backward( Iterable[Iterable[Optional[torch.Tensor]]], Iterable[Iterable[Optional[torch.Tensor]]], ]: + if self.activation_recompute_in_mlp: + raise RuntimeError( + f"{self.__class__.__name__}(activation_recompute_in_mlp=True) requires the " + "fused grouped MLP path." + ) + ctx = basic_op_ctxs[0] input_, scales = ctx.saved_tensors input_ = maybe_dequantize(input_, ctx.dtype) @@ -526,6 +544,9 @@ class ScaledSwiGLU(_ScaledGLU): When set, the GLU activations will use an experimental block interleaved format. See the corresponding option in the SwiGLU operation for more details. + activation_recompute_in_mlp : bool, default = ``False`` + Enable fused grouped MLP kernels to recompute activation outputs + during backward when supported instead of saving them. """ @@ -553,6 +574,9 @@ class ScaledClampedQGeGLU(_ScaledGLU): glu_interleave_size : int, optional When set, the GLU activations will use an experimental block interleaved format. See :class:`ClampedSwiGLU`. + activation_recompute_in_mlp : bool, default = ``False`` + Enable fused grouped MLP kernels to recompute activation outputs + during backward when supported instead of saving them. limit : float, default ``7.0`` Clamp limit (see :class:`ClampedSwiGLU`). alpha : float, default ``1.702`` @@ -564,10 +588,14 @@ def __init__( self, glu_interleave_size: Optional[int] = None, *, + activation_recompute_in_mlp: bool = False, limit: float = 7.0, alpha: float = 1.702, ) -> None: - super().__init__(glu_interleave_size) + super().__init__( + glu_interleave_size, + activation_recompute_in_mlp=activation_recompute_in_mlp, + ) self._clamped: ClampedSwiGLU = ClampedSwiGLU( limit=limit, alpha=alpha, diff --git a/transformer_engine/pytorch/ops/fused/__init__.py b/transformer_engine/pytorch/ops/fused/__init__.py index 19a090f12..b29e35814 100644 --- a/transformer_engine/pytorch/ops/fused/__init__.py +++ b/transformer_engine/pytorch/ops/fused/__init__.py @@ -32,8 +32,10 @@ # Import experimental fusions # Note: Registration logic is non-trivial, so submodule handles it internally. from .forward_grouped_mlp import ( # pylint: disable=wrong-import-position - ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8, + ForwardGroupedMLP_CuTeGEMMGLU_MXFP8, + ForwardGroupedMLP_CuTeGEMMUnary_MXFP8, ) from .backward_grouped_mlp import ( # pylint: disable=wrong-import-position - BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8, + BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8, + BackwardGroupedMLP_CuTeGEMMDUnary_MXFP8, ) diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index a11d0505c..3b6330b22 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -18,15 +18,17 @@ from ...tensor.mxfp8_tensor import MXFP8Quantizer from ...utils import clear_tensor_data, get_cached_ones_tensor, get_device_compute_capability from ...constants import MXFP8_BLOCK_SCALING_SIZE -from ..basic import GroupedLinear, ScaledClampedQGeGLU, ScaledSwiGLU +from ..basic import GroupedLinear, ScaledSReLU, ScaledClampedQGeGLU from ..fuser import register_backward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext from .._common import ( _cudnn_frontend_version_supported, + _cudnn_frontend_supports_grouped_gemm_srelu, fuse_grouped_mlp_ops, get_accumulate_flag_in_param, get_dummy_wgrads_for_params, get_main_grad_from_param, + is_glu_activation, maybe_dequantize, view_main_grad_as_grouped_buffer, validate_grouped_mlp_dims, @@ -248,20 +250,17 @@ def _compute_grad_params( return w_list + bias_list -class BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8(FusedOperation): - """Fused op for MXFP8 GroupedLinear + ScaledSwiGLU or ScaledClampedQGeGLU + GroupedLinear +class _BackwardGroupedMLP_CuTeGEMMDBase_MXFP8(FusedOperation): + """Base fused backward op for MXFP8 GroupedLinear + activation + GroupedLinear. Uses experimental CuTe DSL kernel from cuDNN front-end. """ @classmethod - @functools.lru_cache(maxsize=None) - def grouped_gemm_dglu_kernel(cls) -> Callable: - """Fused kernel for grouped GEMM, GLU activation backward, and scale grad.""" - from cudnn import grouped_gemm_dglu_wrapper_sm100 # pylint: disable=no-name-in-module - - return grouped_gemm_dglu_wrapper_sm100 + def grouped_gemm_dactivation_kernel(cls) -> Callable: + """Fused kernel for grouped GEMM, activation backward, and scale grad.""" + raise NotImplementedError @classmethod @functools.lru_cache(maxsize=None) @@ -296,7 +295,7 @@ def is_supported(cls) -> bool: if not _cudnn_frontend_version_supported(): return False try: - cls.grouped_gemm_dglu_kernel() + cls.grouped_gemm_dactivation_kernel() cls.grouped_gemm_quant_kernel() except ImportError: return False @@ -306,19 +305,26 @@ def __init__( self, *, fc1: GroupedLinear, - swiglu: ScaledSwiGLU | ScaledClampedQGeGLU, + activation: Optional[FusibleOperation], fc2: GroupedLinear, ) -> None: - super().__init__((fc1, swiglu, fc2)) + if activation is None: + raise TypeError("Expected a grouped MLP activation op.") + super().__init__((fc1, activation, fc2)) if not self.is_supported(): - self.grouped_gemm_dglu_kernel() # Try triggering import error + self.grouped_gemm_dactivation_kernel() # Try triggering import error raise RuntimeError(f"{self.__class__.__name__} is not supported on this system.") - validate_grouped_mlp_dims(fc1, swiglu, fc2) - # The cuDNN dgeglu implementation corresponds to ScaledClampedQGeGLU. - # The act_func string should be fixed on the cuDNN FE side. - self._cudnn_dact_func: str = ( - "dgeglu" if isinstance(swiglu, ScaledClampedQGeGLU) else "dswiglu" - ) + validate_grouped_mlp_dims(fc1, activation, fc2) + if not is_glu_activation(activation): + # grouped_gemm_dsrelu_wrapper_sm100 is dSReLU-specific and does not + # take the GLU ``act_func`` selector. + self._cudnn_dact_func: Optional[str] = None + else: + # The cuDNN dgeglu implementation corresponds to ScaledClampedQGeGLU. + # The act_func string should be fixed on the cuDNN FE side. + self._cudnn_dact_func = ( + "dgeglu" if isinstance(activation, ScaledClampedQGeGLU) else "dswiglu" + ) def fuser_backward( self, @@ -333,7 +339,7 @@ def fuser_backward( # Get basic operations fc1_op, _, fc2_op = self.basic_ops - fc1_ctx, swiglu_ctx, fc2_ctx = basic_op_ctxs + fc1_ctx, activation_ctx, fc2_ctx = basic_op_ctxs # Tensor properties fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features) @@ -358,8 +364,11 @@ def fuser_backward( saved_tensors[num_groups:], ) - # Saved tensors from scaled SwiGLU forward - swiglu_in, scales = swiglu_ctx.saved_tensors + # Saved tensors from activation forward + activation_in, scales = activation_ctx.saved_tensors + recompute_fc2_x_from_dsrelu = bool( + getattr(fc2_ctx, "recompute_input_from_dsrelu", False) + ) and bool(fc2_ctx.weight_requires_grad) # Saved tensors from FC2 forward. # Layout: [split_sizes, base_split_offsets, split_points, @@ -446,20 +455,19 @@ def fuser_backward( # Kernel scaling factors alpha_tensor = get_cached_ones_tensor(num_groups, dtype, device) - norm_const_tensor = get_cached_ones_tensor(1, dtype, device) + norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device) current_stream = torch.cuda.current_stream().cuda_stream scales_f32 = scales.detach().to(dtype=torch.float32) scales_tensor = scales_f32.reshape(-1, 1, 1) dscales_tensor = torch.zeros_like(scales_tensor) - fc2_dglu_kwargs = { + fc2_dactivation_kwargs = { "a_tensor": fc2_dy_data, - "c_tensor": swiglu_in.unsqueeze(0).permute(1, 2, 0), + "c_tensor": activation_in.unsqueeze(0).permute(1, 2, 0), "sfa_tensor": fc2_dy_scales, "padded_offsets": split_points, "alpha_tensor": alpha_tensor, - "beta_tensor": alpha_tensor, "prob_tensor": scales_tensor, "dprob_tensor": dscales_tensor, "generate_dbias": fc1_op.has_bias, @@ -469,9 +477,13 @@ def fuser_backward( "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, "current_stream": current_stream, "discrete_col_sfd": True, - "act_func": self._cudnn_dact_func, "use_dynamic_sched": True, } + if self._cudnn_dact_func is not None: + fc2_dactivation_kwargs["beta_tensor"] = alpha_tensor + fc2_dactivation_kwargs["act_func"] = self._cudnn_dact_func + else: + fc2_dactivation_kwargs["use_dsrelu_reuse"] = recompute_fc2_x_from_dsrelu if fc2_op.single_grouped_weight: # Clone and swizzle scales for GEMM @@ -495,8 +507,8 @@ def fuser_backward( ) fc2_w_scales = fc2_w_scales.permute(3, 4, 1, 5, 2, 0) - fc2_dglu_kwargs["b_tensor"] = fc2_w_data - fc2_dglu_kwargs["sfb_tensor"] = fc2_w_scales + fc2_dactivation_kwargs["b_tensor"] = fc2_w_data + fc2_dactivation_kwargs["sfb_tensor"] = fc2_w_scales else: fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sw = tex.get_device_pointer_for_data_and_scales( [w._columnwise_data for w in grouped_fc2_weight], @@ -505,13 +517,13 @@ def fuser_backward( rowwise=False, data_dtype=grouped_fc2_weight[0]._fp8_dtype, ) - fc2_dglu_kwargs["b_ptrs"] = fc2_b_ptrs - fc2_dglu_kwargs["sfb_ptrs"] = fc2_sfb_ptrs - fc2_dglu_kwargs["n"] = fc2_weight_shape[1] - fc2_dglu_kwargs["b_dtype"] = torch.float8_e4m3fn - fc2_dglu_kwargs["b_major"] = "n" + fc2_dactivation_kwargs["b_ptrs"] = fc2_b_ptrs + fc2_dactivation_kwargs["sfb_ptrs"] = fc2_sfb_ptrs + fc2_dactivation_kwargs["n"] = fc2_weight_shape[1] + fc2_dactivation_kwargs["b_dtype"] = torch.float8_e4m3fn + fc2_dactivation_kwargs["b_major"] = "n" - fc2_dgrad_kernel_out = self.grouped_gemm_dglu_kernel()(**fc2_dglu_kwargs) + fc2_dgrad_kernel_out = self.grouped_gemm_dactivation_kernel()(**fc2_dactivation_kwargs) fc1_dy_row_data = fc2_dgrad_kernel_out["d_row_tensor"] fc1_dy_row_data = fc1_dy_row_data.view(out_shape[0], fc1_weight_shape[0]) @@ -523,6 +535,37 @@ def fuser_backward( fc1_dy_col_scale = fc2_dgrad_kernel_out["sfd_col_tensor"].permute(5, 2, 4, 0, 1, 3).view(-1) grad_scales = fc2_dgrad_kernel_out["dprob_tensor"].view(-1) + if recompute_fc2_x_from_dsrelu: + d_srelu_tensor = fc2_dgrad_kernel_out.get("d_srelu_tensor") + if d_srelu_tensor is None: + raise RuntimeError( + "SReLU recompute is enabled, but the DSReLU kernel did not return " + "the recomputed FC2 input tensor." + ) + + sfd_col_d_srelu_tensor = fc2_dgrad_kernel_out.get("sfd_col_d_srelu_tensor") + if sfd_col_d_srelu_tensor is None: + raise RuntimeError( + "SReLU recompute is enabled, but the DSReLU kernel did not return " + "the recomputed FC2 input column scale tensor." + ) + + fc2_x_col_data = d_srelu_tensor.view(out_shape[0], fc2_weight_shape[1]) + fc2_x_col_scale = sfd_col_d_srelu_tensor.permute(5, 2, 4, 0, 1, 3) + grouped_fc2_x = GroupedTensor( + shape=(out_shape[0], fc2_weight_shape[1]), + dtype=dtype, + num_tensors=num_groups, + quantizer=fc2_ctx.input_quantizers[0], + data=None, + columnwise_data=fc2_x_col_data.reshape(-1), + scale_inv=None, + columnwise_scale_inv=fc2_x_col_scale.reshape(-1), + first_dims=split_sizes, + tensor_offsets=base_split_offsets * fc2_weight_shape[1], + with_gemm_swizzled_scales=True, + ) + fc2_bias_grads: Optional[list[Optional[torch.Tensor]]] = None fc2_bias_grad_packed: Optional[torch.Tensor] = None if scale_bias: @@ -547,7 +590,8 @@ def fuser_backward( else: fc2_bias_grads = [fc2_dbias_packed[idx] for idx in range(num_groups)] - grad_scales = grad_scales.to(dtype=dtype) + if grad_scales is not None: + grad_scales = grad_scales.to(dtype=dtype) fc1_bias_grads: Optional[list[Optional[torch.Tensor]]] = None fc1_bias_grad_packed: Optional[torch.Tensor] = None @@ -618,7 +662,7 @@ def fuser_backward( "a_tensor": fc1_dgrad_a_data, "sfa_tensor": fc1_dgrad_a_scales, "padded_offsets": split_points, - "alpha_tensor": alpha_tensor.float(), + "alpha_tensor": alpha_tensor, "norm_const_tensor": None, "prob_tensor": torch.ones((out_shape[0], 1, 1), dtype=torch.float32, device=device), "acc_dtype": torch.float32, @@ -703,13 +747,44 @@ def fuser_backward( ) fc2_grad_extra = (None, None) if fc2_op._scale_bias else (None,) + activation_grad_extra = (grad_scales,) if grad_scales is not None else () return ( grad_input, [fc1_grad_params, (), fc2_grad_params], - [(None,), (grad_scales,), fc2_grad_extra], + [(None,), activation_grad_extra, fc2_grad_extra], ) +class BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8(_BackwardGroupedMLP_CuTeGEMMDBase_MXFP8): + """Fused backward op for GroupedLinear + scaled GLU + GroupedLinear.""" + + @classmethod + @functools.lru_cache(maxsize=None) + def grouped_gemm_dactivation_kernel(cls) -> Callable: + """Fused kernel for grouped GEMM, GLU activation backward, and scale grad.""" + from cudnn import grouped_gemm_dglu_wrapper_sm100 # pylint: disable=no-name-in-module + + return grouped_gemm_dglu_wrapper_sm100 + + +class BackwardGroupedMLP_CuTeGEMMDUnary_MXFP8(_BackwardGroupedMLP_CuTeGEMMDBase_MXFP8): + """Fused backward op for GroupedLinear + scaled unary activation + GroupedLinear.""" + + @classmethod + @functools.lru_cache(maxsize=None) + def is_supported(cls) -> bool: + """Whether the SReLU fused backward operation is supported on the current system.""" + return _cudnn_frontend_supports_grouped_gemm_srelu() and super().is_supported() + + @classmethod + @functools.lru_cache(maxsize=None) + def grouped_gemm_dactivation_kernel(cls) -> Callable: + """Fused kernel for grouped GEMM and dSReLU activation backward.""" + from cudnn import grouped_gemm_dsrelu_wrapper_sm100 # pylint: disable=no-name-in-module + + return grouped_gemm_dsrelu_wrapper_sm100 + + def fuse_backward_ops( ops: list[FusibleOperation], *, @@ -735,10 +810,28 @@ def fuse_backward_ops( return fuse_grouped_mlp_ops( ops, recipe=recipe, - fused_op_cls=BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8, + fused_op_cls=BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8, + ) + + +def fuse_backward_srelu_ops( + ops: list[FusibleOperation], + *, + recipe: Optional[Recipe] = None, + **unused, # pylint: disable=unused-argument +) -> list[FusibleOperation]: + """Apply GroupedLinear + ScaledSReLU + GroupedLinear fusion for backward pass.""" + + return fuse_grouped_mlp_ops( + ops, + recipe=recipe, + fused_op_cls=BackwardGroupedMLP_CuTeGEMMDUnary_MXFP8, + activation_op_types=(ScaledSReLU,), ) # Register fusion if available -if BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8.is_supported(): +if BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8.is_supported(): register_backward_fusion(fuse_backward_ops, prepend=True) +if BackwardGroupedMLP_CuTeGEMMDUnary_MXFP8.is_supported(): + register_backward_fusion(fuse_backward_srelu_ops, prepend=True) diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index 91db2ff9b..f7ac45502 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -19,12 +19,15 @@ from ...tensor.grouped_tensor import GroupedTensor from ...tensor.mxfp8_tensor import MXFP8Quantizer from ...constants import MXFP8_BLOCK_SCALING_SIZE -from ..basic import GroupedLinear, ScaledClampedQGeGLU, ScaledSwiGLU +from ..basic import GroupedLinear, ScaledSReLU, ScaledClampedQGeGLU from ..fuser import register_forward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext from .._common import ( _cudnn_frontend_version_supported, + _cudnn_frontend_supports_grouped_gemm_srelu, + _nvidia_cudnn_frontend_supports_wgrad, fuse_grouped_mlp_ops, + is_glu_activation, is_quantized_tensor, maybe_dequantize, validate_grouped_mlp_dims, @@ -45,20 +48,35 @@ def _pack_grouped_linear_bias_for_cudnn(linear_op: GroupedLinear) -> Optional[to return torch.stack(rows, dim=0).transpose(0, 1) -class ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8(FusedOperation): - """Fused op for MXFP8 GroupedLinear + scaled GLU + GroupedLinear +@functools.lru_cache(maxsize=1) +def _grouped_gemm_dsrelu_backward_supported() -> bool: + """Whether the cuDNN FE grouped GEMM dSReLU backward wrapper is available.""" + if int(os.environ.get("NVTE_CUTEDSL_FUSED_GROUPED_MLP", "0")) <= 0: + return False + if get_device_compute_capability()[0] != 10: + return False + if not _cudnn_frontend_supports_grouped_gemm_srelu(): + return False + try: + from cudnn import ( + grouped_gemm_dsrelu_wrapper_sm100, + ) # pylint: disable=import-outside-toplevel + except ImportError: + return False + return grouped_gemm_dsrelu_wrapper_sm100 is not None + + +class _ForwardGroupedMLP_CuTeGEMMBase_MXFP8(FusedOperation): + """Base fused op for MXFP8 GroupedLinear + activation + GroupedLinear. Uses experimental CuTe DSL kernel from cuDNN front-end. """ @classmethod - @functools.lru_cache(maxsize=None) - def grouped_gemm_glu_kernel(cls) -> Callable: - """Fused kernel for grouped GEMM, GLU activation, and post-multiplication.""" - from cudnn import grouped_gemm_glu_wrapper_sm100 # pylint: disable=no-name-in-module - - return grouped_gemm_glu_wrapper_sm100 + def grouped_gemm_activation_kernel(cls) -> Callable: + """Fused kernel for grouped GEMM, activation, and post-multiplication.""" + raise NotImplementedError @classmethod @functools.lru_cache(maxsize=None) @@ -79,7 +97,7 @@ def is_supported(cls) -> bool: if not _cudnn_frontend_version_supported(): return False try: - cls.grouped_gemm_glu_kernel() + cls.grouped_gemm_activation_kernel() cls.grouped_gemm_quant_kernel() except ImportError: return False @@ -89,17 +107,26 @@ def __init__( self, *, fc1: GroupedLinear, - swiglu: ScaledSwiGLU | ScaledClampedQGeGLU, + activation: Optional[FusibleOperation], fc2: GroupedLinear, ) -> None: - super().__init__((fc1, swiglu, fc2)) + if activation is None: + raise TypeError("Expected a grouped MLP activation op.") + super().__init__((fc1, activation, fc2)) if not self.is_supported(): - self.grouped_gemm_glu_kernel() # Try triggering import error + self.grouped_gemm_activation_kernel() # Try triggering import error raise RuntimeError(f"{self.__class__.__name__} is not supported on this system.") - validate_grouped_mlp_dims(fc1, swiglu, fc2) - # The cuDNN geglu implementation corresponds to ScaledClampedQGeGLU. - # The act_func string should be fixed on the cuDNN FE side. - self._cudnn_act_func: str = "geglu" if isinstance(swiglu, ScaledClampedQGeGLU) else "swiglu" + validate_grouped_mlp_dims(fc1, activation, fc2) + if not is_glu_activation(activation): + # grouped_gemm_srelu_wrapper_sm100 is SReLU-specific and does not + # take the GLU ``act_func`` selector. + self._cudnn_act_func: Optional[str] = None + else: + # The cuDNN geglu implementation corresponds to ScaledClampedQGeGLU. + # The act_func string should be fixed on the cuDNN FE side. + self._cudnn_act_func = ( + "geglu" if isinstance(activation, ScaledClampedQGeGLU) else "swiglu" + ) def fuser_forward( self, @@ -113,7 +140,7 @@ def fuser_forward( ) -> tuple[torch.Tensor, Iterable[Iterable[torch.Tensor]]]: # Get basic operations fc1_op, _, fc2_op = self.basic_ops - fc1_ctx, swiglu_ctx, fc2_ctx = basic_op_ctxs + fc1_ctx, activation_ctx, fc2_ctx = basic_op_ctxs # Tensor properties fc1_weight_shape = (fc1_op.out_features, fc1_op.in_features) @@ -164,7 +191,7 @@ def fuser_forward( split_points = base_split_offsets[1:].to(dtype=torch.int) fc2_x_tensor_offsets = base_split_offsets * fc2_weight_shape[1] - # Extract post-scales from extra input + # Extract per-row activation probabilities from the middle op. scales = basic_op_extra_inputs[1][0] # Prepare FC1 grouped weight tensor for fused kernels. @@ -281,13 +308,13 @@ def fuser_forward( fc1_x_scales = fc1_x_scales.permute(3, 4, 1, 5, 2, 0) alpha_tensor = get_cached_ones_tensor(num_groups, dtype, device) - norm_const_tensor = get_cached_ones_tensor(1, dtype, device) + norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device) current_stream = torch.cuda.current_stream().cuda_stream fc1_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc1_op) fc2_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc2_op) - fc1_glu_kwargs = { + fc1_activation_kwargs = { "a_tensor": fc1_x_data, "sfa_tensor": fc1_x_scales, "padded_offsets": split_points, @@ -302,9 +329,10 @@ def fuser_forward( "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, "current_stream": current_stream, "discrete_col_sfd": True, - "act_func": self._cudnn_act_func, "use_dynamic_sched": True, } + if self._cudnn_act_func is not None: + fc1_activation_kwargs["act_func"] = self._cudnn_act_func if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM. @@ -329,8 +357,8 @@ def fuser_forward( ) fc1_w_scales = fc1_w_scales.permute(3, 4, 1, 5, 2, 0) - fc1_glu_kwargs["b_tensor"] = fc1_w_data - fc1_glu_kwargs["sfb_tensor"] = fc1_w_scales + fc1_activation_kwargs["b_tensor"] = fc1_w_data + fc1_activation_kwargs["sfb_tensor"] = fc1_w_scales else: # Discrete-weight kernel: per-expert data/scale pointers fc1_b_ptrs, fc1_sfb_ptrs, _fc1_sw = tex.get_device_pointer_for_data_and_scales( @@ -340,13 +368,13 @@ def fuser_forward( rowwise=True, data_dtype=grouped_fc1_weight[0]._fp8_dtype, ) - fc1_glu_kwargs["b_ptrs"] = fc1_b_ptrs - fc1_glu_kwargs["sfb_ptrs"] = fc1_sfb_ptrs - fc1_glu_kwargs["n"] = fc1_weight_shape[0] - fc1_glu_kwargs["b_dtype"] = torch.float8_e4m3fn - fc1_glu_kwargs["b_major"] = "k" + fc1_activation_kwargs["b_ptrs"] = fc1_b_ptrs + fc1_activation_kwargs["sfb_ptrs"] = fc1_sfb_ptrs + fc1_activation_kwargs["n"] = fc1_weight_shape[0] + fc1_activation_kwargs["b_dtype"] = torch.float8_e4m3fn + fc1_activation_kwargs["b_major"] = "k" - fc1_kernel_out = self.grouped_gemm_glu_kernel()(**fc1_glu_kwargs) + fc1_kernel_out = self.grouped_gemm_activation_kernel()(**fc1_activation_kwargs) # Unpack kernel outputs # Note: Fused kernel outputs tensors with non-contiguous @@ -357,8 +385,8 @@ def fuser_forward( # Column-wise data logical shape: (sum(m_splits), k, 1) # Column-wise scale logical shape: (32 (block col), 4 (block col), # k/128, 4 (block row), sum(m_splits)/128, 1) - swiglu_in = fc1_kernel_out["c_tensor"] - swiglu_in = swiglu_in.view(in_shape[0], fc1_weight_shape[0]) + activation_in = fc1_kernel_out["c_tensor"] + activation_in = activation_in.view(in_shape[0], fc1_weight_shape[0]) fc2_in_row_data = fc1_kernel_out["d_tensor"] fc2_in_row_data = fc2_in_row_data.view(in_shape[0], fc2_weight_shape[1]) fc2_in_row_scale = fc1_kernel_out["sfd_row_tensor"] @@ -397,7 +425,7 @@ def fuser_forward( "a_tensor": fc1_kernel_out["d_tensor"], "sfa_tensor": fc1_kernel_out["sfd_row_tensor"], "padded_offsets": split_points, - "alpha_tensor": alpha_tensor.float(), + "alpha_tensor": alpha_tensor, "bias_tensor": fc2_bias_packed, "norm_const_tensor": None, "prob_tensor": fc2_scales_tensor, @@ -450,10 +478,23 @@ def fuser_forward( # Save state for backward pass if requires_grad: - mark_grouped_tensor(grouped_fc1_x, swiglu_in, scales, grouped_fc2_x) + mark_grouped_tensor(grouped_fc1_x, activation_in, scales, grouped_fc2_x) + activation_op = self.basic_ops[1] + activation_is_srelu = isinstance(activation_op, ScaledSReLU) + activation_recompute_in_mlp = bool( + getattr(activation_op, "activation_recompute_in_mlp", False) + ) + recompute_srelu_fc2_x = ( + activation_is_srelu + and activation_recompute_in_mlp + and weight_requires_grad + and _grouped_gemm_dsrelu_backward_supported() + and _nvidia_cudnn_frontend_supports_wgrad() + ) + saved_grouped_fc2_x = None if recompute_srelu_fc2_x else grouped_fc2_x # Save the input ``GroupedTensor``s themselves for the activations. - for grouped_fc_x in (grouped_fc1_x, grouped_fc2_x): + for grouped_fc_x in (grouped_fc1_x, saved_grouped_fc2_x): if grouped_fc_x is not None: grouped_fc_x.rowwise_data = None grouped_fc_x.scale_inv = None @@ -481,11 +522,11 @@ def fuser_forward( fc1_ctx.input_requires_grad = input_requires_grad fc1_ctx.weight_requires_grad = weight_requires_grad - # Scaled SwiGLU - swiglu_ctx.save_for_backward(swiglu_in, scales) - swiglu_ctx.input_requires_grad = True - swiglu_ctx.extra_input_requires_grad = True - swiglu_ctx.dtype = dtype + # Activation + activation_ctx.save_for_backward(activation_in, scales) + activation_ctx.extra_input_requires_grad = True + activation_ctx.input_requires_grad = True + activation_ctx.dtype = dtype # FC2 saved-tensor layout. Matches the unfused # ``GroupedLinear._fuser_forward_grouped_tensor`` layout so the @@ -504,7 +545,7 @@ def fuser_forward( ] if fc2_op._scale_bias: fc2_saved.append(fc2_scales) - fc2_saved.append(grouped_fc2_x) + fc2_saved.append(saved_grouped_fc2_x) fc2_saved.extend(fc2_weight_tensors) fc2_ctx.save_for_backward(*fc2_saved) fc2_ctx.use_grouped_tensor_path = True @@ -516,10 +557,41 @@ def fuser_forward( fc2_ctx.dtype = dtype fc2_ctx.input_requires_grad = input_requires_grad fc2_ctx.weight_requires_grad = weight_requires_grad + fc2_ctx.recompute_input_from_dsrelu = recompute_srelu_fc2_x return fc2_out, [(), (), ()] +class ForwardGroupedMLP_CuTeGEMMGLU_MXFP8(_ForwardGroupedMLP_CuTeGEMMBase_MXFP8): + """Fused op for MXFP8 GroupedLinear + scaled GLU + GroupedLinear.""" + + @classmethod + @functools.lru_cache(maxsize=None) + def grouped_gemm_activation_kernel(cls) -> Callable: + """Fused kernel for grouped GEMM, GLU activation, and post-multiplication.""" + from cudnn import grouped_gemm_glu_wrapper_sm100 # pylint: disable=no-name-in-module + + return grouped_gemm_glu_wrapper_sm100 + + +class ForwardGroupedMLP_CuTeGEMMUnary_MXFP8(_ForwardGroupedMLP_CuTeGEMMBase_MXFP8): + """Fused op for MXFP8 GroupedLinear + scaled unary activation + GroupedLinear.""" + + @classmethod + @functools.lru_cache(maxsize=None) + def is_supported(cls) -> bool: + """Whether the SReLU fused operation is supported on the current system.""" + return _cudnn_frontend_supports_grouped_gemm_srelu() and super().is_supported() + + @classmethod + @functools.lru_cache(maxsize=None) + def grouped_gemm_activation_kernel(cls) -> Callable: + """Fused kernel for grouped GEMM, SReLU activation, and post-multiplication.""" + from cudnn import grouped_gemm_srelu_wrapper_sm100 # pylint: disable=no-name-in-module + + return grouped_gemm_srelu_wrapper_sm100 + + def fuse_forward_ops( ops: list[FusibleOperation], *, @@ -545,10 +617,28 @@ def fuse_forward_ops( return fuse_grouped_mlp_ops( ops, recipe=recipe, - fused_op_cls=ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8, + fused_op_cls=ForwardGroupedMLP_CuTeGEMMGLU_MXFP8, + ) + + +def fuse_forward_srelu_ops( + ops: list[FusibleOperation], + *, + recipe: Optional[Recipe] = None, + **unused, # pylint: disable=unused-argument +) -> list[FusibleOperation]: + """Apply GroupedLinear + ScaledSReLU + GroupedLinear fusion for forward pass.""" + + return fuse_grouped_mlp_ops( + ops, + recipe=recipe, + fused_op_cls=ForwardGroupedMLP_CuTeGEMMUnary_MXFP8, + activation_op_types=(ScaledSReLU,), ) # Register fusion if available -if ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported(): +if ForwardGroupedMLP_CuTeGEMMGLU_MXFP8.is_supported(): register_forward_fusion(fuse_forward_ops, prepend=True) +if ForwardGroupedMLP_CuTeGEMMUnary_MXFP8.is_supported(): + register_forward_fusion(fuse_forward_srelu_ops, prepend=True) From 86ade9ea8efc94df850017eecbd91b7235c67845 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Thu, 21 May 2026 16:56:54 -0700 Subject: [PATCH 076/180] CP Tests batching using subprocess worker pool (#2993) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Batch CP attention tests via a persistent NCCL pool The existing test path spawns one torchrun per parametrized case, paying NCCL init + CUDA context + Python startup on every call. With ~hundreds of cases the launch overhead dominates wall time and was a primary driver of the L3 timeout that prior batching PRs worked around. This change replaces the per-case subprocess with one long-lived torchrun per (world_size). NCCL is initialized once at session start and reused across cases. Pytest sends one JSON request per case over rank-0 stdin; the worker dispatches to run_dpa_with_cp(**kwargs), gathers (ok, error) from every rank, and writes one JSON response on rank-0 stdout. run_attention_with_cp.py is left almost untouched; a new NVTE_CP_POOL_PG=1 env var gates the dist.init_process_group() and dist.destroy_process_group() calls so the function reuses the pool's main PG instead of creating its own. The per-case cp_comm_group (and a2a+p2p sub-groups) are explicitly destroyed at function exit to prevent communicator leakage across cases. The PoolWorker class adds two pieces of error recovery that the prior subprocess-per-case design got for free: a select-based per-call timeout (default 600s, NVTE_CP_POOL_TIMEOUT_SEC) and auto-respawn on worker death or timeout. A test-level exception is reported as an AssertionError and the pool keeps running for the next case. Two pool sizes are needed because cp_comm_type='a2a+p2p' requires world_size=4 and the others use world_size=2; you can't resize an active PG. Pools are spawned lazily so a 2-GPU-only run never pays the 4-GPU init. Signed-off-by: Sudhakar Singh * Reset FP8 state and barrier between pool cases Two resilience fixes carried over from the existing batching PR (sudhakars/cp_test_batching_pr) without which the pool will cascade-fail FP8 tests and silently propagate NCCL desync. 1. FP8GlobalStateManager.reset() between cases. FP8 quantizer state (recipe handles, autocast counters) lives in module-level globals. Reusing one Python process across cases otherwise carries that state forward. The prior batching PR landed an explicit fix for the same issue ("Fix FP8 cascade failures") after observing real test failures from this. 2. dist.barrier() after each case. If one rank's case errored before its last collective, the others can be stuck waiting on a comm that will never complete. The barrier here surfaces that immediately as a timeout in this case rather than letting the corruption leak into the next case's collectives. Also pops the transient NVTE_* env vars run_dpa_with_cp sets at the top of each call. run_dpa_with_cp already sets them unconditionally so this is defensive, but cheap insurance against future variants that might not. Signed-off-by: Sudhakar Singh * Deep-copy ModelConfig in run_dpa_with_cp The model_configs_{flash,fused}_attn dicts are module-level and shared across pool cases. The THD branch below rewrites config.attn_mask_type in place (causal -> padding_causal, no_mask -> padding). With the persistent-pool runner, the next case looking up the same model key gets the mutated config and fails the "causal or no_mask only" assert. Caught at benchmark time on cp_2_0 + thd, identical to the cascade the existing batching PR (sudhakars/cp_test_batching_pr) hit and fixed the same way in commit 6355f620. Signed-off-by: Sudhakar Singh * Skip deterministic configs incompatible with FusedAttention Mirrors the two pre-emptive skips on the PR-batching branch: * non-vanilla softmax with FusedAttention is not deterministic * post_scale_bias with requires_grad is not deterministic Without these skips, the corresponding configs propagate into the pool worker under NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 and fail inside run_dpa_with_cp instead of being marked SKIPPED. Signed-off-by: Sudhakar Singh * Reseed RNG between pool cases; reset before, not after The pool worker reused RNG state across cases, which produced small numerical drift on some non-FP8 fused-attention configs (cp_1_0 + thd/p2p, cp_1_0 + sbhd/all_gather) compared to the single-shot worker. Matches the per-case startup of the single-shot worker: torch.manual_seed(1234) + torch.cuda.manual_seed(1234) at the start of every case, alongside the existing FP8 / env / cache resets. Moved the reset call from the post-case finally block to the start of _run_one so the first case is also seeded consistently with subsequent cases. Otherwise the first case would inherit the process-default RNG and only the second-and-later cases would be deterministic. Validated locally: 38 passed, 0 failed (was 36 passed, 2 failed). Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Robustify pool: capture worker stderr, tighten timeout, add timing knob Three changes that bring the pool's failure semantics on par with the per-batch torchrun approach in PR #2965 and remove a couple of footguns: 1. Capture pool-worker stderr into a ring buffer and attach the tail to crash-path AssertionErrors. Equivalent in spirit to PR #2965's run_distributed() — CI JUnit XML now shows the actual cause (NCCL error, Python traceback, OOM) inline with the failing test, instead of just "pool worker died mid-request" / "timed out". A daemon drainer thread reads stderr line-by-line into a deque(maxlen=200) and also echoes to sys.stderr so pytest's per-test capture still gets every line. Maximum buffered footprint ~40 KB. 2. Tighten POOL_SUBMIT_TIMEOUT_SEC default 600 -> 90. On H100 the slowest observed per-case wall is ~15 s (p99 also 15 s, p50 ~5 s). 90 s gives ~6x headroom over the worst observed case while still detecting a genuine hang within ~1.5 min instead of ~10 min. Env var still overrides for slower machines or expanded test matrices. 3. Optional per-case wall-time logging (NVTE_CP_POOL_TIMING=1) prints "[POOL-TIMING] case_idx=N world_size=W wall_s=X.XXX ok=B" to stderr on rank 0 only. Grep-friendly; lets future tuning recalibrate the timeout against the observed distribution. Off by default so normal runs stay quiet. Validated: 38 passed / 0 failed in 248 s on H100, test_essential=True, with no perf regression vs the un-patched 256 s. Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address PR review: NCCL leak, stdout protocol, Windows note Three fixes responding to https://github.com/NVIDIA/TransformerEngine/pull/2993 review comments: P1: NCCL communicator leak on exception (run_attention_with_cp.py) run_dpa_with_cp() created cp_comm_group (and optionally cp_comm_sub_groups) near the top, but the destroy_process_group() calls ran only on the success path at the end of the function. Any exception in between (tensor assertion, OOM, NCCL error) skipped the cleanup, leaking communicators in pool mode. Long sessions with repeated failures could exhaust NCCL internal tracking. Wrap the test work in try/finally so the destroy logic always runs. Initialise cp_comm_sub_groups = [] unconditionally so the finally block is safe even when cp_comm_type != "a2a+p2p" (or when an assert fires before the populate loop). Each destroy is itself try/except so a destroy failure on one group doesn't leak the others. P2: stdout protocol can be corrupted by interleaved chatter torchrun and ranks 1..N share rank 0's stdout fd. Any non-rank-0 print, NCCL debug line, or torchrun status output interleaves with the JSON response and breaks json.loads, killing the pool with a misleading "json decode error". Prefix every response with "[CP_POOL_RESP] " in run_attention_with_cp_pool.py and have PoolWorker.submit() scan stdout for sentinel-prefixed lines, echoing non-protocol lines to stderr for visibility. Bounded scan (MAX_NOISE_LINES=1000) so a chatty worker can't stall the parent. P2 (doc): select.select on a pipe fd is Linux/macOS only Added a short comment noting Windows portability. CP attention tests run on Linux GPU hosts; this is a documentation issue, not a real bug. Validated: 38 passed / 0 failed in 270 s on H100, test_essential=True (was 248 s pre-P2 — the +22 s is the new sentinel-scan loop's per-line overhead at ~600 ms/case, within noise). Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Fix stream race on max_logit_per_step in all-gather CP forward In AttnFuncWithCPAndKVAllGather.forward, max_logit_per_step[i] is written inside `with torch.cuda.stream(flash_attn_streams[i])`. For i=1, flash_attn_streams[1] is cp_stream — i.e. *not* the default stream. Later, at loop iteration i=2, the code reads max_logit_per_step[1] via `torch.maximum(max_logit, max_logit_per_step[i-1])` which runs on the default stream. Without an explicit wait_stream, this is a read-after-write race across streams. The post-loop `current_stream().wait_stream(cp_stream)` is too late — the race has already fired. The race is latent: outcome depends on stream scheduling. In a fresh-process subprocess (one-torchrun-per-test path), streams are cleanly initialised and timing happens to put the write before the read. In a long-running persistent-worker process — exposed by PR #2993's pool design — prior workloads shape stream state differently, the read can fire before the write completes, and max_logit ends up with stale values in some heads (~0.3 abs diff, 3/12 elements wrong on the H100 matrix). Fix: insert `current_stream().wait_stream(flash_attn_streams[i-1])` before the torch.maximum read. No-op when the streams are identical (i=1 case, where flash_attn_streams[0] is current_stream), only fires when reading from cp_stream (i=2 case). Validated: 8xH100, test_essential=False, 348 passed / 0 failed in 27m 10s (was 323 passed + 5 failed at this commit's parent, all 5 failing on cp_comm_type=all_gather with mismatched max_logit). The failing configs (all_gather + cp_1_0/cp_1_1 + bshd or fp16) now pass under the pool — confirming the race was the sole root cause. Signed-off-by: Sudhakar Singh * Address PR review (R2): drop dead code in pool worker and PoolWorker Line-level cleanups from the second reviewer pass on PR #2993. Each item is dead/redundant; none changes behaviour. Full-matrix test_essential=False on 8xH100 still passes 348/0 in 26m 23s after these. run_attention_with_cp_pool.py: - Drop _TRANSIENT_ENV_KEYS tuple + pop loop. run_dpa_with_cp already re-sets NVTE_FUSED_ATTN/NVTE_FLASH_ATTN unconditionally at the top and pops the FP8 ones itself. The pop loop was defensive against a hypothetical "future caller that doesn't re-set them" that doesn't exist. - Drop gc.collect() after torch.cuda.empty_cache(). The cases create no Python reference cycles between iterations and empty_cache only frees CUDA blocks PyTorch already considers free; the combination was no-op here. - Drop dist.barrier() after dist.gather_object(). gather_object is itself a collective synchronization point — if every rank reaches it, none is ahead. The "surface a wedged communicator here" comment was wishful: a wedged communicator would already wedge the gather. test_attention_with_cp.py (PoolWorker): - Drop _MAX_NOISE_LINES = 1000 + the scanned counter + the unreachable post-loop "1000+ lines" branch. select()'s deadline already bounds the loop; the line-count cap was redundant and the over-limit branch was unreachable in practice. - Inline _stderr_tail() into _diag(). Single caller, single use. - Drop the _stderr_thread attribute. The drainer is daemon and self-terminates when the pipe closes; we never read the field anywhere, so initialising and nulling it was bookkeeping for no reason. - Drop the dead assert in submit() — _ensure_alive() on the prior line already guarantees proc/stdin/stdout exist. Deferred to a follow-up: - L8 (drop try/except around dist.destroy_process_group). Real semantic change: hides errors that occur when a previous test wedged the communicator. Worth doing but needs its own validation. - R1 medium items M1 (module-level flag vs NVTE_CP_POOL_PG env var), M2 (redirect rank>0 stdout vs sentinel scan), M3 (explicit CUDA_VISIBLE_DEVICES per pool). Same reasoning — separate PRs. Signed-off-by: Sudhakar Singh * Address PR review (items 2+3): reuse CP groups across pool cases world_size and the rank set don't change for the lifetime of one pool, so recreating the world group and a2a+p2p sub-groups per case wastes ~50-100 ms of NCCL setup each. Pre-create them once in the pool worker (new helper _create_cp_comm_groups), stash on the run_attention_with_cp module via module-level _pool_cp_comm_group / _pool_cp_comm_sub_groups pointers, and reuse them from run_dpa_with_cp in pool mode. Pool teardown destroys them once at shutdown. Also move per-case dist.new_group() calls inside the try/finally in run_dpa_with_cp: a failure mid-loop in the a2a+p2p sub_group population otherwise leaks every communicator created before the failure. The finally now only destroys groups we created locally (cp_comm_group / sub_groups populated in the else-branch), leaving pool-owned groups alone for reuse. cyanguwa's review feedback on PR #2993. Signed-off-by: Sudhakar Singh * Flatten try/finally wrap in run_dpa_with_cp The Round-1 P1 NCCL-communicator-leak fix (e162a9ec) wrapped the ~540-line body of run_dpa_with_cp in try/finally. The wrap itself was tiny but it re-indented every line of the body by one level, inflating the PR diff of run_attention_with_cp.py to ~1000 lines against origin/main. Items 2+3 (d15bfce3) since made the wrap unnecessary: - In pool mode, cp_comm_group and cp_comm_sub_groups are owned by the pool worker (which destroys them once at pool shutdown). run_dpa_with_cp neither creates nor destroys them, so an in-body exception can't leak communicators. - In single-shot mode, groups are still created locally, but the subprocess exits at function return; NCCL releases everything at process teardown, so a stray exception leaks communicators only for the milliseconds before the process dies — a bounded one-off cost, not the unbounded accumulation that Round-1 flagged for pool mode. Removing the wrap drops the run_attention_with_cp.py diff against origin/main from ~1000 lines to ~120 lines without changing observable behaviour. Smoke-tested: 4 representative cases pass. Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Set test_essential=True to match shipping default Round-3 review (greptile, discussion_r3250016711) flagged that the working tree had test_essential=False — i.e. the full ~328-config matrix instead of the ~38-config essential subset that the rest of the CI matrix expects. Flipping back to True so CI doesn't regress baseline on the known H1-style cascade configs that only appear in the full matrix. Signed-off-by: Sudhakar Singh * Retry once on pool-infrastructure failures with stderr-logged flake trace The pool worker subprocess can die mid-case due to async NCCL aborts or flaky 4-GPU collective state that doesn't reproduce on a fresh pool. Without retry, these manifest as one-off CI failures attributable to infrastructure, not the PR's content. Add a single-attempt retry around PoolWorker.submit() that fires only on infrastructure failure modes (pool-worker-died, timeout, broken-pipe-pre-send). Test-assertion failures from the worker (resp["error"]) carry full per-rank tracebacks and propagate without retry — so a real bug still surfaces as FAILED. Visibility: every retry attempt writes a [POOL-RETRY] line to stderr. pytest captures per-test stderr and writes it into JUnit /. A flaky test will appear as PASSED in the case row but with a [POOL-RETRY] line in — visible to the reviewer, and queryable by CI dashboards looking for flake patterns (e.g. "same test_id retries across multiple CI runs"). If both attempts die, a [POOL-RETRY-FAIL] line is also logged with the first error's headline, then the second attempt's full traceback propagates as the test failure. Smoke-tested: 3 representative cases (p2p, a2a flash; p2p fused) still PASS in 19 s. Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Pool: redirect non-rank-0 stdout to /dev/null; drop sentinel Replaces the [CP_POOL_RESP] sentinel-prefix protocol with a stronger fix at the source: on rank>0, close stdout at the fd level via dup2 to /dev/null at worker startup. Catches both Python `print` writes and C-level (NCCL, libc, etc.) writes that the sentinel could only mitigate by scanning + skipping non-protocol lines. With non-rank-0 stdout silenced, rank 0's JSON line is the only thing that reaches the parent's pipe, so PoolWorker._submit_once collapses from a sentinel-scanning while loop to a single select + readline + json.loads. Closes follow-up M2 from the PR description; addresses greptile's review comment on stdout pollution. Validated on 8xH100 with the test_essential=True flash-attn pool path (9 passed / 55 skipped / 0 failed in 56s; no JSONDecodeError, no protocol corruption). Signed-off-by: Sudhakar Singh * Address PR review (R3): backend-cache, pool isolation, group-kill, decode-safety - Invalidate DotProductAttention._attention_backends between pool cases so per-case NVTE_FLASH_ATTN/NVTE_FUSED_ATTN toggles take effect instead of reusing the previous case's resolved backend. - torch.cuda.empty_cache() after each case so a 2-GPU pool doesn't squat on GPUs that an overlapping 4-GPU pool needs. - PoolWorker subprocess uses start_new_session=True; _kill() uses killpg on the whole process group so torchrun's rank workers don't survive as orphans holding CUDA/NCCL state. - On a failed worker response, kill the pool before raising so half-aborted CUDA/NCCL/FP8 state from a failed case doesn't leak into the next. - Guard json.loads with a try/except + diagnostic so any rank-0 stdout pollution surfaces as a clear test failure rather than a silent protocol desync. Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Sudhakar Singh Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../attention/run_attention_with_cp.py | 102 ++++-- .../attention/run_attention_with_cp_pool.py | 221 ++++++++++++ .../attention/test_attention_with_cp.py | 317 +++++++++++++++--- .../dot_product_attention/context_parallel.py | 6 + 4 files changed, 576 insertions(+), 70 deletions(-) create mode 100644 tests/pytorch/attention/run_attention_with_cp_pool.py diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 8dfea644a..9f6b4944e 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -2,6 +2,7 @@ # # See LICENSE for license information. +import copy import os import sys import logging @@ -29,6 +30,15 @@ ) from utils import ModelConfig, compare_and_assert +# Pool mode (NVTE_CP_POOL_PG=1) only: shared CP collective groups, created once +# per pool by run_attention_with_cp_pool.main() and reused across every case in +# that pool. world_size and the rank set don't change per case, so re-creating +# these per call would be wasted NCCL setup (~50-100 ms each). Single-shot +# subprocess mode leaves these None / [] and run_dpa_with_cp creates/destroys +# its own groups inline. +_pool_cp_comm_group = None +_pool_cp_comm_sub_groups: list = [] + dtypes = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp8": torch.bfloat16} @@ -209,10 +219,13 @@ def run_dpa_with_cp( os.environ["NVTE_FUSED_ATTN"] = "0" if kernel_backend == "FlashAttention": os.environ["NVTE_FLASH_ATTN"] = "1" - config = model_configs_flash_attn[model] + # Deep-copy: the module-level dict is shared across pool cases; the + # THD branch below rewrites attn_mask_type in place, which would + # otherwise leak into subsequent cases reusing the same model key. + config = copy.deepcopy(model_configs_flash_attn[model]) if kernel_backend == "FusedAttention": os.environ["NVTE_FUSED_ATTN"] = "1" - config = model_configs_fused_attn[model] + config = copy.deepcopy(model_configs_fused_attn[model]) assert config.attn_mask_type in [ "causal", "no_mask", @@ -226,6 +239,9 @@ def run_dpa_with_cp( # set up distributed group rank = int(os.getenv("RANK", "0")) world_size = int(os.getenv("WORLD_SIZE", "1")) + # When NVTE_CP_POOL_PG=1, the pool runner owns the lifecycle of the main + # process group across many cases; here we only reuse it. + _pool_managed_pg = os.getenv("NVTE_CP_POOL_PG", "0") == "1" if dist.is_initialized(): world_size = dist.get_world_size() rank = dist.get_rank() @@ -234,25 +250,35 @@ def run_dpa_with_cp( device = rank % device_count torch.cuda.set_device(device) logging.info(f"[Rank {rank}] Setup: world_size {world_size}") - dist.init_process_group(backend="nccl", world_size=world_size, rank=rank) - - # set up communication group for CP + if not _pool_managed_pg: + dist.init_process_group(backend="nccl", world_size=world_size, rank=rank) + + # Set up communication group for CP. In pool mode, the pool worker has + # already pre-created world-scoped and a2a+p2p sub-groups once and stashed + # them in module-level pointers; we reuse those and the pool destroys them + # at shutdown. In single-shot mode we create them per call and destroy in + # the finally below. cp_comm_ranks = range(world_size) assert rank in cp_comm_ranks - cp_comm_group = dist.new_group(cp_comm_ranks, backend="nccl") - if cp_comm_type == "a2a+p2p": - assert world_size % 2 == 0, ( - "{cp_comm_type=} requires world_size % 2 = 0 as it assumes the a2a level has cp_size" - " = 2." - ) - cp_comm_sub_ranks = [range(i * 2, (i + 1) * 2) for i in range(world_size // 2)] - cp_comm_sub_ranks += [range(i, world_size, 2) for i in range(2)] - cp_comm_sub_groups = [] - for sub_ranks in cp_comm_sub_ranks: - sub_group = dist.new_group(sub_ranks, backend="nccl") - if rank in sub_ranks: - cp_comm_sub_groups.append(sub_group) - + _reusing_pool_groups = _pool_managed_pg and _pool_cp_comm_group is not None + cp_comm_group = None + cp_comm_sub_groups: list = [] + if _reusing_pool_groups: + cp_comm_group = _pool_cp_comm_group + cp_comm_sub_groups = _pool_cp_comm_sub_groups if cp_comm_type == "a2a+p2p" else [] + else: + cp_comm_group = dist.new_group(cp_comm_ranks, backend="nccl") + if cp_comm_type == "a2a+p2p": + assert world_size % 2 == 0, ( + "{cp_comm_type=} requires world_size % 2 = 0 as it assumes the a2a level has" + " cp_size = 2." + ) + cp_comm_sub_ranks = [range(i * 2, (i + 1) * 2) for i in range(world_size // 2)] + cp_comm_sub_ranks += [range(i, world_size, 2) for i in range(2)] + for sub_ranks in cp_comm_sub_ranks: + sub_group = dist.new_group(sub_ranks, backend="nccl") + if rank in sub_ranks: + cp_comm_sub_groups.append(sub_group) if dtype == "fp8": if scaling_mode == "delayed": fp8_recipe = DelayedScaling(fp8_dpa=fp8_dpa, fp8_mha=fp8_mha) @@ -564,7 +590,10 @@ def run_dpa_with_cp( seq_kv_size = dbias.shape[-1] # Reshape to split seq_q dimension dbias = dbias.view( - *shape_before_seq, 2 * world_size, seq_q_size // (2 * world_size), seq_kv_size + *shape_before_seq, + 2 * world_size, + seq_q_size // (2 * world_size), + seq_kv_size, ) # Index select on the newly created dimension (now at position seq_q_dim) dbias = dbias.index_select(seq_q_dim, seq_idx) @@ -754,7 +783,14 @@ def run_dpa_with_cp( ) elif qkv_format == "thd": compare_and_assert( - t, tensors_cp[i], names_no_cp[i], names_cp[i], atol, rtol, rmse_tol, is_fp8 + t, + tensors_cp[i], + names_no_cp[i], + names_cp[i], + atol, + rtol, + rmse_tol, + is_fp8, ) else: compare_and_assert( @@ -762,8 +798,28 @@ def run_dpa_with_cp( ) logging.info(f"[Rank {rank}] CP vs no-CP: {names[i]} matches") - # destroy distribution group - dist.destroy_process_group() + # Teardown on the success path. Pool mode: cp_comm_group / cp_comm_sub_groups + # point at pool-shared groups owned by the pool runner (which destroys them + # at pool shutdown), and the main PG is also pool-owned — both branches + # below are no-ops. Single-shot mode: destroy what we created here. If the + # body above raises, we skip this — the subprocess dies at function return + # and NCCL releases the communicators with the process. + if not _reusing_pool_groups: + if cp_comm_group is not None: + try: + dist.destroy_process_group(cp_comm_group) + except Exception: + pass + for g in cp_comm_sub_groups: + try: + dist.destroy_process_group(g) + except Exception: + pass + if not _pool_managed_pg: + try: + dist.destroy_process_group() + except Exception: + pass def main(**kwargs): diff --git a/tests/pytorch/attention/run_attention_with_cp_pool.py b/tests/pytorch/attention/run_attention_with_cp_pool.py new file mode 100644 index 000000000..3e5f64a42 --- /dev/null +++ b/tests/pytorch/attention/run_attention_with_cp_pool.py @@ -0,0 +1,221 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +""" +Persistent worker for batched CP attention tests. + +Launched ONCE per (pytest session, world_size) by torchrun. All ranks init +NCCL, then enter a dispatch loop: + + rank 0: + read one JSON request line from stdin + broadcast it to all ranks + all ranks: + call run_dpa_with_cp(**kwargs) — the same work function the + per-case subprocess design uses, with NVTE_CP_POOL_PG=1 so the + function reuses our PG instead of re-initing it + torch.cuda.empty_cache() per case + all ranks gather (ok, error_msg) to rank 0 + rank 0: + write one JSON response line to stdout + +Protocol (line-delimited JSON over rank-0 stdio): + request : {"op": "run", "kwargs": {...}} + {"op": "shutdown"} + response: {"ok": true} + {"ok": false, "error": "first failing rank's traceback"} +""" +import json +import os +import sys +import time +import traceback + +import torch +import torch.distributed as dist + +# Make sibling modules importable when launched directly. +sys.path.append(os.path.dirname(os.path.abspath(__file__))) + +from run_attention_with_cp import run_dpa_with_cp +from transformer_engine.pytorch.quantization import FP8GlobalStateManager + + +def _recv_request(rank: int) -> dict: + box = [None] + if rank == 0: + line = sys.stdin.readline() + box[0] = {"op": "shutdown"} if not line else json.loads(line) + dist.broadcast_object_list(box, src=0) + return box[0] + + +def _send_response(rank: int, payload: dict) -> None: + if rank == 0: + sys.stdout.write(json.dumps(payload) + "\n") + sys.stdout.flush() + + +def _silence_non_rank0_stdout(rank: int) -> None: + """Redirect non-rank-0 stdout to /dev/null at fd level. + + All ranks share rank 0's stdout fd (torchrun inherits it from the launcher), + so Python/library writes on rank>0 would interleave with rank 0's JSON + protocol on the parent's pipe. Closing fd 1 at the OS level on rank>0 + catches both Python (``print``) and C-level (NCCL, etc.) writes. + """ + if rank == 0: + return + devnull = os.open(os.devnull, os.O_WRONLY) + os.dup2(devnull, 1) + os.close(devnull) + sys.stdout = open(1, "w", closefd=False) + + +def _reset_between_cases() -> None: + """Drop state that would otherwise cascade across cases. + + Matches the per-case startup of the single-shot worker + (``_run_single_config`` on the per-case-subprocess branch): identical RNG + seed at the start of every case, FP8 state cleared, allocator clean. + ``run_dpa_with_cp`` re-sets ``NVTE_FUSED_ATTN``/``NVTE_FLASH_ATTN`` + unconditionally and pops the other transient env vars itself, so no + explicit pop is needed here. + """ + torch.manual_seed(1234) + torch.cuda.manual_seed(1234) + FP8GlobalStateManager.reset() + torch.cuda.empty_cache() + # Invalidate DPA's module-level backend cache so the per-case + # NVTE_FLASH_ATTN/NVTE_FUSED_ATTN env-var toggle actually takes effect + # instead of reusing the previous case's resolved backend. + try: + from transformer_engine.pytorch.attention.dot_product_attention import dot_product_attention + + dot_product_attention._attention_backends["backend_selection_requires_update"] = True + except (ImportError, AttributeError, KeyError): + pass + + +_case_counter = 0 + + +def _run_one(req: dict, rank: int) -> tuple[bool, str]: + global _case_counter + op = req["op"] + if op != "run": + return False, f"unknown op: {op}" + # Reset BEFORE the case so the first case also starts from a known RNG seed + # and clean FP8 state — same as the single-shot worker's per-process startup. + _reset_between_cases() + t0 = time.monotonic() + ok = True + err = "" + try: + run_dpa_with_cp(**req.get("kwargs", {})) + except Exception: + ok = False + err = f"[Rank {rank}] {traceback.format_exc()}" + wall = time.monotonic() - t0 + # Per-case wall time on rank 0, opt-in via NVTE_CP_POOL_TIMING=1. + # Used to tune POOL_SUBMIT_TIMEOUT_SEC against the observed distribution. + if rank == 0 and int(os.environ.get("NVTE_CP_POOL_TIMING", "0")): + _case_counter += 1 + sys.stderr.write( + f"[POOL-TIMING] case_idx={_case_counter} " + f"world_size={int(os.environ.get('WORLD_SIZE', 0))} " + f"wall_s={wall:.3f} ok={ok}\n" + ) + sys.stderr.flush() + return ok, err + + +def _create_cp_comm_groups(rank: int, world_size: int) -> tuple: + """Pre-create the CP collective groups for this pool. + + world_size and the rank set are constant for the lifetime of one pool, so + the world group and the a2a+p2p sub-groups are deterministic. Creating + them once here and reusing them across every case eliminates ~50-100 ms + of NCCL setup per case (cyanguwa's review feedback on PR #2993). + + Returns ``(world_group, a2a_p2p_sub_groups)``. ``a2a_p2p_sub_groups`` is + empty when world_size is too small to support a2a+p2p (needs an even + world_size ≥ 4); cases with cp_comm_type='a2a+p2p' wouldn't be routed to + such a pool anyway. + """ + world_group = dist.new_group(range(world_size), backend="nccl") + sub_groups: list = [] + if world_size >= 4 and world_size % 2 == 0: + # Mirror the layout in run_attention_with_cp.py: cp_size/2 pairs along + # axis 0, plus 2 stride-2 groups along axis 1. + cp_comm_sub_ranks = [range(i * 2, (i + 1) * 2) for i in range(world_size // 2)] + cp_comm_sub_ranks += [range(i, world_size, 2) for i in range(2)] + for sub_ranks in cp_comm_sub_ranks: + sub_group = dist.new_group(sub_ranks, backend="nccl") + if rank in sub_ranks: + sub_groups.append(sub_group) + return world_group, sub_groups + + +def main() -> None: + rank = int(os.environ["RANK"]) + world_size = int(os.environ["WORLD_SIZE"]) + _silence_non_rank0_stdout(rank) + torch.cuda.set_device(rank % torch.cuda.device_count()) + dist.init_process_group(backend="nccl", rank=rank, world_size=world_size) + os.environ["NVTE_CP_POOL_PG"] = "1" + + # Stash pool-shared CP groups on the run_attention_with_cp module so + # run_dpa_with_cp can read them per case. Imported here (after the env var + # is set) to keep import-time side effects minimal. + import run_attention_with_cp as _rac + + _rac._pool_cp_comm_group, _rac._pool_cp_comm_sub_groups = _create_cp_comm_groups( + rank, world_size + ) + + try: + while True: + req = _recv_request(rank) + if req.get("op") == "shutdown": + break + + ok, msg = _run_one(req, rank) + + gathered: list[tuple[bool, str]] = [None] * world_size # type: ignore[list-item] + # gather_object is itself a collective synchronization point — if + # every rank reached it, none is ahead. No extra barrier needed. + dist.gather_object((ok, msg), gathered if rank == 0 else None, dst=0) + + if rank == 0: + all_ok = all(o for o, _ in gathered) + if all_ok: + _send_response(rank, {"ok": True}) + else: + first_err = next(m for o, m in gathered if not o) + _send_response(rank, {"ok": False, "error": first_err}) + # Release the allocator cache so this pool doesn't squat on + # GPUs that an overlapping different-world-size pool needs. + torch.cuda.empty_cache() + finally: + # Tear down pool-shared CP groups before the main PG (NCCL requires + # sub-groups to be destroyed first). Each destroy is independently + # guarded so a wedged communicator on one group doesn't leak the rest. + if _rac._pool_cp_comm_group is not None: + try: + dist.destroy_process_group(_rac._pool_cp_comm_group) + except Exception: + pass + for g in _rac._pool_cp_comm_sub_groups: + try: + dist.destroy_process_group(g) + except Exception: + pass + _rac._pool_cp_comm_group = None + _rac._pool_cp_comm_sub_groups = [] + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index 23d1bfdd8..f0d2c27c1 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -2,12 +2,18 @@ # # See LICENSE for license information. +import json import os +import select +import signal import subprocess import sys +import threading +import time import pathlib import logging import copy +from collections import deque import pytest import torch from transformer_engine.pytorch import ( @@ -24,7 +30,7 @@ _current_file = pathlib.Path(__file__).resolve() sys.path.append(str(_current_file.parent.parent)) -from utils import ModelConfig, get_available_attention_backends, run_distributed +from utils import ModelConfig, get_available_attention_backends pytest_logging_level = logging.getLevelName(logging.root.level) @@ -60,19 +66,228 @@ } -def get_bash_arguments(num_gpus_per_node, **kwargs): - args = [ - "python3", - "-m", - "torch.distributed.launch", - "--nproc-per-node=" + str(num_gpus_per_node), - ] - te_path = os.getenv("TE_PATH", "/opt/transformerengine") - script_path = os.path.join(te_path, "tests/pytorch/attention/run_attention_with_cp.py") - args.append(script_path) - for k, v in kwargs.items(): - args.append(f"{k}={v}") - return args +# --- Persistent pool runner ----------------------------------------------- +# +# Each (world_size) is served by one long-lived torchrun running +# run_attention_with_cp_pool.py. We submit one work item per pytest case over +# rank-0 stdin and read one JSON response from rank-0 stdout. Replaces +# the per-case torchrun launch path; init/destroy NCCL once per pool, not +# once per case. +# +# Why two pool sizes: cp_comm_type="a2a+p2p" needs world_size=4; everything +# else uses world_size=2. We can't resize an active PG, so we keep one pool +# per world_size and route each case to the right one. Pools are spawned +# lazily on first use so a session that only exercises 2-GPU cases never +# pays the 4-GPU init cost. + +# Per-case wall is ~5 s p50 / ~15 s max on H100 (test_essential=True). +# 90 s gives ~6× headroom over the slowest observed case while still detecting +# a genuine hang within ~1.5 min instead of ~10 min. Override with the env var +# if a slower machine or expanded test matrix needs more room. +POOL_SUBMIT_TIMEOUT_SEC = float(os.getenv("NVTE_CP_POOL_TIMEOUT_SEC", "90")) + + +class PoolWorker: + # Crash-path AssertionErrors include the tail of the worker's stderr so CI + # JUnit XML shows the actual failure cause (NCCL/CUDA messages, Python + # traceback) inline with the failing test, not just "pool worker died". + # Equivalent in spirit to PR #2965's run_distributed() stderr capture. + _STDERR_BUFFER_LINES = 200 # ring cap (~40 KB ceiling) + _STDERR_TAIL_CHARS = 4000 # how much to attach to the AssertionError + + def __init__(self, world_size: int): + self.world_size = world_size + self.proc: subprocess.Popen | None = None + self._stderr_buf: deque[str] = deque(maxlen=self._STDERR_BUFFER_LINES) + + def _spawn(self) -> None: + te_path = os.getenv("TE_PATH", "/opt/transformerengine") + worker = os.path.join(te_path, "tests/pytorch/attention/run_attention_with_cp_pool.py") + cmd = [ + sys.executable, + "-m", + "torch.distributed.run", + f"--nproc-per-node={self.world_size}", + "--standalone", # picks a free rendezvous port + worker, + ] + # stderr=PIPE so we can capture the tail for crash-path AssertionErrors; + # a daemon drainer thread also echoes each line to sys.stderr so pytest's + # per-test stderr capture still works. The thread is daemon, so it + # self-terminates when the pipe closes — no tracking needed. + self.proc = subprocess.Popen( + cmd, + stdin=subprocess.PIPE, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1, + env={**os.environ, "PYTHONUNBUFFERED": "1"}, + # Own process group so _kill can killpg all ranks in one shot; + # without this, terminating the launcher PID leaves rank workers + # as orphans holding CUDA/NCCL state. + start_new_session=True, + ) + self._stderr_buf.clear() + threading.Thread(target=self._drain_stderr, daemon=True).start() + + def _drain_stderr(self) -> None: + proc = self.proc + if proc is None or proc.stderr is None: + return + for line in iter(proc.stderr.readline, ""): + self._stderr_buf.append(line) + sys.stderr.write(line) + sys.stderr.flush() + + def _diag(self, msg: str) -> str: + tail = "".join(self._stderr_buf)[-self._STDERR_TAIL_CHARS :] + if not tail.strip(): + return msg + return f"{msg}\n\n--- pool worker stderr (tail) ---\n{tail}" + + def _ensure_alive(self) -> None: + if self.proc is None or self.proc.poll() is not None: + self._spawn() + + def _killpg(self, sig: int) -> None: + try: + os.killpg(self.proc.pid, sig) + except ProcessLookupError: + pass + + def _kill(self) -> None: + # Kill the whole process group so rank workers don't survive as orphans. + if self.proc and self.proc.poll() is None: + self._killpg(signal.SIGTERM) + try: + self.proc.wait(timeout=5) + except subprocess.TimeoutExpired: + self._killpg(signal.SIGKILL) + self.proc.wait() + self.proc = None + + # One retry on pool-infrastructure failures (worker died / timed out / broken + # pipe). Test-assertion failures from the worker carry the full per-rank + # traceback in resp["error"] and propagate without retry. Every retry leaves + # a [POOL-RETRY] line in stderr so pytest's capture surfaces + # flake patterns in JUnit XML for offline analysis. + _MAX_RETRIES = 1 + + def submit(self, kwargs: dict, timeout: float = POOL_SUBMIT_TIMEOUT_SEC) -> None: + first_err = None + for attempt in range(self._MAX_RETRIES + 1): + try: + return self._submit_once(kwargs, timeout) + except AssertionError as e: + msg_head = str(e).splitlines()[0] + infrastructure_flake = ( + "pool worker died" in msg_head + or "timed out" in msg_head + or "before request could be sent" in msg_head + ) + if not infrastructure_flake or attempt == self._MAX_RETRIES: + if first_err is not None: + sys.stderr.write( + f"[POOL-RETRY-FAIL] world_size={self.world_size}: " + "both attempts died; first error was: " + f"{str(first_err).splitlines()[0]!r}\n" + ) + sys.stderr.flush() + raise + first_err = e + sys.stderr.write( + f"[POOL-RETRY] world_size={self.world_size} attempt {attempt + 1} " + f"died: {msg_head!r}; respawning pool and retrying\n" + ) + sys.stderr.flush() + raise first_err # unreachable; loop either returns or raises + + def _submit_once(self, kwargs: dict, timeout: float) -> None: + self._ensure_alive() + req = json.dumps({"op": "run", "kwargs": kwargs}) + "\n" + try: + self.proc.stdin.write(req) + self.proc.stdin.flush() + except BrokenPipeError: + msg = self._diag("pool worker died before request could be sent") + self._kill() + raise AssertionError(msg) + + # Worker redirects non-rank-0 stdout to /dev/null at fd level, so + # rank 0's JSON line is the only thing that arrives on this pipe. + # select() on a pipe fd is Linux/macOS only — on Windows the select + # module only accepts sockets. CP attention tests run on Linux GPU + # hosts so this is fine; flag if portability is ever needed. + ready, _, _ = select.select([self.proc.stdout], [], [], timeout) + if not ready: + msg = self._diag( + f"pool worker (world_size={self.world_size}) timed out after " + f"{timeout}s; pool killed and will be respawned for the next case" + ) + self._kill() + raise AssertionError(msg) + + line = self.proc.stdout.readline() + if not line: + msg = self._diag("pool worker died mid-request") + self._kill() + raise AssertionError(msg) + + # A stray non-JSON line from rank 0 would desynchronize the protocol; + # turn it into a clear test failure rather than a raw JSONDecodeError. + try: + resp = json.loads(line) + except json.JSONDecodeError as e: + self._kill() + raise AssertionError( + self._diag(f"pool worker JSON protocol broke: {e!r}; line={line!r}") + ) + + if not resp["ok"]: + # Discard the pool so half-aborted CUDA/NCCL/FP8 state from the + # failed case doesn't leak into the next. resp["error"] already + # carries the per-rank traceback via gather_object. + self._kill() + raise AssertionError(resp["error"]) + + def shutdown(self) -> None: + if self.proc and self.proc.poll() is None: + try: + self.proc.stdin.write(json.dumps({"op": "shutdown"}) + "\n") + self.proc.stdin.flush() + self.proc.stdin.close() + except BrokenPipeError: + pass + try: + self.proc.wait(timeout=30) + except subprocess.TimeoutExpired: + self._kill() + self.proc = None + + +@pytest.fixture(scope="session") +def cp_pool(): + """Returns a callable: cp_pool(world_size) -> PoolWorker.""" + pools: dict[int, PoolWorker] = {} + + def _get(world_size: int) -> PoolWorker: + if world_size > torch.cuda.device_count(): + pytest.skip(f"Test requires {world_size} GPUs, but found {torch.cuda.device_count()}") + if world_size not in pools: + pools[world_size] = PoolWorker(world_size) + return pools[world_size] + + yield _get + for p in pools.values(): + p.shutdown() + + +def _submit(pool: PoolWorker, **kwargs) -> None: + # run_dpa_with_cp expects all kwargs as strings (it does e.g. + # `fp8_bwd == "True"`), matching the old argv-based path. Serialize + # everything as strings so we don't accidentally change semantics. + pool.submit({k: str(v) for k, v in kwargs.items()}) dtypes = ["bf16", "fp16"] @@ -91,10 +306,9 @@ def get_bash_arguments(num_gpus_per_node, **kwargs): @pytest.mark.parametrize("model", model_configs_flash_attn.keys()) @pytest.mark.parametrize("qkv_format", qkv_formats) @pytest.mark.parametrize("cp_comm_type", cp_comm_types) -def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): +def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type): num_gpus = 4 if cp_comm_type == "a2a+p2p" else 2 - if num_gpus > torch.cuda.device_count(): - pytest.skip(f"Test requires {num_gpus} GPUs, but found {torch.cuda.device_count()}") + pool = cp_pool(num_gpus) config = model_configs_flash_attn[model] config.context_parallel = True @@ -140,16 +354,14 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): if not flash_attn_supported: pytest.skip("No attention backend available.") - run_distributed( - get_bash_arguments( - num_gpus_per_node=num_gpus, - dtype=dtype, - model=model, - qkv_format=qkv_format, - kernel_backend="FlashAttention", - cp_comm_type=cp_comm_type, - log_level=pytest_logging_level, - ), + _submit( + pool, + dtype=dtype, + model=model, + qkv_format=qkv_format, + kernel_backend="FlashAttention", + cp_comm_type=cp_comm_type, + log_level=pytest_logging_level, ) @@ -274,15 +486,23 @@ def test_cp_with_flash_attention(dtype, model, qkv_format, cp_comm_type): @pytest.mark.parametrize("scaling_mode", [None, "delayed", "current", "mxfp8"]) @pytest.mark.parametrize("f16_O", [True, False]) def test_cp_with_fused_attention( - dtype, model, qkv_format, cp_comm_type, fp8_bwd, fp8_mha, fp8_dpa, scaling_mode, f16_O + cp_pool, + dtype, + model, + qkv_format, + cp_comm_type, + fp8_bwd, + fp8_mha, + fp8_dpa, + scaling_mode, + f16_O, ): config = model_configs_fused_attn[model] config.context_parallel = True config.cp_comm_type = cp_comm_type num_gpus = 4 if cp_comm_type == "a2a+p2p" else 2 - if num_gpus > torch.cuda.device_count(): - pytest.skip(f"Test requires {num_gpus} GPUs, but found {torch.cuda.device_count()} GPUs.") + pool = cp_pool(num_gpus) if get_device_compute_capability() < (9, 0) and qkv_format == "thd": pytest.skip("Only sm90+ architectures support THD format!") @@ -404,21 +624,24 @@ def test_cp_with_fused_attention( if not fused_attn_supported: pytest.skip("No attention backend available.") - run_distributed( - get_bash_arguments( - num_gpus_per_node=num_gpus, - dtype=dtype, - model=model, - qkv_format=qkv_format, - kernel_backend="FusedAttention", - cp_comm_type=cp_comm_type, - fp8_bwd=fp8_bwd, - fp8_dpa=fp8_dpa, - fp8_mha=fp8_mha, - scaling_mode=scaling_mode, - f16_O=f16_O, - is_training=is_training, - deterministic=_deterministic, - log_level=pytest_logging_level, - ), + if _deterministic and config.softmax_type != "vanilla": + pytest.skip("Deterministic mode does not support non-vanilla softmax with FusedAttention") + if _deterministic and config.attn_bias_type == "post_scale_bias" and is_training: + pytest.skip("Deterministic mode does not support post_scale_bias with requires_grad") + + _submit( + pool, + dtype=dtype, + model=model, + qkv_format=qkv_format, + kernel_backend="FusedAttention", + cp_comm_type=cp_comm_type, + fp8_bwd=fp8_bwd, + fp8_dpa=fp8_dpa, + fp8_mha=fp8_mha, + scaling_mode=scaling_mode, + f16_O=f16_O, + is_training=is_training, + deterministic=_deterministic, + log_level=pytest_logging_level, ) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 3db0417bd..35684625a 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -3277,6 +3277,12 @@ def forward( elif o_format == "sbhd": out_f16[i - 1].copy_(out_per_step[i - 1]) if return_max_logit: + # max_logit_per_step[i-1] was written on flash_attn_streams[i-1] + # (cp_stream for i-1=1). The torch.maximum below runs on the + # default stream, so without this wait the read can race with + # the write. The post-loop wait_stream(cp_stream) is too late. + # No-op when flash_attn_streams[i-1] is current_stream(). + torch.cuda.current_stream().wait_stream(flash_attn_streams[i - 1]) max_logit = torch.maximum(max_logit, max_logit_per_step[i - 1]) torch.cuda.current_stream().wait_stream(cp_stream) From 856d075cdd1b923e5b12658e120f8d9c37518123 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Thu, 21 May 2026 18:06:28 -0700 Subject: [PATCH 077/180] Update cudnn-frontend to 1.24.0 (#3016) update cudnn-fe 1.24 Signed-off-by: Sudhakar Singh --- 3rdparty/cudnn-frontend | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/cudnn-frontend b/3rdparty/cudnn-frontend index fb682ce76..c4a97621e 160000 --- a/3rdparty/cudnn-frontend +++ b/3rdparty/cudnn-frontend @@ -1 +1 @@ -Subproject commit fb682ce761a2705e40f9b5d528737a3e0eb33cec +Subproject commit c4a97621eca52fa0c3a1862a411a16be580b25c6 From 9af70a8bbba216b43cca6cc6428ac97a7e978acf Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Fri, 22 May 2026 08:07:18 -0700 Subject: [PATCH 078/180] [Pytorch][Bug] DCP Checkpoint Loading Fixes for FSDP2 with QuantizedModelInit (#2974) * all changes in Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * simplify Signed-off-by: Varun Thumbe * address review comment Signed-off-by: Varun Thumbe * fix CI, Test for CPU quantized tensor Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * add things thats just necessary Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix test Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fix errors Signed-off-by: Varun Thumbe * address review comments Signed-off-by: Varun Thumbe --------- Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../fsdp2_tests/run_fsdp2_fused_adam.py | 41 +----- .../fsdp2_tests/run_fsdp2_model.py | 9 -- tests/pytorch/test_quantized_tensor.py | 50 +++++++ .../common/util/pybind_helper.h | 11 +- transformer_engine/pytorch/__init__.py | 58 ++++++++ transformer_engine/pytorch/module/base.py | 5 +- .../pytorch/quantized_tensor.py | 62 ++++++++- .../pytorch/tensor/_quantization_helpers.py | 1 + .../pytorch/tensor/float8_blockwise_tensor.py | 127 +++++++++++------- .../pytorch/tensor/float8_tensor.py | 81 +++++++---- .../pytorch/tensor/mxfp8_tensor.py | 114 +++++++++++----- .../pytorch/tensor/nvfp4_tensor.py | 123 ++++++++++++----- .../tensor/storage/float8_tensor_storage.py | 5 - .../tensor/storage/mxfp8_tensor_storage.py | 16 ++- .../tensor/storage/nvfp4_tensor_storage.py | 17 ++- 15 files changed, 501 insertions(+), 219 deletions(-) diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py index ecda481ed..1abb49e98 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_fused_adam.py @@ -228,13 +228,6 @@ def test_fused_adam_fp8_master_weights_no_meta(recipe_name): """ recipe = get_recipe_from_string(recipe_name) - if recipe_name in ("MXFP8BlockScaling", "Float8BlockScaling", "NVFP4BlockScaling"): - pytest.xfail( - f"{recipe_name}: FSDP2 all-gather hooks for block-scaling QuantizedTensor " - "subclasses fail when parameters are initialized on CUDA. " - "Use device='meta' + reset_parameters() after sharding." - ) - world_size, device = _get_dist_info() model = _build_model(fp8_init=True, recipe=recipe, use_meta_device=False) @@ -604,12 +597,6 @@ def test_safetensors_fp32_export(recipe_name): - Saved tensor shapes match expected (unsharded) shapes """ recipe = get_recipe_from_string(recipe_name) - if recipe_name == "MXFP8BlockScaling": - pytest.xfail( - "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " - "MXFP8 quantized tensors, causing illegal memory access. " - "Fixed by https://github.com/NVIDIA/TransformerEngine/pull/2789." - ) from safetensors.torch import load_file, save_file from torch.distributed.checkpoint.state_dict import ( @@ -692,26 +679,7 @@ def test_dcp_output_parity(recipe_name, async_save): """ recipe = get_recipe_from_string(recipe_name) - if recipe_name == "MXFP8BlockScaling": - pytest.xfail( - "MXFP8BlockScaling: FusedAdam CUDA kernel does not support " - "MXFP8 quantized tensors, causing illegal memory access: " - "/transformer_engine/common/multi_tensor/multi_tensor_apply.cuh:92 in function " - "multi_tensor_apply: CUDA Error: an illegal memory access was encountered. " - "Fixed by https://github.com/NVIDIA/TransformerEngine/pull/2789." - ) - - if recipe_name == "NVFP4BlockScaling": - pytest.xfail( - "NVFP4BlockScaling: DCP load_state_dict triggers reset_sharded_param() " - "which calls data_ptr() on NVFP4Tensor wrapper subclass with invalid storage" - ) - - if ( - recipe_name == "Float8BlockScaling" - and not async_save - and torch.cuda.get_device_capability()[0] == 12 - ): + if recipe_name == "Float8BlockScaling" and torch.cuda.get_device_capability()[0] == 12: pytest.xfail( "Float8BlockScaling is failing on SM120 with RuntimeError: " "transformer_engine/common/transpose/quantize_transpose_vector_blockwise.cu:534 " @@ -719,13 +687,6 @@ def test_dcp_output_parity(recipe_name, async_save): "Blackwell and newer, the FP8 block scaling recipe is emulated with MXFP8, which " "requires using power of two scaling factors." ) - if recipe_name == "Float8BlockScaling" and async_save: - pytest.xfail( - "Float8BlockScaling: async DCP save/load round-trip produces different model " - "outputs — quantization metadata (scales) is not correctly persisted through " - "async distributed checkpointing. On SM120, additionally fails with pow2_scale " - "assertion in quantize_transpose_vector_blockwise." - ) import torch.distributed.checkpoint as dcp diff --git a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py index 6342e63e7..36ac307b9 100644 --- a/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py +++ b/tests/pytorch/distributed/fsdp2_tests/run_fsdp2_model.py @@ -379,20 +379,11 @@ def test_distributed(recipe_name, fp8_init, sharding_dims, layer_type): "sending only 1 tensor (scale is per-tensor metadata). Fix: concatenate MXFP8 " "data and scale_inv into a single buffer in pre_all_gather, split in post." ) - if recipe_name == "Float8BlockScaling" and fp8_init: pytest.xfail( "Float8BlockScaling + fp8_init: scale inverse padding is not handled " "correctly during FSDP2 all-gather slice ops." ) - if recipe_name == "NVFP4BlockScaling" and fp8_init and layer_type == "TransformerLayer": - pytest.xfail( - "NVFP4BlockScaling + fp8_init + TransformerLayer: " - "_check_fp8_fsdp2_allgather numerical error compounds across multiple " - "linear layers in the transformer block (up to ~1e-2 max abs diff). " - "LayerNormLinear passes with relaxed tolerances. " - "NVFP4 + FSDP2 training is validated by run_fsdp2_fused_adam.py." - ) torch.manual_seed(42) torch.cuda.manual_seed(42) diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index 526045e43..119914fbc 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -616,6 +616,56 @@ def test_identity_op( torch.testing.assert_close(y_test, y_ref, **tols) torch.testing.assert_close(dx_test, dx_ref, **tols) + @pytest.mark.parametrize("quantization", _quantization_list) + def test_cpu_dequantize( + self, + *, + quantization: str, + shape: Iterable[int] = (128, 128), + dtype: torch.dtype = torch.bfloat16, + ) -> None: + """Dequantize on a CPU-resident QuantizedTensor.""" + + # Construct a quantized tensor on CUDA. + _, x_cuda = make_reference_and_test_tensors( + shape=shape, + quantization=quantization, + test_dtype=dtype, + requires_grad=False, + ) + assert isinstance(x_cuda, QuantizedTensor) + assert x_cuda.device.type == "cuda" + + # Reference: dequantize on CUDA, then move the dense result to CPU. + ref_cpu = x_cuda.dequantize().to(device="cpu") + + # Move the QuantizedTensor itself to CPU and dequantize there. + # ``.cpu()`` routes through ``aten._to_copy.default`` so all inner + # buffers (data, scales, amax) are moved to CPU. + x_cpu = x_cuda.cpu() + assert isinstance(x_cpu, QuantizedTensor) + assert x_cpu.device.type == "cpu" + for attr in ( + "_data", + "_rowwise_data", + "_columnwise_data", + "_rowwise_scale_inv", + "_columnwise_scale_inv", + "_amax_rowwise", + "_amax_columnwise", + ): + buf = getattr(x_cpu, attr, None) + if buf is not None: + assert buf.device.type == "cpu", f"{attr} did not move to CPU" + + # Dequantize the CPU tensor. Implementation may bounce through CUDA + # internally, but must return a CPU tensor. + y_cpu = x_cpu.dequantize() + assert y_cpu.device.type == "cpu" + assert y_cpu.dtype == ref_cpu.dtype + assert y_cpu.shape == ref_cpu.shape + torch.testing.assert_close(y_cpu, ref_cpu, rtol=0, atol=0) + @pytest.mark.parametrize("quantization", _quantization_list) @pytest.mark.parametrize("dim", [0, 1]) def test_chunk( diff --git a/transformer_engine/common/util/pybind_helper.h b/transformer_engine/common/util/pybind_helper.h index ef7687e3e..ed48fe4d6 100644 --- a/transformer_engine/common/util/pybind_helper.h +++ b/transformer_engine/common/util/pybind_helper.h @@ -23,7 +23,16 @@ .value("kBFloat16", transformer_engine::DType::kBFloat16) \ .value("kFloat8E4M3", transformer_engine::DType::kFloat8E4M3) \ .value("kFloat8E5M2", transformer_engine::DType::kFloat8E5M2) \ - .value("kFloat4E2M1", transformer_engine::DType::kFloat4E2M1); \ + .value("kFloat4E2M1", transformer_engine::DType::kFloat4E2M1) \ + .def("__reduce_ex__", \ + [](transformer_engine::DType self, pybind11::object /*protocol*/) { \ + return pybind11::make_tuple(pybind11::type::of(pybind11::cast(self)), \ + pybind11::make_tuple(static_cast(self))); \ + }) \ + .def("__reduce__", [](transformer_engine::DType self) { \ + return pybind11::make_tuple(pybind11::type::of(pybind11::cast(self)), \ + pybind11::make_tuple(static_cast(self))); \ + }); \ pybind11::enum_(m, "NVTE_Bias_Type", pybind11::module_local()) \ .value("NVTE_NO_BIAS", NVTE_Bias_Type::NVTE_NO_BIAS) \ .value("NVTE_PRE_SCALE_BIAS", NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS) \ diff --git a/transformer_engine/pytorch/__init__.py b/transformer_engine/pytorch/__init__.py index 3ff0d75ee..7653d5992 100644 --- a/transformer_engine/pytorch/__init__.py +++ b/transformer_engine/pytorch/__init__.py @@ -89,8 +89,66 @@ from transformer_engine.pytorch.tensor import MXFP8Tensor from transformer_engine.pytorch.tensor import Float8BlockwiseQTensor from transformer_engine.pytorch.tensor import NVFP4Tensor +from transformer_engine.pytorch.tensor.float8_tensor import ( + _make_float8_tensor_in_reduce_ex, +) +from transformer_engine.pytorch.tensor.mxfp8_tensor import ( + _make_mxfp8_tensor_in_reduce_ex, +) +from transformer_engine.pytorch.tensor.nvfp4_tensor import ( + _make_nvfp4_tensor_in_reduce_ex, +) +from transformer_engine.pytorch.tensor.float8_blockwise_tensor import ( + _make_float8_blockwise_tensor_in_reduce_ex, +) try: torch._dynamo.config.error_on_nested_jit_trace = False except AttributeError: pass # error_on_nested_jit_trace was added in PyTorch 2.2.0 + +# To allow for safe unpickling of QuantizedTensors when using DCP +# checkpointing with FSDP2. ``tex.DType`` (the pybind11 enum) has its +# ``__reduce_ex__`` / ``__reduce__`` overridden in the C++ binding (see +# ``transformer_engine/common/util/pybind_helper.h``) so its pickle +# stream encodes as ``(tex.DType, (int,))`` and only the class itself +# needs to be allow-listed below. +try: + from torch.serialization import add_safe_globals + import transformer_engine_torch as tex + + add_safe_globals( + [ + # Storage mixins (used during pickling of internal-only tensors) + QuantizedTensorStorage, + Float8TensorStorage, + MXFP8TensorStorage, + NVFP4TensorStorage, + Float8BlockwiseQTensorStorage, + # Quantizer types embedded in metadata + Quantizer, + Float8Quantizer, + Float8CurrentScalingQuantizer, + MXFP8Quantizer, + NVFP4Quantizer, + Float8BlockQuantizer, + # pybind11 enum used as Quantizer.dtype + tex.DType, + # __reduce_ex__ reconstructors (module-level functions). + _make_float8_tensor_in_reduce_ex, + _make_mxfp8_tensor_in_reduce_ex, + _make_nvfp4_tensor_in_reduce_ex, + _make_float8_blockwise_tensor_in_reduce_ex, + ] + ) +except (ImportError, AttributeError): + import warnings as _warnings + + _warnings.warn( + "transformer_engine: torch.serialization.add_safe_globals is " + "unavailable on this PyTorch version (added in 2.4). DCP " + "checkpointing of QuantizedTensor weights with FSDP2 will not " + "work; upgrade to PyTorch >= 2.4 to enable it.", + RuntimeWarning, + stacklevel=2, + ) diff --git a/transformer_engine/pytorch/module/base.py b/transformer_engine/pytorch/module/base.py index 746177ec7..a1213fe49 100644 --- a/transformer_engine/pytorch/module/base.py +++ b/transformer_engine/pytorch/module/base.py @@ -44,6 +44,7 @@ from ..quantized_tensor import QuantizedTensor, QuantizedTensorStorage, Quantizer from ..tensor.float8_tensor import Float8Quantizer, Float8CurrentScalingQuantizer from ..tensor.mxfp8_tensor import MXFP8Quantizer +from ..tensor.nvfp4_tensor import NVFP4Quantizer from ..tensor.float8_blockwise_tensor import Float8BlockQuantizer from ..tensor.storage.float8_tensor_storage import Float8TensorStorage from ..tensor.storage.mxfp8_tensor_storage import MXFP8TensorStorage @@ -1641,7 +1642,9 @@ def reset_parameters(self, defer_init: Optional[bool] = False) -> None: raise RuntimeError("Weight quantizer has not been initialized") quantizer.set_usage(rowwise=True, columnwise=torch.is_grad_enabled()) quantizer.internal = False - if is_dtensor and isinstance(quantizer, Float8CurrentScalingQuantizer): + if is_dtensor and isinstance( + quantizer, (Float8CurrentScalingQuantizer, NVFP4Quantizer) + ): device_mesh = dtensor_param.device_mesh amax_reduction_group = ( device_mesh.get_group(mesh_dim="shard") diff --git a/transformer_engine/pytorch/quantized_tensor.py b/transformer_engine/pytorch/quantized_tensor.py index 7163e2b17..404796fd6 100644 --- a/transformer_engine/pytorch/quantized_tensor.py +++ b/transformer_engine/pytorch/quantized_tensor.py @@ -552,9 +552,26 @@ def half(self) -> torch.Tensor: # pylint: disable=missing-function-docstring return self.dequantize(dtype=torch.float16) - def cpu(self, memory_format=torch.preserve_format) -> torch.Tensor: + def cpu(self, memory_format=torch.preserve_format) -> QuantizedTensor: + """Move tensor to CPU while preserving the QuantizedTensor type. + + Routes through ``aten._to_copy.default`` so the subclass-preserving + handler in ``__torch_dispatch__`` runs (rather than dequantizing). + + """ # pylint: disable=missing-function-docstring - return self.dequantize().cpu(memory_format=memory_format) + return self.to(device=torch.device("cpu"), memory_format=memory_format) + + def untyped_storage(self) -> torch.UntypedStorage: + """Return an empty UntypedStorage on the tensor's device. + + ``QuantizedTensor`` is a ``_make_wrapper_subclass`` and has no real + backing storage of its own; the actual bytes live in the inner + buffers (e.g. ``_rowwise_data`` / ``_columnwise_data``) which are + an implementation detail of the quantization scheme. Need to define + this method to avoid DCP staging errors with FSDP2. + """ + return torch.UntypedStorage(0, device=self.device) def expand_as(self, other: torch.Tensor) -> torch.Tensor: # pylint: disable=missing-function-docstring @@ -608,6 +625,36 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): dst.copy_(src) return None + # _to_copy op (used by .to(device=...), .cpu(), DCP staging). + # Preserve the QuantizedTensor subclass and move all internal + # buffers (data, scales, etc.) to the requested device. + if func == torch.ops.aten._to_copy.default: + tensor = args[0] + kw = dict(kwargs) if kwargs else {} + dtype = kw.get("dtype", None) + if dtype is None or dtype == tensor.dtype: + target_device = kw.get("device", tensor.device) or tensor.device + target_device = torch.device(target_device) + pin_memory = bool(kw.get("pin_memory", False)) + non_blocking = bool(kw.get("non_blocking", False)) + new_metadata = {"device": target_device} + # Update tensor storage metadata + for key, value in tensor.get_metadata().items(): + if isinstance(value, torch.Tensor): + value = value.to(device=target_device, non_blocking=non_blocking) + if pin_memory and target_device.type == "cpu": + value = value.pin_memory() + new_metadata[key] = value + # Update torch Tensor metadata + new_metadata.update( + { + "dtype": tensor.dtype, + "shape": tensor.shape, + "requires_grad": tensor.requires_grad, + } + ) + return type(tensor)(**new_metadata) + # View op if func == torch.ops.aten.view.default: raise NotImplementedError("{cls.__name__} class does not support tensor views") @@ -748,14 +795,19 @@ def make_like( """Create new quantized tensor By default, new tensor has the same attributes and underlying - data. This function is intended to create view of tensors. - + data. This function is intended to create a view of ``tensor``, """ shape = shape if shape is not None else tensor.shape dtype = dtype if dtype is not None else tensor.dtype kwargs = tensor.get_metadata() kwargs["fake_dtype"] = dtype - return cls(shape=shape, dtype=dtype, requires_grad=requires_grad, **kwargs) + return cls( + shape=shape, + dtype=dtype, + requires_grad=requires_grad, + device=tensor.device, + **kwargs, + ) def to_dtype(self, dtype: torch.dtype) -> QuantizedTensor: """Create `QuantizedTensor` with given nominal dtype diff --git a/transformer_engine/pytorch/tensor/_quantization_helpers.py b/transformer_engine/pytorch/tensor/_quantization_helpers.py index ba3407e13..56cf50363 100644 --- a/transformer_engine/pytorch/tensor/_quantization_helpers.py +++ b/transformer_engine/pytorch/tensor/_quantization_helpers.py @@ -61,6 +61,7 @@ def forward( kwargs = tensor.get_metadata() for key, val in init_kwargs.items(): kwargs[key] = val + kwargs["device"] = tensor.device return type(tensor)(tensor.shape, tensor.dtype, **kwargs) @staticmethod diff --git a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py index d0296902a..e091e27e5 100644 --- a/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_blockwise_tensor.py @@ -333,21 +333,6 @@ def reshape(self, *shape: Tuple[int]) -> Float8BlockwiseQTensor: # pylint: disable=missing-function-docstring return _ReshapeFunc.apply(self, shape) - def untyped_storage(self) -> torch.UntypedStorage: - """Return the underlying UntypedStorage of the FP8 data. - - Note that FP8 block-scaled tensor may involve multiple - buffers: row-wise FP8 data, row-wise scales, column-wise FP8 - data, column-wise scales. The UntypedStorage of the row-wise - FP8 data is returned if it exists, and otherwise the - UntypedStorage of the column-wise FP8 data. - - """ - data = self._rowwise_data if self._rowwise_data is not None else self._columnwise_data - if data is not None: - return data.untyped_storage() - return torch.UntypedStorage(0, device=self.device) - @classmethod def __torch_dispatch__(cls, func, types, args, kwargs=None): @@ -432,42 +417,10 @@ def contiguous( return self raise ValueError("Float8BlockwiseQTensor does not support different memory formats!") - @classmethod - def _make_in_reduce_ex( - cls, - shape: torch.Size, - rowwise_data: torch.Tensor, - rowwise_scale_inv: torch.Tensor, - columnwise_data: torch.Tensor, - columnwise_scale_inv: torch.Tensor, - fp8_dtype: TE_DType, - dtype: torch.dtype, - quantizer: Quantizer, - is_2D_scaled: bool, - data_format: Any = None, # pylint: disable=unused-argument - ) -> Float8BlockwiseQTensor: - """Build Float8BlockwiseQTensor, for use in __reduce__ - - __reduce_ex__ assumes object constructor has positional - arguments. - - """ - return Float8BlockwiseQTensor( - shape=shape, - rowwise_data=rowwise_data, - rowwise_scale_inv=rowwise_scale_inv, - fp8_dtype=fp8_dtype, - columnwise_data=columnwise_data, - columnwise_scale_inv=columnwise_scale_inv, - dtype=dtype, - quantizer=quantizer, - is_2D_scaled=is_2D_scaled, - ) - def __reduce_ex__(self, protocol: int) -> tuple: """Custom pickling to remove references to FP8 metadata objects""" return ( - Float8BlockwiseQTensor._make_in_reduce_ex, + _make_float8_blockwise_tensor_in_reduce_ex, ( self.shape, self._rowwise_data, @@ -482,6 +435,45 @@ def __reduce_ex__(self, protocol: int) -> tuple: ), ) + @classmethod + def _make_in_reduce_ex( + cls, + shape: torch.Size, + rowwise_data: torch.Tensor, + rowwise_scale_inv: torch.Tensor, + columnwise_data: torch.Tensor, + columnwise_scale_inv: torch.Tensor, + fp8_dtype: TE_DType, + dtype: torch.dtype, + quantizer: Quantizer, + is_2D_scaled: bool, + data_format: Any = None, + ) -> Float8BlockwiseQTensor: + """This classmethod is kept for backward compatibility only. + ``__reduce_ex__`` used to point at this classmethod, but bound + classmethods pickle as ``(getattr, (cls, name))`` which adds an + extra reduction step to the pickle stream. The current + ``__reduce_ex__`` references the module-level + ``_make_float8_blockwise_tensor_in_reduce_ex`` instead so the + pickle stream uses a single ``GLOBAL`` opcode. This classmethod + is retained so that previously pickled ``Float8BlockwiseQTensor`` + payloads (which still reference + ``Float8BlockwiseQTensor._make_in_reduce_ex``) can still be + unpickled. + """ + return _make_float8_blockwise_tensor_in_reduce_ex( + shape, + rowwise_data, + rowwise_scale_inv, + columnwise_data, + columnwise_scale_inv, + fp8_dtype, + dtype, + quantizer, + is_2D_scaled, + data_format, + ) + def _get_data(self) -> Float8BlockwiseQTensor: """Get tensor data property""" return self @@ -653,6 +645,7 @@ def fsdp_post_all_gather( columnwise_scale_inv=None, quantizer=self._quantizer, is_2D_scaled=is_2D_scaled, + device=rowwise_data.device, ) # For 2D block scaling, derive columnwise data and scales from rowwise @@ -668,6 +661,40 @@ def fsdp_post_all_gather( return out, all_gather_outputs +def _make_float8_blockwise_tensor_in_reduce_ex( + shape: torch.Size, + rowwise_data: torch.Tensor, + rowwise_scale_inv: torch.Tensor, + columnwise_data: torch.Tensor, + columnwise_scale_inv: torch.Tensor, + fp8_dtype: TE_DType, + dtype: torch.dtype, + quantizer: Quantizer, + is_2D_scaled: bool, + data_format: Any = None, # pylint: disable=unused-argument +) -> Float8BlockwiseQTensor: + """Reconstruct a ``Float8BlockwiseQTensor`` from its ``__reduce_ex__`` payload.""" + # Infer device from inner buffers so the wrapper subclass stays + # consistent with its data (e.g. CPU after DCP staging deserialize). + device = None + if rowwise_data is not None: + device = rowwise_data.device + elif columnwise_data is not None: + device = columnwise_data.device + return Float8BlockwiseQTensor( + shape=shape, + rowwise_data=rowwise_data, + rowwise_scale_inv=rowwise_scale_inv, + fp8_dtype=fp8_dtype, + columnwise_data=columnwise_data, + columnwise_scale_inv=columnwise_scale_inv, + dtype=dtype, + quantizer=quantizer, + is_2D_scaled=is_2D_scaled, + device=device, + ) + + class _ViewFunc(torch.autograd.Function): """View function @@ -749,6 +776,7 @@ def forward( quantizer=tensor._quantizer, is_2D_scaled=tensor._is_2D_scaled, requires_grad=tensor.requires_grad, + device=tensor.device, ) @staticmethod @@ -778,6 +806,7 @@ def backward( quantizer=grad._quantizer, is_2D_scaled=grad._is_2D_scaled, requires_grad=grad.requires_grad, + device=grad.device, ) return dgrad, None return grad.view(ctx.shape), None @@ -863,6 +892,7 @@ def forward( quantizer=tensor._quantizer, is_2D_scaled=tensor._is_2D_scaled, requires_grad=tensor.requires_grad, + device=tensor.device, ) @staticmethod @@ -891,6 +921,7 @@ def backward( quantizer=grad._quantizer, is_2D_scaled=grad._is_2D_scaled, requires_grad=grad.requires_grad, + device=grad.device, ) return dgrad, None return grad.view(ctx.shape), None diff --git a/transformer_engine/pytorch/tensor/float8_tensor.py b/transformer_engine/pytorch/tensor/float8_tensor.py index c4c5934f9..7842ccc12 100644 --- a/transformer_engine/pytorch/tensor/float8_tensor.py +++ b/transformer_engine/pytorch/tensor/float8_tensor.py @@ -154,6 +154,7 @@ def create_tensor_from_data( requires_grad=requires_grad, data_transpose=None, quantizer=self, + device=data.device, ) def onnx_quantize(self, tensor: torch.Tensor) -> QuantizedTensor: @@ -335,6 +336,7 @@ def create_tensor_from_data( requires_grad=requires_grad, data_transpose=None, quantizer=self, + device=data.device, ) def get_columnwise_shape(self, rowwise_data_shape: Iterable[int]) -> Tuple[int, ...]: @@ -355,6 +357,7 @@ def onnx_quantize(self, tensor: torch.Tensor) -> QuantizedTensor: requires_grad=False, data_transpose=None, quantizer=self, + device=data.device, ) def onnx_dequantize(self, tensor: QuantizedTensor) -> torch.Tensor: @@ -587,6 +590,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): fp8_dtype=tensor._fp8_dtype, data_transpose=out_transpose, quantizer=tensor._quantizer, + device=tensor.device, ) if func in (aten.slice.Tensor, aten.select.int): @@ -687,6 +691,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): fp8_scale_inv=scale_inv, data_transpose=func_transposed_out, quantizer=quantizer, + device=tensor.device, ) return out_tensor @@ -860,6 +865,7 @@ def fsdp_post_all_gather( "quantizer": self._quantizer, "requires_grad": False, "data": data, + "device": data.device, } out = Float8Tensor(**fp8_args) @@ -898,6 +904,20 @@ def is_cpu(self): return self._transpose.is_cpu raise RuntimeError("Both data and transpose are None") + def __reduce_ex__(self, protocol: int) -> tuple: + """Custom pickling to remove references to FP8 metadata objects. + + Always serializes the underlying FP8 buffers (no dequantization + fallback for CPU tensors) so that DCP async-staging round-trips + preserve bitwise-identical data. ``Float8Tensor`` is registered + with ``torch.serialization.add_safe_globals`` to keep + ``torch.load(weights_only=True)`` compatibility. + """ + return ( + _make_float8_tensor_in_reduce_ex, + (self._data, self._fp8_dtype, self._scale_inv, self.dtype, self.shape), + ) + @classmethod def _make_in_reduce_ex( cls, @@ -905,37 +925,20 @@ def _make_in_reduce_ex( fp8_dtype: TE_DType, fp8_scale_inv: torch.Tensor, dtype: torch.dtype, - shape: torch.shape, + shape: torch.Size, ) -> Float8Tensor: - """Build Float8Tensor, for use in __reduce__ - - __reduce_ex__ assumes object constructor has positional - arguments. - + """This classmethod is kept for backward compatibility only. + ``__reduce_ex__`` used to point at this classmethod, but bound + classmethods pickle as ``(getattr, (cls, name))`` which adds an + extra reduction step to the pickle stream. The current + ``__reduce_ex__`` references the module-level + ``_make_float8_tensor_in_reduce_ex`` instead so the pickle stream + uses a single ``GLOBAL`` opcode. This classmethod is retained so + that previously pickled ``Float8Tensor`` payloads (which still + reference ``Float8Tensor._make_in_reduce_ex``) can still be + unpickled. """ - return Float8Tensor( - data=data, - fp8_dtype=fp8_dtype, - fp8_scale_inv=fp8_scale_inv, - dtype=dtype, - shape=shape, - ) - - def __reduce_ex__(self, protocol: int) -> tuple: - """Custom pickling to remove references to FP8 metadata objects - - CPU Float8Tensors are serialized as dequantized plain tensors - for compatibility with torch.load(weights_only=True), which is - used by DCP async save staging. - """ - data_is_cpu = self._data is not None and self._data.is_cpu - transpose_is_cpu = self._transpose is not None and self._transpose.is_cpu - if data_is_cpu or transpose_is_cpu: - return self.dequantize(dtype=self.dtype).__reduce_ex__(protocol) - return ( - Float8Tensor._make_in_reduce_ex, - (self._data, self._fp8_dtype, self._scale_inv, self.dtype, self.shape), - ) + return _make_float8_tensor_in_reduce_ex(data, fp8_dtype, fp8_scale_inv, dtype, shape) def _get_data(self) -> Float8Tensor: """Get tensor data property""" @@ -1000,6 +1003,24 @@ def _set_data(self, tensor: torch.Tensor) -> None: data = property(_get_data, _set_data) +def _make_float8_tensor_in_reduce_ex( + data: torch.Tensor, + fp8_dtype: TE_DType, + fp8_scale_inv: torch.Tensor, + dtype: torch.dtype, + shape: torch.Size, +) -> Float8Tensor: + """Reconstruct a ``Float8Tensor`` from its ``__reduce_ex__`` payload.""" + return Float8Tensor( + data=data, + fp8_dtype=fp8_dtype, + fp8_scale_inv=fp8_scale_inv, + dtype=dtype, + shape=shape, + device=data.device if data is not None else None, + ) + + class _ViewFunc(torch.autograd.Function): """View function @@ -1036,6 +1057,7 @@ def forward( fp8_dtype=tensor._fp8_dtype, data_transpose=out_transpose, quantizer=tensor._quantizer, + device=tensor.device, ) @staticmethod @@ -1083,6 +1105,7 @@ def forward( fp8_dtype=tensor._fp8_dtype, data_transpose=out_transpose, quantizer=tensor._quantizer, + device=tensor.device, ) @staticmethod diff --git a/transformer_engine/pytorch/tensor/mxfp8_tensor.py b/transformer_engine/pytorch/tensor/mxfp8_tensor.py index 134f8b5a6..2815aaa96 100644 --- a/transformer_engine/pytorch/tensor/mxfp8_tensor.py +++ b/transformer_engine/pytorch/tensor/mxfp8_tensor.py @@ -161,6 +161,7 @@ def create_tensor_from_data( fp8_dtype=fp8_dtype, quantizer=self, with_gemm_swizzled_scales=False, + device=data.device, ) def onnx_quantize(self, tensor: torch.Tensor) -> QuantizedTensor: @@ -346,6 +347,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): requires_grad=False, fp8_dtype=tensor._fp8_dtype, with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, + device=tensor.device, ) if func == torch.ops.aten.copy_.default: @@ -452,6 +454,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): requires_grad=False, fp8_dtype=tensor._fp8_dtype, with_gemm_swizzled_scales=False, + device=tensor.device, ) for splitted_tensor_data in zip(*out_data) ] @@ -541,6 +544,7 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): requires_grad=False, fp8_dtype=tensor._fp8_dtype, with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, + device=tensor.device, ) # Default case @@ -692,45 +696,17 @@ def fsdp_post_all_gather( shape=(rowwise_data.shape if rowwise_data is not None else columnwise_data.shape), quantizer=self._quantizer, with_gemm_swizzled_scales=False, + device=( + rowwise_data.device if rowwise_data is not None else columnwise_data.device + ), ) out._quantizer.set_usage(rowwise=rowwise_usage, columnwise=columnwise_usage) return out, all_gather_outputs - @classmethod - def _make_in_reduce_ex( - cls, - rowwise_data: torch.Tensor, - rowwise_scale_inv: torch.Tensor, - columnwise_data: torch.Tensor, - columnwise_scale_inv: torch.Tensor, - fp8_dtype: TE_DType, - dtype: torch.dtype, - shape: torch.shape, - quantizer: Optional[Quantizer] = None, - with_gemm_swizzled_scales: bool = False, - ) -> MXFP8Tensor: - """Build MXFP8Tensor, for use in __reduce__ - - __reduce_ex__ assumes object constructor has positional - arguments. - - """ - return MXFP8Tensor( - rowwise_data=rowwise_data, - rowwise_scale_inv=rowwise_scale_inv, - fp8_dtype=fp8_dtype, - columnwise_data=columnwise_data, - columnwise_scale_inv=columnwise_scale_inv, - dtype=dtype, - shape=shape, - quantizer=quantizer, - with_gemm_swizzled_scales=with_gemm_swizzled_scales, - ) - def __reduce_ex__(self, protocol: int) -> tuple: """Custom pickling""" return ( - MXFP8Tensor._make_in_reduce_ex, + _make_mxfp8_tensor_in_reduce_ex, ( self._rowwise_data, self._rowwise_scale_inv, @@ -744,6 +720,42 @@ def __reduce_ex__(self, protocol: int) -> tuple: ), ) + @classmethod + def _make_in_reduce_ex( + cls, + rowwise_data: torch.Tensor, + rowwise_scale_inv: torch.Tensor, + columnwise_data: torch.Tensor, + columnwise_scale_inv: torch.Tensor, + fp8_dtype: TE_DType, + dtype: torch.dtype, + shape: torch.Size, + quantizer: Optional[Quantizer] = None, + with_gemm_swizzled_scales: bool = False, + ) -> MXFP8Tensor: + """This classmethod is kept for backward compatibility only. + ``__reduce_ex__`` used to point at this classmethod, but bound + classmethods pickle as ``(getattr, (cls, name))`` which adds an + extra reduction step to the pickle stream. The current + ``__reduce_ex__`` references the module-level + ``_make_mxfp8_tensor_in_reduce_ex`` instead so the pickle stream + uses a single ``GLOBAL`` opcode. This classmethod is retained so + that previously pickled ``MXFP8Tensor`` payloads (which still + reference ``MXFP8Tensor._make_in_reduce_ex``) can still be + unpickled. + """ + return _make_mxfp8_tensor_in_reduce_ex( + rowwise_data, + rowwise_scale_inv, + columnwise_data, + columnwise_scale_inv, + fp8_dtype, + dtype, + shape, + quantizer, + with_gemm_swizzled_scales, + ) + def _get_data(self) -> MXFP8Tensor: """Get tensor data property""" return super().data @@ -832,6 +844,40 @@ def is_cuda(self): raise RuntimeError("MXFP8Tensor has no data!") +def _make_mxfp8_tensor_in_reduce_ex( + rowwise_data: torch.Tensor, + rowwise_scale_inv: torch.Tensor, + columnwise_data: torch.Tensor, + columnwise_scale_inv: torch.Tensor, + fp8_dtype: TE_DType, + dtype: torch.dtype, + shape: torch.Size, + quantizer: Optional[Quantizer] = None, + with_gemm_swizzled_scales: bool = False, +) -> MXFP8Tensor: + """Reconstruct an ``MXFP8Tensor`` from its ``__reduce_ex__`` payload.""" + # Infer device from inner buffers so the wrapper subclass stays + # consistent with its data (CPU after DCP staging deserialize, + # CUDA after the usual quantize path). + device = None + if rowwise_data is not None: + device = rowwise_data.device + elif columnwise_data is not None: + device = columnwise_data.device + return MXFP8Tensor( + rowwise_data=rowwise_data, + rowwise_scale_inv=rowwise_scale_inv, + fp8_dtype=fp8_dtype, + columnwise_data=columnwise_data, + columnwise_scale_inv=columnwise_scale_inv, + dtype=dtype, + shape=shape, + quantizer=quantizer, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + device=device, + ) + + class _ViewFunc(torch.autograd.Function): """View function @@ -891,6 +937,7 @@ def forward( fp8_dtype=tensor._fp8_dtype, quantizer=tensor._quantizer, with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, + device=tensor.device, ) @staticmethod @@ -918,6 +965,7 @@ def backward( fp8_dtype=grad._fp8_dtype, quantizer=grad._quantizer, with_gemm_swizzled_scales=grad._with_gemm_swizzled_scales, + device=grad.device, ) return dgrad, None return grad.view(ctx.shape), None @@ -979,6 +1027,7 @@ def forward( fp8_dtype=tensor._fp8_dtype, quantizer=tensor._quantizer, with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, + device=tensor.device, ) @staticmethod @@ -1004,6 +1053,7 @@ def backward( columnwise_scale_inv=grad._columnwise_scale_inv, fp8_dtype=grad._fp8_dtype, quantizer=grad._quantizer, + device=grad.device, with_gemm_swizzled_scales=grad._with_gemm_swizzled_scales, ) return dgrad, None diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index df7a2b4bd..2ebefefaa 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -584,6 +584,7 @@ def fsdp_post_all_gather( quantizer=self._quantizer, requires_grad=False, with_gemm_swizzled_scales=False, + device=rowwise_data.device, ) # Derive columnwise data locally via transpose instead of all-gathering it @@ -722,51 +723,16 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): quantizer=tensor._quantizer, requires_grad=tensor.requires_grad, with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, + device=tensor.device, ) # Default case return super().__torch_dispatch__(func, types, args, kwargs) - @classmethod - def _make_in_reduce_ex( - cls, - shape: torch.Size, - rowwise_data: torch.Tensor, - rowwise_scale_inv: torch.Tensor, - columnwise_data: torch.Tensor, - columnwise_scale_inv: torch.Tensor, - amax_rowwise: torch.Tensor, - amax_columnwise: torch.Tensor, - fp4_dtype: TE_DType, - dtype: torch.dtype, - quantizer: Quantizer, - with_gemm_swizzled_scales: bool = False, - ) -> NVFP4Tensor: - """Build NVFP4Tensor, for use in __reduce__ - - __reduce_ex__ assumes object constructor has positional - arguments. - - """ - return NVFP4Tensor( - shape=shape, - dtype=dtype, - fp4_dtype=fp4_dtype, - rowwise_data=rowwise_data, - rowwise_scale_inv=rowwise_scale_inv, - columnwise_data=columnwise_data, - columnwise_scale_inv=columnwise_scale_inv, - amax_rowwise=amax_rowwise, - amax_columnwise=amax_columnwise, - quantizer=quantizer, - requires_grad=False, - with_gemm_swizzled_scales=with_gemm_swizzled_scales, - ) - def __reduce_ex__(self, protocol: int) -> tuple: """Custom pickling""" return ( - NVFP4Tensor._make_in_reduce_ex, + _make_nvfp4_tensor_in_reduce_ex, ( self.shape, self._rowwise_data, @@ -782,6 +748,46 @@ def __reduce_ex__(self, protocol: int) -> tuple: ), ) + @classmethod + def _make_in_reduce_ex( + cls, + shape: torch.Size, + rowwise_data: torch.Tensor, + rowwise_scale_inv: torch.Tensor, + columnwise_data: torch.Tensor, + columnwise_scale_inv: torch.Tensor, + amax_rowwise: torch.Tensor, + amax_columnwise: torch.Tensor, + fp4_dtype: TE_DType, + dtype: torch.dtype, + quantizer: Quantizer, + with_gemm_swizzled_scales: bool = False, + ) -> NVFP4Tensor: + """This classmethod is kept for backward compatibility only. + ``__reduce_ex__`` used to point at this classmethod, but bound + classmethods pickle as ``(getattr, (cls, name))`` which adds an + extra reduction step to the pickle stream. The current + ``__reduce_ex__`` references the module-level + ``_make_nvfp4_tensor_in_reduce_ex`` instead so the pickle stream + uses a single ``GLOBAL`` opcode. This classmethod is retained so + that previously pickled ``NVFP4Tensor`` payloads (which still + reference ``NVFP4Tensor._make_in_reduce_ex``) can still be + unpickled. + """ + return _make_nvfp4_tensor_in_reduce_ex( + shape, + rowwise_data, + rowwise_scale_inv, + columnwise_data, + columnwise_scale_inv, + amax_rowwise, + amax_columnwise, + fp4_dtype, + dtype, + quantizer, + with_gemm_swizzled_scales, + ) + def _get_data(self) -> NVFP4Tensor: """Get tensor data property""" return super().data @@ -872,6 +878,45 @@ def is_cuda(self): raise RuntimeError("NVFP4Tensor has no data!") +def _make_nvfp4_tensor_in_reduce_ex( + shape: torch.Size, + rowwise_data: torch.Tensor, + rowwise_scale_inv: torch.Tensor, + columnwise_data: torch.Tensor, + columnwise_scale_inv: torch.Tensor, + amax_rowwise: torch.Tensor, + amax_columnwise: torch.Tensor, + fp4_dtype: TE_DType, + dtype: torch.dtype, + quantizer: Quantizer, + with_gemm_swizzled_scales: bool = False, +) -> NVFP4Tensor: + """Reconstruct an ``NVFP4Tensor`` from its ``__reduce_ex__`` payload.""" + # Infer device from whichever inner buffer is populated so the wrapper + # subclass stays consistent with its data buffers (e.g. CPU after DCP + # async-staging deserialize, CUDA after the usual quantize path). + device = None + if rowwise_data is not None: + device = rowwise_data.device + elif columnwise_data is not None: + device = columnwise_data.device + return NVFP4Tensor( + shape=shape, + dtype=dtype, + fp4_dtype=fp4_dtype, + rowwise_data=rowwise_data, + rowwise_scale_inv=rowwise_scale_inv, + columnwise_data=columnwise_data, + columnwise_scale_inv=columnwise_scale_inv, + amax_rowwise=amax_rowwise, + amax_columnwise=amax_columnwise, + quantizer=quantizer, + requires_grad=False, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + device=device, + ) + + class _ViewFunc(torch.autograd.Function): """View function @@ -951,6 +996,7 @@ def forward( fp4_dtype=tensor._fp4_dtype, requires_grad=tensor.requires_grad, with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, + device=tensor.device, ) @staticmethod @@ -993,6 +1039,7 @@ def backward( fp4_dtype=grad._fp4_dtype, requires_grad=grad.requires_grad, with_gemm_swizzled_scales=grad._with_gemm_swizzled_scales, + device=grad.device, ) return dgrad, None return grad.view(ctx.shape), None @@ -1077,6 +1124,7 @@ def forward( fp4_dtype=tensor._fp4_dtype, requires_grad=tensor.requires_grad, with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, + device=tensor.device, ) @staticmethod @@ -1119,6 +1167,7 @@ def backward( fp4_dtype=grad._fp4_dtype, requires_grad=grad.requires_grad, with_gemm_swizzled_scales=grad._with_gemm_swizzled_scales, + device=grad.device, ) return dgrad, None return grad.view(ctx.shape), None diff --git a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py index de7f8f58e..3a72ec5d1 100644 --- a/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/float8_tensor_storage.py @@ -139,11 +139,6 @@ def get_metadata(self) -> Dict[str, Any]: "fp8_dtype": self._fp8_dtype, "data_transpose": self._transpose, "quantizer": self._quantizer, - "device": ( - self._data.device - if self._data is not None - else (self._transpose.device if self._transpose is not None else None) - ), "fake_dtype": self._dtype, } diff --git a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py index 842f42838..874555f46 100644 --- a/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/mxfp8_tensor_storage.py @@ -35,12 +35,16 @@ def forward( if tensor._columnwise_data is not None and tensor._columnwise_data.numel() == 0: return torch.empty(tensor.size(), dtype=dtype, device=tensor.device) - dtype = torch_to_transformer_engine_dtype[dtype] - - # Make sure FP8 data is in expected format - if tensor._rowwise_data is not None or tensor._columnwise_data is not None: - return tex.dequantize(tensor, dtype) - raise ValueError("Cannot dequantize MXFP8 tensor with no data") + if tensor._rowwise_data is None and tensor._columnwise_data is None: + raise ValueError("Cannot dequantize MXFP8 tensor with no data") + te_dtype = torch_to_transformer_engine_dtype[dtype] + # ``tex.dequantize`` requires CUDA-resident buffers. + src_device = tensor.device + if src_device.type != "cuda": + cuda_tensor = tensor.to(device=torch.device("cuda")) + result = tex.dequantize(cuda_tensor, te_dtype) + return result.to(device=src_device) + return tex.dequantize(tensor, te_dtype) @staticmethod def backward( diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index e51acb71e..490184e5f 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -47,13 +47,18 @@ def forward( if tensor._columnwise_data is not None and tensor._columnwise_data.numel() == 0: return torch.empty(tensor.size(), dtype=dtype, device=tensor.device) - # Dequantize row-wise data - if tensor._rowwise_data is not None: - return tex.dequantize(tensor, torch_to_transformer_engine_dtype[dtype]) - - if tensor._columnwise_data is not None: + if tensor._rowwise_data is None and tensor._columnwise_data is None: + raise ValueError("Attempted to dequantize NVFP4 tensor with no data") + if tensor._rowwise_data is None and tensor._columnwise_data is not None: raise NotImplementedError("Dequantizing column-wise NVFP4 data is not implemented yet!") - raise ValueError("Attempted to dequantize NVFP4 tensor with no data") + + # ``tex.dequantize`` requires CUDA-resident buffers. If the tensor has + src_device = tensor.device + if src_device.type != "cuda": + cuda_tensor = tensor.to(device=torch.device("cuda")) + result = tex.dequantize(cuda_tensor, torch_to_transformer_engine_dtype[dtype]) + return result.to(device=src_device) + return tex.dequantize(tensor, torch_to_transformer_engine_dtype[dtype]) @staticmethod def backward( From dc9af4abc6bad7b81d01ece364c01db9ff8a0e65 Mon Sep 17 00:00:00 2001 From: Ziang Li Date: Fri, 22 May 2026 15:58:11 -0700 Subject: [PATCH 079/180] Implement 4over6 NVFP4 recipe (#2972) * Initial implementation Signed-off-by: Ziang Li * Make 4over6 compile time for dequant Signed-off-by: Ziang Li * Expand 1d fwd+bwd test Signed-off-by: Ziang Li * Refactor Signed-off-by: Ziang Li * Clean up Signed-off-by: Ziang Li * Clean up Signed-off-by: Ziang Li * Add gemm test Signed-off-by: Ziang Li * Add more tests and fix offload Signed-off-by: Ziang Li * Fix offload Signed-off-by: Ziang Li * Clean up arg Signed-off-by: Ziang Li * Add more test Signed-off-by: Ziang Li * Add more tests Signed-off-by: Ziang Li * Clean up test Signed-off-by: Ziang Li * Refactor cuh kernel impl Signed-off-by: Ziang Li * Further extract Signed-off-by: Ziang Li * Clean up Signed-off-by: Ziang Li * Add recipe_id Signed-off-by: Ziang Li * Fix failing unit tests Signed-off-by: Ziang Li * Clean up test Signed-off-by: Ziang Li * Clean up Signed-off-by: Ziang Li * Refactor ref Signed-off-by: Ziang Li * Update comments and docs Signed-off-by: Ziang Li * Drop unnecessary test_sanity workaround The following tests passed: `NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 python3 -m pytest --tb=auto tests/pytorch/test_sanity.py ` `NVTE_GROUPED_LINEAR_SINGLE_PARAM=1 NVTE_TEST_NVINSPECT_ENABLED=1 NVTE_TEST_NVINSPECT_CONFIG_FILE=tests/pytorch/debug/test_configs/dummy_feature.yaml NVTE_TEST_NVINSPECT_FEATURE_DIRS=transformer_engine/debug/features PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest --tb=auto tests/pytorch/test_sanity.py ` Signed-off-by: Ziang Li * Refactor `QuantizerRole` Signed-off-by: Ziang Li * Allow separate recipe 4over6 config Signed-off-by: Ziang Li * Support 2d Signed-off-by: Ziang Li * Refactor 2d Signed-off-by: Ziang Li * Clean up anti pattern Signed-off-by: Ziang Li * Enforce 4over6 consistency Signed-off-by: Ziang Li * Update comments Signed-off-by: Ziang Li * Update docs Signed-off-by: Ziang Li * Fix test Signed-off-by: Ziang Li * Drop test_fusible_ops Signed-off-by: Ziang Li * Revert "Drop test_fusible_ops" This reverts commit 69f9ccc36a9c459f50c2f00b6cd6a62c5e1bdf13. Signed-off-by: Ziang Li * Refactor test_fusible_ops Signed-off-by: Ziang Li * Refactor ref and extend cpp test Signed-off-by: Ziang Li * Clean up cpp test Signed-off-by: Ziang Li * Minor comment Signed-off-by: Ziang Li * Drop doc Signed-off-by: Ziang Li * Explicit handle conditional smem buffer Signed-off-by: Ziang Li * Further clean up Signed-off-by: Ziang Li * More templates Signed-off-by: Ziang Li * Simplify cpp Signed-off-by: Ziang Li * Drop write back lifting Signed-off-by: Ziang Li * Add MAE and dedicated fast math env var Signed-off-by: Ziang Li * Harden cpp test Signed-off-by: Ziang Li * Add warning and err fast math coverage Signed-off-by: Ziang Li * Fold test case and clean up cpp test Signed-off-by: Ziang Li * Initial 448 vs 256 implementation Signed-off-by: Ziang Li * Use e4m3 max instead of boolean, more template Signed-off-by: Ziang Li * Add benchmark script and minor optimization Signed-off-by: Ziang Li * Use standalone kernels Signed-off-by: Ziang Li * Use cp async Signed-off-by: Ziang Li * Add benchmark script Signed-off-by: Ziang Li * Minor fix after rebase Signed-off-by: Ziang Li * Naming consistency Signed-off-by: Ziang Li * Remove 4over6 benchmark Signed-off-by: Ziang Li * Refactor modes Signed-off-by: Ziang Li * Relax tol for `test_layernorm_mlp` for `nvfp4_4over6` Signed-off-by: Ziang Li * Minor fix recipe naming Signed-off-by: Ziang Li * Remove gradient 4over6 quantization and partially allow SR/RHT Signed-off-by: Ziang Li * Allow RHT in pytorch ref Signed-off-by: Ziang Li * Update transformer_engine/pytorch/csrc/quantizer.cpp Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> * Minor fix TODO lint Signed-off-by: Ziang Li * Use standard nvfp4 for grad ref in test_fusible_ops.py since 4over6 is not applied to gradient quantizers Signed-off-by: Ziang Li * Minor fix test-fusible_ops 4over6 helper Signed-off-by: Ziang Li * Default to 256 for 4over6 Signed-off-by: Ziang Li * Reset RNG state for each TE ops test Adding tests affected RNG in unrelated tests. Signed-off-by: Tim Moon * Remove loosened NVFP4 tols in layernorm MLP test. Make sure tensors are representable in quantized format. Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Ziang Li Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Signed-off-by: Tim Moon Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: Tim Moon Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- docs/envvars.rst | 24 + .../cpp/operator/test_cast_nvfp4_transpose.cu | 616 +++++++++++++--- tests/cpp/operator/test_dequantize_nvfp4.cu | 68 +- tests/cpp/test_common.cu | 12 + tests/cpp/test_common.h | 3 + tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py | 52 ++ .../nvfp4/test_nvfp4_quantize_exact.py | 65 +- tests/pytorch/test_backward_override.py | 21 +- tests/pytorch/test_cpu_offloading.py | 91 ++- tests/pytorch/test_cuda_graphs.py | 21 +- tests/pytorch/test_fusible_ops.py | 100 ++- tests/pytorch/test_numerics.py | 72 +- tests/pytorch/test_quantized_tensor.py | 21 +- tests/pytorch/test_recipe.py | 134 +++- tests/pytorch/test_sanity.py | 64 +- tests/pytorch/test_torch_compile.py | 43 +- tests/pytorch/utils.py | 39 +- .../common/cast/dispatch/quantize.cuh | 41 +- .../common/cast/nvfp4/core_nvfp4.cuh | 8 +- .../common/cast/nvfp4/dequantize_nvfp4.cuh | 40 +- .../cast/nvfp4/quantize_4over6_nvfp4.cuh | 668 ++++++++++++++++++ .../comm_gemm_overlap/comm_gemm_overlap.cpp | 4 + transformer_engine/common/common.h | 16 +- .../transformer_engine/transformer_engine.h | 55 ++ transformer_engine/common/recipe/__init__.py | 30 + transformer_engine/common/recipe/nvfp4.cu | 13 +- .../common/transformer_engine.cpp | 28 + transformer_engine/pytorch/csrc/common.h | 4 + .../pytorch/csrc/extensions/cast.cpp | 70 +- transformer_engine/pytorch/csrc/quantizer.cpp | 49 +- .../pytorch/csrc/type_converters.cpp | 2 + .../custom_recipes/quantization_ref_nvfp4.py | 190 ++++- transformer_engine/pytorch/quantization.py | 39 +- .../pytorch/tensor/grouped_tensor.py | 6 + .../pytorch/tensor/nvfp4_tensor.py | 72 +- .../tensor/storage/grouped_tensor_storage.py | 49 ++ .../tensor/storage/nvfp4_tensor_storage.py | 16 + 37 files changed, 2595 insertions(+), 251 deletions(-) create mode 100644 transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh diff --git a/docs/envvars.rst b/docs/envvars.rst index ffbad409d..bd62ccac4 100644 --- a/docs/envvars.rst +++ b/docs/envvars.rst @@ -287,6 +287,30 @@ Kernel Configuration :Default: ``0`` :Description: Enable row-scaled NVFP4 tensors for forward activation quantizers in the ``NVFP4BlockScaling`` recipe. When set to ``1`` (or when ``NVFP4BlockScaling(row_scaled_activation=True)`` is used), rowwise ``amax`` metadata is stored as one FP32 value per tensor row instead of a single scalar. +.. envvar:: NVTE_NVFP4_4OVER6 + + :Type: ``str`` (``none``, ``weights``, ``activations``, or ``all``) + :Default: ``none`` + :Description: Enable 4over6 adaptive NVFP4 block scaling for weights, activations, or both in the ``NVFP4BlockScaling`` recipe. For each selected FP4 block, quantization compares map-to-4 and map-to-6 candidates and stores the candidate with lower configured error. ``none`` keeps standard NVFP4. Current 4over6 support targets RL and post-training scenarios; pre-training paths that combine 4over6 with RHT are not yet implemented. + +.. envvar:: NVTE_NVFP4_4OVER6_E4M3_USE_256 + + :Type: ``str`` (``none``, ``weights``, ``activations``, or ``all``) + :Default: ``all`` + :Description: Select NVFP4 4over6 quantizers that use 256 instead of 448 as the global E4M3 scale bound. By default, all 4over6 quantizers use 256. Set the env var to ``none`` (or set ``NVFP4BlockScaling(nvfp4_4over6_e4m3_use_256="none")``) to use the standard NVFP4 448 bound for all 4over6 quantizers. This option is only meaningful for tensor roles that also enable :envvar:`NVTE_NVFP4_4OVER6`. + +.. envvar:: NVTE_NVFP4_4OVER6_ERR_MODE + + :Type: ``str`` (``MAE`` or ``MSE``) + :Default: ``MAE`` + :Description: Select the input-domain error metric used by NVFP4 4over6 map-to-4 versus map-to-6 candidate selection in the ``NVFP4BlockScaling`` recipe. + +.. envvar:: NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH + + :Type: ``int`` (0 or 1) + :Default: ``0`` + :Description: Allow the NVFP4 4over6 candidate error computation to use faster non-strict floating-point expressions. By default, 4over6 error comparison uses strict expressions; ``NVTE_USE_FAST_MATH`` does not control this error-comparison path. + Torch Compilation and Fusion ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu index a8f58f859..d6ab4b674 100644 --- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu +++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu @@ -62,12 +62,14 @@ std::vector create_transpose(const InputType* const input, const size } // Compute the global encode scale factor for a given global amax -float compute_global_encode_scaling_factor_FP4(const float global_amax, const bool use_fast_math) { - constexpr float fp8_max = 448.0f; // 448.0f; +float compute_global_encode_scaling_factor_FP4(const float global_amax, const bool use_fast_math, + const int e4m3_max = 448) { + NVTE_CHECK(e4m3_max == 448 || e4m3_max == 256, "Unsupported NVFP4 E4M3 max."); + const float fp8_max = static_cast(e4m3_max); constexpr float fp4_max = 6.0f; // 6.0f; float global_encode_scale = fp8_max * fp4_max / global_amax; // If scale is infinity, return the max normalized value - const float max_norm_clamp = use_fast_math + const float max_norm_clamp = (use_fast_math && e4m3_max == 448) ? Numeric_Traits::maxNorm : Numeric_Traits::maxNorm; @@ -79,6 +81,103 @@ float compute_global_encode_scaling_factor_FP4(const float global_amax, const bo return global_encode_scale; } +struct NVFP4FourOverSixQuantization { + fp8e4m3 scale_map4; + fp8e4m3 scale_map6; + float reciprocal_map4; + float reciprocal_map6; + fp4e2m1x2 quantized_map4; + fp4e2m1x2 quantized_map6; +}; + +enum class NVFP4FourOverSixCandidate { + Map4, + Map6, +}; + +enum class NVFP4ScalingMode { + Block1D, + RowScaled1D, + Block2D, +}; + +struct NVFP4FourOverSixTestConfig { + NVTENVFP44Over6Mode mode = kNVTENVFP44Over6Disabled; + int e4m3_max = 448; + bool err_use_fast_math = false; +}; + +bool use_2d_quantization(const NVFP4ScalingMode scaling_mode) { + return scaling_mode == NVFP4ScalingMode::Block2D; +} + +NVFP4FourOverSixQuantization compute_4over6_quantization_scales( + const float block_amax, const float global_encode_scale) { + constexpr float fp4_max = 6.0f; + constexpr float fp8_max = 448.0f; + constexpr float scale_expansion_factor = 1.5f; + const float base_sf_high_precision = block_amax / fp4_max * global_encode_scale; + const float sf_high_precision_map4 = + fminf(base_sf_high_precision * scale_expansion_factor, fp8_max); + const float sf_high_precision_map6 = fminf(base_sf_high_precision, fp8_max); + const fp8e4m3 scale_map4 = static_cast(sf_high_precision_map4); + const fp8e4m3 scale_map6 = static_cast(sf_high_precision_map6); + + const float global_decode_scale = 1.0f / global_encode_scale; + const float scale_map4_fp32 = static_cast(scale_map4); + const float reciprocal_map4 = + fminf(1.0f / (scale_map4_fp32 * global_decode_scale), Numeric_Traits::maxNorm); + const float scale_map6_fp32 = static_cast(scale_map6); + const float reciprocal_map6 = + fminf(1.0f / (scale_map6_fp32 * global_decode_scale), Numeric_Traits::maxNorm); + + const float2 zero = {0.0f, 0.0f}; + return { + scale_map4, + scale_map6, + reciprocal_map4, + reciprocal_map6, + fp4e2m1x2(zero), + fp4e2m1x2(zero), + }; +} + +fp8e4m3 select_4over6_scale(const NVFP4FourOverSixQuantization& quantization, + const NVFP4FourOverSixCandidate candidate) { + if (candidate == NVFP4FourOverSixCandidate::Map4) { + return quantization.scale_map4; + } + return quantization.scale_map6; +} + +fp4e2m1x2 select_4over6_quantized_pair(const NVFP4FourOverSixQuantization& quantization, + const NVFP4FourOverSixCandidate candidate) { + if (candidate == NVFP4FourOverSixCandidate::Map4) { + return quantization.quantized_map4; + } + return quantization.quantized_map6; +} + +NVFP4FourOverSixQuantization quantize_4over6_pair( + const float x, const float y, const NVFP4FourOverSixQuantization& quantization) { + const float2 scaled_map4 = {x * quantization.reciprocal_map4, + y * quantization.reciprocal_map4}; + const fp4e2m1x2 quantized_map4(scaled_map4); + + const float2 scaled_map6 = {x * quantization.reciprocal_map6, + y * quantization.reciprocal_map6}; + const fp4e2m1x2 quantized_map6(scaled_map6); + + return { + quantization.scale_map4, + quantization.scale_map6, + quantization.reciprocal_map4, + quantization.reciprocal_map6, + quantized_map4, + quantized_map6, + }; +} + // 1D Scaling: Original implementation with 1x16 blocks template void quantize_nvfp4_1d(float (*OP)(const float), @@ -89,10 +188,15 @@ void quantize_nvfp4_1d(float (*OP)(const float), const size_t cols, const size_t scales_stride, const float global_amax, - const bool use_fast_math) { + const bool use_fast_math, + const bool use_4over6 = false, + const int e4m3_max = 448, + const NVFP4FourOverSixCandidate four_over_six_candidate = + NVFP4FourOverSixCandidate::Map6) { // Compute a global encoding/decoding scaling factor for all S_dec_b - const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax, use_fast_math); + const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax, use_fast_math, + e4m3_max); constexpr size_t block_size_X = 16; const size_t blocks_X = divide_round_up(cols, block_size_X); @@ -122,6 +226,27 @@ void quantize_nvfp4_1d(float (*OP)(const float), block_amax = std::max(block_amax, std::abs(elt)); } + const size_t scale_idx = i * scales_stride + block_X; + + if (use_4over6) { + const NVFP4FourOverSixQuantization quantization = + compute_4over6_quantization_scales(block_amax, S_enc); + scales[scale_idx] = select_4over6_scale(quantization, four_over_six_candidate); + + for (size_t j = j_min; j < j_max; j += 2) { + const int idx_pair = (i * cols + j) / 2; + const int cache_idx_x = j - j_min; + const int cache_idx_y = cache_idx_x + 1; + const float cached_x = cache_buffer[cache_idx_x]; + const float cached_y = cache_buffer[cache_idx_y]; + const NVFP4FourOverSixQuantization pair_quantization = + quantize_4over6_pair(cached_x, cached_y, quantization); + output[idx_pair] = + select_4over6_quantized_pair(pair_quantization, four_over_six_candidate); + } + continue; + } + // Compute and store the per-block FP8 decode scale const float S_dec_b = block_amax * (S_enc * (1.0f / 6.0f)); const fp8e4m3 S_dec_b_fp8 = static_cast(fminf(S_dec_b, Numeric_Traits::maxNorm)); @@ -131,7 +256,6 @@ void quantize_nvfp4_1d(float (*OP)(const float), const float S_enc_b_fp8 = S_dec_b_fp32 == 0.f ? 0.f : fminf(1.0f / (S_dec_b_fp32 * (1.0f / S_enc)), Numeric_Traits::maxNorm); - const size_t scale_idx = i * scales_stride + block_X; scales[scale_idx] = S_dec_b_fp8; float scale_reciprocal = S_enc_b_fp8; @@ -167,9 +291,14 @@ void compute_2d_mathematical_scales(float (*OP)(const float), const size_t cols, const float global_amax, std::vector>& math_scales, - const bool use_fast_math) { - - const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax, use_fast_math); + const bool use_fast_math, + const bool use_4over6 = false, + const int e4m3_max = 448, + const NVFP4FourOverSixCandidate four_over_six_candidate = + NVFP4FourOverSixCandidate::Map6) { + + const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax, use_fast_math, + e4m3_max); constexpr size_t block_size_Y = 16; constexpr size_t block_size_X = 16; const size_t blocks_Y = divide_round_up(rows, block_size_Y); @@ -197,9 +326,16 @@ void compute_2d_mathematical_scales(float (*OP)(const float), } // Compute E4M3 scaling factor for this 16x16 block - const float S_dec_b = block_amax / 6.0f; - const fp8e4m3 S_dec_b_fp8 = static_cast(S_dec_b * S_enc); - math_scales[block_Y][block_X] = S_dec_b_fp8; + if (use_4over6) { + const NVFP4FourOverSixQuantization quantization = + compute_4over6_quantization_scales(block_amax, S_enc); + math_scales[block_Y][block_X] = + select_4over6_scale(quantization, four_over_six_candidate); + } else { + const float S_dec_b = block_amax / 6.0f * S_enc; + const fp8e4m3 S_dec_b_fp8_map6 = static_cast(S_dec_b); + math_scales[block_Y][block_X] = S_dec_b_fp8_map6; + } } } } @@ -214,13 +350,19 @@ void quantize_nvfp4_2d(float (*OP)(const float), const size_t cols, const size_t scales_stride, const float global_amax, - const bool use_fast_math) { + const bool use_fast_math, + const bool use_4over6 = false, + const int e4m3_max = 448, + const NVFP4FourOverSixCandidate four_over_six_candidate = + NVFP4FourOverSixCandidate::Map6) { // Step 1: Compute mathematical 8x8 scaling factors std::vector> math_scales; - compute_2d_mathematical_scales(OP, input, rows, cols, global_amax, math_scales, use_fast_math); + compute_2d_mathematical_scales(OP, input, rows, cols, global_amax, math_scales, use_fast_math, + use_4over6, e4m3_max, four_over_six_candidate); - const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax, use_fast_math); + const float S_enc = compute_global_encode_scaling_factor_FP4(global_amax, use_fast_math, + e4m3_max); constexpr size_t block_size_Y = 16; constexpr size_t block_size_X = 16; const size_t blocks_Y = divide_round_up(rows, block_size_Y); @@ -250,7 +392,7 @@ void quantize_nvfp4_2d(float (*OP)(const float), // Get the scaling factor for this block const float S_dec_b_fp8 = static_cast(math_scales[block_Y][block_X]); - const float S_enc_b_fp8 = S_dec_b_fp8 == 0 ? 0.f : S_enc / S_dec_b_fp8; + const float S_enc_b_fp8 = S_dec_b_fp8 == 0.0f ? 0.0f : S_enc / S_dec_b_fp8; const float scale_reciprocal = S_enc_b_fp8; // Process and cache data for this 16x16 block @@ -302,11 +444,17 @@ void quantize_nvfp4(float (*OP)(const float), const size_t scales_stride, const float global_amax, const bool use_fast_math, - const bool use_2d_quantization = false) { + const bool use_2d_quantization = false, + const bool use_4over6 = false, + const int e4m3_max = 448, + const NVFP4FourOverSixCandidate four_over_six_candidate = + NVFP4FourOverSixCandidate::Map6) { if (use_2d_quantization) { - quantize_nvfp4_2d(OP, input, output, scales, rows, cols, scales_stride, global_amax, use_fast_math); + quantize_nvfp4_2d(OP, input, output, scales, rows, cols, scales_stride, global_amax, + use_fast_math, use_4over6, e4m3_max, four_over_six_candidate); } else { - quantize_nvfp4_1d(OP, input, output, scales, rows, cols, scales_stride, global_amax, use_fast_math); + quantize_nvfp4_1d(OP, input, output, scales, rows, cols, scales_stride, global_amax, + use_fast_math, use_4over6, e4m3_max, four_over_six_candidate); } } @@ -324,7 +472,11 @@ void compute_ref(float (*OP)(const float), const size_t scales_stride_t, const bool use_fast_math, const bool use_2d_quantization = false, - const bool row_scaled_nvfp4 = false) + const bool row_scaled_nvfp4 = false, + const bool use_4over6 = false, + const int e4m3_max = 448, + const NVFP4FourOverSixCandidate four_over_six_candidate = + NVFP4FourOverSixCandidate::Map6) { std::vector input_t = create_transpose(input, rows, cols); NVTE_CHECK(!(use_2d_quantization && row_scaled_nvfp4), @@ -334,7 +486,8 @@ void compute_ref(float (*OP)(const float), if (use_2d_quantization) { // Step 1: Compute mathematical 8×8 scaling factors std::vector> math_scales; - compute_2d_mathematical_scales(OP, input, rows, cols, *amax, math_scales, use_fast_math); + compute_2d_mathematical_scales(OP, input, rows, cols, *amax, math_scales, use_fast_math, + use_4over6, e4m3_max, four_over_six_candidate); constexpr size_t block_size_Y = 16; constexpr size_t block_size_X = 16; @@ -362,9 +515,11 @@ void compute_ref(float (*OP)(const float), // Step 4: Process quantized outputs using the same algorithm as quantize_nvfp4_2d // (This part processes the actual FP4 data using the mathematical scaling factors) quantize_nvfp4_2d(OP, input, output, nullptr, rows, cols, scales_stride, *amax, - use_fast_math); // scales already filled + use_fast_math, use_4over6, e4m3_max, + four_over_six_candidate); // scales already filled quantize_nvfp4_2d(OP, input_t.data(), output_t, nullptr, cols, rows, scales_stride_t, *amax, - use_fast_math); // scales_t already filled + use_fast_math, use_4over6, e4m3_max, + four_over_six_candidate); // scales_t already filled return; } @@ -381,16 +536,21 @@ void compute_ref(float (*OP)(const float), scales_stride, amax[row], use_fast_math, - use_2d_quantization); + use_2d_quantization, + use_4over6, + e4m3_max, + four_over_six_candidate); } return; } // Ref impl for basic NVFP4 quantize_nvfp4(OP, input, output, scales, rows, cols, scales_stride, *amax, - use_fast_math, use_2d_quantization); + use_fast_math, use_2d_quantization, use_4over6, e4m3_max, + four_over_six_candidate); quantize_nvfp4(OP, input_t.data(), output_t, scales_t, cols, rows, scales_stride_t, *amax, - use_fast_math, use_2d_quantization); + use_fast_math, use_2d_quantization, use_4over6, e4m3_max, + four_over_six_candidate); } void compare_nvfp4_tensors(const std::string& name, @@ -515,6 +675,92 @@ void compareResults_nvfp4(Tensor &test, } } +template +bool bitwise_equal(const T& x, const T& y) { + const auto *x_bytes = reinterpret_cast(&x); + const auto *y_bytes = reinterpret_cast(&y); + for (size_t i = 0; i < sizeof(T); ++i) { + if (x_bytes[i] != y_bytes[i]) { + return false; + } + } + return true; +} + +bool nvfp4_output_block_matches(const fp4e2m1x2* const test_data, + const fp4e2m1x2* const ref_data, + const size_t row, + const size_t cols, + const size_t block_x) { + constexpr size_t block_size_X = 16; + const size_t j_min = block_x * block_size_X; + const size_t j_max = std::min(j_min + block_size_X, cols); + for (size_t j = j_min; j < j_max; j += 2) { + const size_t idx_pair = (row * cols + j) / 2; + if (!bitwise_equal(test_data[idx_pair], ref_data[idx_pair])) { + return false; + } + } + return true; +} + +void compare_nvfp4_4over6_candidates(const std::string& name, + const fp4e2m1* const test_data, + const fp8e4m3* const test_scales, + const fp4e2m1x2* const ref_data_map4, + const fp8e4m3* const ref_scales_map4, + const fp4e2m1x2* const ref_data_map6, + const fp8e4m3* const ref_scales_map6, + const size_t rows, + const size_t cols, + const size_t blocks_X, + const size_t scales_stride) { + constexpr int max_mismatches_to_print = 3; + const auto* const test_data_pairs = reinterpret_cast(test_data); + size_t total_mismatches = 0; + + for (size_t row = 0; row < rows; ++row) { + for (size_t block_x = 0; block_x < blocks_X; ++block_x) { + const size_t scale_idx = row * scales_stride + block_x; + const bool scale_matches_map4 = + bitwise_equal(test_scales[scale_idx], ref_scales_map4[scale_idx]); + const bool data_matches_map4 = + nvfp4_output_block_matches(test_data_pairs, ref_data_map4, row, cols, block_x); + const bool scale_matches_map6 = + bitwise_equal(test_scales[scale_idx], ref_scales_map6[scale_idx]); + const bool data_matches_map6 = + nvfp4_output_block_matches(test_data_pairs, ref_data_map6, row, cols, block_x); + + if ((scale_matches_map4 && data_matches_map4) || + (scale_matches_map6 && data_matches_map6)) { + continue; + } + + ++total_mismatches; + if (total_mismatches <= max_mismatches_to_print) { + std::cout << "Error in tensor " << name << ": 4over6 block mismatch at row " + << row << ", block_x " << block_x + << ". The output did not match either map-to-4 or map-to-6 exactly." + << std::endl; + } + } + } + + std::cout << "=== SUMMARY for tensor " << name << " ===" << std::endl; + std::cout << "Total 4over6 blocks checked: " << (rows * blocks_X) << std::endl; + if (total_mismatches > 0) { + std::cout << "STATUS: FAILED for output" << std::endl; + std::cout << "Total mismatched 4over6 blocks found: " << total_mismatches << std::endl; + std::cout << "============================" << std::endl; + GTEST_FAIL() << "Found " << total_mismatches << " 4over6 block mismatches in tensor " + << name; + } + + std::cout << "STATUS: PASSED for output" << std::endl; + std::cout << "Each 4over6 block matched either map-to-4 or map-to-6 exactly" << std::endl; + std::cout << "============================" << std::endl; +} + void compare_rowwise_amax(Tensor &output, const std::vector &ref_amax) { ASSERT_EQ(output.rowwise_amax_size(), ref_amax.size()); const auto *amax_ptr = output.cpu_rowwise_amax_ptr(); @@ -529,12 +775,25 @@ template void performTest(float (*OP)(const float), const std::vector& shape, const bool use_fast_math, - const bool row_scaled_nvfp4 = false) { + const NVFP4ScalingMode scaling_mode = NVFP4ScalingMode::Block1D, + const NVTENVFP44Over6Mode mode = kNVTENVFP44Over6Disabled, + const int e4m3_max = 448, + const bool use_4over6_err_use_fast_math = false) { using namespace test; + const bool use_4over6 = mode != kNVTENVFP44Over6Disabled; + + if (use_4over6 && use_fast_math) { + std::cout << "WARNING: Plain NVFP4 fast math is ignored for 4over6. " + "Use use_4over6_err_use_fast_math to test the 4over6 candidate " + "error fast-math path." + << std::endl; + } DType itype = TypeInfo::dtype; DType otype = DType::kFloat4E2M1; + const bool is_2d_quantization = use_2d_quantization(scaling_mode); + const bool row_scaled_nvfp4 = scaling_mode == NVFP4ScalingMode::RowScaled1D; const bool rowwise = true; const bool columnwise = !row_scaled_nvfp4; @@ -560,14 +819,52 @@ void performTest(float (*OP)(const float), Tensor input("input", shape, itype); Tensor output("output", shape, otype, rowwise, columnwise, NVTE_NVFP4_1D_SCALING); + output.set_nvfp4_e4m3_max(e4m3_max); std::unique_ptr ref_output = std::make_unique(rows * (cols / 2)); std::unique_ptr ref_output_t = std::make_unique(cols * (rows / 2)); std::unique_ptr ref_scales = std::make_unique(blocks_Y * blocks_X); std::unique_ptr ref_scales_t = std::make_unique(blocks_Y_t * blocks_X_t); + std::unique_ptr ref_output_map6; + std::unique_ptr ref_output_t_map6; + std::unique_ptr ref_scales_map6; + std::unique_ptr ref_scales_t_map6; fillCase(&input, InputsFillCase::uniform); + if (use_4over6 && row_scaled_nvfp4) { + const float target_row_amax = static_cast(e4m3_max) * 6.0f * 8.0f; + auto *input_vals = input.rowwise_cpu_dptr(); + for (size_t row = 0; row < rows; ++row) { + float row_amax = 0.0f; + size_t max_col = 0; + for (size_t col = 0; col < cols; ++col) { + const float val = static_cast(input_vals[row * cols + col]); + const float abs_val = fabsf(val); + if (abs_val > row_amax) { + row_amax = abs_val; + max_col = col; + } + } + + if (row_amax == 0.0f) { + continue; + } + + const float row_scale = target_row_amax / row_amax; + for (size_t col = 0; col < cols; ++col) { + float scaled = static_cast(input_vals[row * cols + col]) * row_scale; + scaled = fminf(fmaxf(scaled, -target_row_amax), target_row_amax); + input_vals[row * cols + col] = static_cast(scaled); + } + + const float max_val = static_cast(input_vals[row * cols + max_col]); + input_vals[row * cols + max_col] = + static_cast(max_val < 0.0f ? -target_row_amax : target_row_amax); + } + input.from_cpu(); + } + // Compute 2nd stage NVFP4 scaling factor std::vector ref_amax; if (row_scaled_nvfp4) { @@ -587,7 +884,11 @@ void performTest(float (*OP)(const float), output.set_row_scaled_nvfp4(row_scaled_nvfp4); } else { // Golden value of amax chosen to make the 2nd-stage scaling mantissa zero and avoid rounding issues - ref_amax.assign(1, 448.0f * 6.0f * 8.0f); + if (use_4over6) { + ref_amax.assign(1, static_cast(e4m3_max) * 6.0f * 8.0f); + } else { + ref_amax.assign(1, 448.0f * 6.0f * 8.0f); + } // Update tensor if (rowwise) { @@ -599,22 +900,63 @@ void performTest(float (*OP)(const float), output.from_cpu(); } - // Reference implementation - bool use_2d_quantization = false; - compute_ref(OP, - input.rowwise_cpu_dptr(), - ref_output.get(), - ref_output_t.get(), - ref_scales.get(), - ref_scales_t.get(), - ref_amax.data(), - rows, - cols, - scales_stride, - scales_stride_t, - use_fast_math, - use_2d_quantization, - row_scaled_nvfp4); + if (use_4over6) { + ref_output_map6 = std::make_unique(rows * (cols / 2)); + ref_output_t_map6 = std::make_unique(cols * (rows / 2)); + ref_scales_map6 = std::make_unique(blocks_Y * blocks_X); + ref_scales_t_map6 = std::make_unique(blocks_Y_t * blocks_X_t); + + compute_ref(OP, + input.rowwise_cpu_dptr(), + ref_output.get(), + ref_output_t.get(), + ref_scales.get(), + ref_scales_t.get(), + ref_amax.data(), + rows, + cols, + scales_stride, + scales_stride_t, + use_fast_math, + is_2d_quantization, + row_scaled_nvfp4, + use_4over6, + e4m3_max, + NVFP4FourOverSixCandidate::Map4); + compute_ref(OP, + input.rowwise_cpu_dptr(), + ref_output_map6.get(), + ref_output_t_map6.get(), + ref_scales_map6.get(), + ref_scales_t_map6.get(), + ref_amax.data(), + rows, + cols, + scales_stride, + scales_stride_t, + use_fast_math, + is_2d_quantization, + row_scaled_nvfp4, + use_4over6, + e4m3_max, + NVFP4FourOverSixCandidate::Map6); + } else { + compute_ref(OP, + input.rowwise_cpu_dptr(), + ref_output.get(), + ref_output_t.get(), + ref_scales.get(), + ref_scales_t.get(), + ref_amax.data(), + rows, + cols, + scales_stride, + scales_stride_t, + use_fast_math, + is_2d_quantization, + row_scaled_nvfp4, + use_4over6); + } // Initialize stochastic rounding Tensor rng_state("rng_state", std::vector{2}, DType::kInt64); @@ -624,10 +966,12 @@ void performTest(float (*OP)(const float), // Quantization options QuantizationConfigWrapper quant_config; - quant_config.set_use_fast_math(use_fast_math); + quant_config.set_use_fast_math(use_fast_math && !use_4over6); quant_config.set_stochastic_rounding(false); quant_config.set_rng_state(rng_state.data()); - quant_config.set_nvfp4_2d_quantization(use_2d_quantization); + quant_config.set_nvfp4_2d_quantization(is_2d_quantization); + quant_config.set_nvfp4_4over6_mode(mode); + quant_config.set_nvfp4_4over6_err_use_fast_math(use_4over6 && use_4over6_err_use_fast_math); // Call appropriate function based on operation type // Activation functions take 3 parameters (input, output, stream) @@ -656,21 +1000,50 @@ void performTest(float (*OP)(const float), const double atol = 1.0E-6; const double rtol = 1.0E-6; - // Set dump_data=true to enable dumping tensor data to files for analysis - compareResults_nvfp4(output, ref_output.get(), ref_output_t.get(), rows, cols, atol, rtol, true, - false, !row_scaled_nvfp4); - - size_t scale_mismatches_num = 0; - compare_scaling_factors("scales", output.rowwise_cpu_scale_inv_ptr(), - ref_scales.get(), - unpadded_blocks_Y, unpadded_blocks_X, scales_stride, - scale_mismatches_num); - - if (!row_scaled_nvfp4) { - compare_scaling_factors("scales_t", output.columnwise_cpu_scale_inv_ptr(), - ref_scales_t.get(), - unpadded_blocks_Y_t, unpadded_blocks_X_t, scales_stride_t, + if (use_4over6) { + output.to_cpu(); + compare_nvfp4_4over6_candidates("output", + output.rowwise_cpu_dptr(), + output.rowwise_cpu_scale_inv_ptr(), + ref_output.get(), + ref_scales.get(), + ref_output_map6.get(), + ref_scales_map6.get(), + rows, + cols, + unpadded_blocks_X, + scales_stride); + if (!row_scaled_nvfp4) { + compare_nvfp4_4over6_candidates("output_t", + output.columnwise_cpu_dptr(), + output.columnwise_cpu_scale_inv_ptr(), + ref_output_t.get(), + ref_scales_t.get(), + ref_output_t_map6.get(), + ref_scales_t_map6.get(), + cols, + rows, + unpadded_blocks_X_t, + scales_stride_t); + } + } else { + // Set dump_data=true to enable dumping tensor data to files for analysis + compareResults_nvfp4(output, ref_output.get(), ref_output_t.get(), rows, cols, atol, rtol, + true, false, !row_scaled_nvfp4); + + size_t scale_mismatches_num = 0; + compare_scaling_factors("scales", output.rowwise_cpu_scale_inv_ptr(), + ref_scales.get(), + unpadded_blocks_Y, unpadded_blocks_X, scales_stride, scale_mismatches_num); + + if (!row_scaled_nvfp4) { + compare_scaling_factors("scales_t", + output.columnwise_cpu_scale_inv_ptr(), + ref_scales_t.get(), + unpadded_blocks_Y_t, unpadded_blocks_X_t, + scales_stride_t, scale_mismatches_num); + } } compare_rowwise_amax(output, ref_amax); @@ -707,7 +1080,8 @@ class FusedCastTransposeNVFP4TestSuite : public ::testing::TestWithParam std::vector, transformer_engine::DType, bool, - bool>> {}; + NVFP4ScalingMode, + NVFP4FourOverSixTestConfig>> {}; TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { // Skip tests for pre-Blackwell architectures @@ -722,7 +1096,8 @@ TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { const auto tensor_dims = std::get<1>(GetParam()); const DType input_type = std::get<2>(GetParam()); const bool use_fast_math = std::get<3>(GetParam()); - const bool row_scaled_nvfp4 = std::get<4>(GetParam()); + const NVFP4ScalingMode scaling_mode = std::get<4>(GetParam()); + const NVFP4FourOverSixTestConfig config = std::get<5>(GetParam()); // Skip tests if the input tensor is 1D if (tensor_dims.size() < 2) { @@ -740,7 +1115,9 @@ TEST_P(FusedCastTransposeNVFP4TestSuite, TestFusedCastTransposeNVFP4) { } TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(input_type, InputType, - performTest(OP, tensor_dims, use_fast_math, row_scaled_nvfp4); + performTest(OP, tensor_dims, use_fast_math, scaling_mode, config.mode, + config.e4m3_max, + config.err_use_fast_math); ); } @@ -756,49 +1133,96 @@ std::string to_string(const ActivationType Act_type) { } } +std::string to_string(const NVFP4ScalingMode scaling_mode) { + switch (scaling_mode) { + case NVFP4ScalingMode::Block1D: return ""; + case NVFP4ScalingMode::RowScaled1D: return "XROW_SCALED"; + case NVFP4ScalingMode::Block2D: return "X2D"; + default: return ""; + } +} + +std::string test_name(const FusedCastTransposeNVFP4TestSuite::ParamType& param) { + std::string name = to_string(std::get<0>(param)); + const auto& shape = std::get<1>(param); + for (const auto& s: shape) { + name += "X" + std::to_string(s); + } + name += "X" + test::typeName(std::get<2>(param)); + if (std::get<3>(param)) { + name += "X_FAST_SCALING"; + } + name += to_string(std::get<4>(param)); + const NVFP4FourOverSixTestConfig& config = std::get<5>(param); + if (config.mode != kNVTENVFP44Over6Disabled) { + name += "X4OVER6"; + if (config.e4m3_max == 448) { + name += "XE4M3_MAX_448"; + } else { + name += "XE4M3_MAX_256"; + } + if (config.mode == kNVTENVFP44Over6MinMSE) { + name += "XMSE"; + } else if (config.mode == kNVTENVFP44Over6MinMAE) { + name += "XMAE"; + } else { + name += "XINVALID_MODE"; + } + if (config.err_use_fast_math) { + name += "XERR_USE_FAST_MATH"; + } + } + return name; +} + INSTANTIATE_TEST_SUITE_P( OperatorTest, FusedCastTransposeNVFP4TestSuite, ::testing::Combine( - ::testing::ValuesIn(Activation_types), - ::testing::ValuesIn(tensor_dims), - ::testing::Values(DType::kBFloat16), - ::testing::Values(false), - ::testing::Values(false)), + ::testing::ValuesIn(Activation_types), // activation_type + ::testing::ValuesIn(tensor_dims), // tensor_dims + ::testing::Values(DType::kBFloat16), // input_type + ::testing::Values(false), // use_fast_math + ::testing::Values(NVFP4ScalingMode::Block1D), // scaling_mode + ::testing::Values(NVFP4FourOverSixTestConfig{})), // four_over_six_config [](const testing::TestParamInfo& info) { - std::string name = to_string(std::get<0>(info.param)); - const auto& shape = std::get<1>(info.param); - for ( const auto& s: shape) { - name += "X" + std::to_string(s); - } - name += "X" + test::typeName(std::get<2>(info.param)); - if (std::get<3>(info.param)) { - name += "X_FAST_SCALING"; - } - return name; + return test_name(info.param); }); INSTANTIATE_TEST_SUITE_P( OperatorTestRowScaled, FusedCastTransposeNVFP4TestSuite, ::testing::Combine( - ::testing::Values(ActivationType::Identity), - ::testing::Values(tensor_dims[4], tensor_dims[9], tensor_dims[12]), - ::testing::Values(DType::kBFloat16, DType::kFloat32), - ::testing::Values(false), - ::testing::Values(true)), + ::testing::ValuesIn(Activation_types), // activation_type + ::testing::ValuesIn(tensor_dims), // tensor_dims + ::testing::Values(DType::kBFloat16, DType::kFloat32), // input_type + ::testing::Values(false), // use_fast_math + ::testing::Values(NVFP4ScalingMode::RowScaled1D), // scaling_mode + ::testing::Values(NVFP4FourOverSixTestConfig{})), // four_over_six_config [](const testing::TestParamInfo& info) { - std::string name = to_string(std::get<0>(info.param)); - const auto& shape = std::get<1>(info.param); - for (const auto& s: shape) { - name += "X" + std::to_string(s); - } - name += "X" + test::typeName(std::get<2>(info.param)); - if (std::get<3>(info.param)) { - name += "X_FAST_SCALING"; - } - if (std::get<4>(info.param)) { - name += "XROW_SCALED"; - } - return name; + return test_name(info.param); + }); + +INSTANTIATE_TEST_SUITE_P( + OperatorTest4Over6, + FusedCastTransposeNVFP4TestSuite, + ::testing::Combine( + ::testing::ValuesIn(Activation_types), // activation_type + ::testing::ValuesIn(tensor_dims), // tensor_dims + ::testing::Values(DType::kBFloat16, DType::kFloat32), // input_type + ::testing::Values(false), // use_fast_math + ::testing::Values(NVFP4ScalingMode::Block1D, + NVFP4ScalingMode::RowScaled1D, + NVFP4ScalingMode::Block2D), // scaling_mode + ::testing::Values( + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMAE, 448, false}, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMAE, 448, true}, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 448, false}, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 448, true}, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMAE, 256, false}, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMAE, 256, true}, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 256, false}, + NVFP4FourOverSixTestConfig{kNVTENVFP44Over6MinMSE, 256, true})), // four_over_six_config + [](const testing::TestParamInfo& info) { + return test_name(info.param); }); diff --git a/tests/cpp/operator/test_dequantize_nvfp4.cu b/tests/cpp/operator/test_dequantize_nvfp4.cu index eb9e8bce2..40c1fbd23 100644 --- a/tests/cpp/operator/test_dequantize_nvfp4.cu +++ b/tests/cpp/operator/test_dequantize_nvfp4.cu @@ -46,8 +46,9 @@ void compute_ref_dequantize_nvfp4(const uint8_t *packed_data, OType *output, size_t rows, size_t cols, - size_t scale_stride) { - constexpr float factor_inv = 1.0f / (6.0f * 448.0f); + size_t scale_stride, + int e4m3_max) { + const float factor_inv = 1.0f / (6.0f * static_cast(e4m3_max)); constexpr size_t BLOCK_SIZE = 16; const size_t Mread = cols / BLOCK_SIZE; const size_t bytes_per_block = BLOCK_SIZE / 2; @@ -86,11 +87,18 @@ float compute_amax(test::Tensor &t, size_t rows, size_t cols) { return amax; } +struct NVFP4DequantizeTestConfig { + NVTENVFP44Over6Mode mode = kNVTENVFP44Over6Disabled; + int e4m3_max = 448; +}; + // Quantize a high-precision input to NVFP4, then dequantize and compare // against a CPU reference computed from the quantized data. template void performTest_dequantize_nvfp4(const size_t rows, const size_t cols, - const bool row_scaled_nvfp4) { + const bool row_scaled_nvfp4, + const NVTENVFP44Over6Mode mode, + const int e4m3_max) { using namespace test; DType otype = TypeInfo::dtype; @@ -105,6 +113,8 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols, // Configure quantized tensor amax size_t amax_size = 1; + quantized.set_nvfp4_e4m3_max(e4m3_max); + ASSERT_EQ(quantized.nvfp4_e4m3_max(), e4m3_max); if (row_scaled_nvfp4) { quantized.set_row_scaled_nvfp4(true); amax_size = rows; @@ -116,7 +126,9 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols, // Quantize if (rows > 0 && cols > 0) { - nvte_quantize(input.data(), quantized.data(), 0); + QuantizationConfigWrapper quant_config; + quant_config.set_nvfp4_4over6_mode(mode); + nvte_quantize_v2(input.data(), quantized.data(), quant_config, 0); cudaDeviceSynchronize(); auto err = cudaGetLastError(); ASSERT_EQ(err, cudaSuccess) << cudaGetErrorString(err); @@ -146,7 +158,7 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols, std::make_unique(rows * cols); compute_ref_dequantize_nvfp4( fp4_data, scales, amax_vals, ref_output.get(), - rows, cols, scale_stride); + rows, cols, scale_stride, e4m3_max); // Compare results from TE and reference impls auto [atol, rtol] = getTolerances(otype); @@ -156,7 +168,9 @@ void performTest_dequantize_nvfp4(const size_t rows, const size_t cols, // Dequantize NVFP4 with GEMM-swizzled scales and compare against compact path. template void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols, - const bool row_scaled_nvfp4) { + const bool row_scaled_nvfp4, + const NVTENVFP44Over6Mode mode, + const int e4m3_max) { using namespace test; DType otype = TypeInfo::dtype; @@ -165,6 +179,8 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols, Tensor quantized_compact("quantized_compact", std::vector{rows, cols}, DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING); + quantized_compact.set_nvfp4_e4m3_max(e4m3_max); + ASSERT_EQ(quantized_compact.nvfp4_e4m3_max(), e4m3_max); if (row_scaled_nvfp4) { quantized_compact.set_row_scaled_nvfp4(true); } else if (rows > 0 && cols > 0) { @@ -174,7 +190,9 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols, } if (rows > 0 && cols > 0) { - nvte_quantize(input.data(), quantized_compact.data(), 0); + QuantizationConfigWrapper quant_config; + quant_config.set_nvfp4_4over6_mode(mode); + nvte_quantize_v2(input.data(), quantized_compact.data(), quant_config, 0); cudaDeviceSynchronize(); } @@ -186,6 +204,8 @@ void performTest_dequantize_nvfp4_swizzled(const size_t rows, const size_t cols, // Create tensor with same FP4 data but swizzled scales Tensor quantized_swizzled("quantized_swizzled", std::vector{rows, cols}, DType::kFloat4E2M1, true, false, NVTE_NVFP4_1D_SCALING); + quantized_swizzled.set_nvfp4_e4m3_max(e4m3_max); + ASSERT_EQ(quantized_swizzled.nvfp4_e4m3_max(), e4m3_max); if (row_scaled_nvfp4) { quantized_swizzled.set_row_scaled_nvfp4(true); } else { @@ -260,7 +280,8 @@ std::vector> nvfp4_tensor_dims = { class DequantizeNVFP4TestSuite : public ::testing::TestWithParam , transformer_engine::DType, - bool>> {}; + bool, + NVFP4DequantizeTestConfig>> {}; TEST_P(DequantizeNVFP4TestSuite, TestDequantizeNVFP4) { @@ -271,10 +292,12 @@ TEST_P(DequantizeNVFP4TestSuite, TestDequantizeNVFP4) const auto tensor_size = std::get<0>(GetParam()); const DType output_type = std::get<1>(GetParam()); const bool row_scaled_nvfp4 = std::get<2>(GetParam()); + const NVFP4DequantizeTestConfig config = std::get<3>(GetParam()); TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(output_type, OutputType, performTest_dequantize_nvfp4( - tensor_size.first, tensor_size.second, row_scaled_nvfp4); + tensor_size.first, tensor_size.second, row_scaled_nvfp4, config.mode, + config.e4m3_max); ); } @@ -284,13 +307,20 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Combine( ::testing::ValuesIn(nvfp4_tensor_dims), ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16), - ::testing::Bool()), + ::testing::Bool(), + ::testing::Values(NVFP4DequantizeTestConfig{}, + NVFP4DequantizeTestConfig{kNVTENVFP44Over6MinMAE, 448}, + NVFP4DequantizeTestConfig{kNVTENVFP44Over6MinMAE, 256})), [](const testing::TestParamInfo& info) { + const NVFP4DequantizeTestConfig config = std::get<3>(info.param); + const bool use_4over6 = config.mode != kNVTENVFP44Over6Disabled; std::string name = std::to_string(std::get<0>(info.param).first) + "X" + std::to_string(std::get<0>(info.param).second) + "X" + test::typeName(std::get<1>(info.param)) + "X" + - (std::get<2>(info.param) ? "RowScaled" : "PerTensor"); + (std::get<2>(info.param) ? "RowScaled" : "PerTensor") + "X" + + (use_4over6 ? "FourOverSix" : "Default") + "X" + + (config.e4m3_max == 256 ? "E4M3Max256" : "E4M3Max448"); return name; } ); @@ -298,7 +328,8 @@ INSTANTIATE_TEST_SUITE_P( class DequantizeNVFP4SwizzledTestSuite : public ::testing::TestWithParam , transformer_engine::DType, - bool>> {}; + bool, + NVFP4DequantizeTestConfig>> {}; TEST_P(DequantizeNVFP4SwizzledTestSuite, TestDequantizeNVFP4Swizzled) { @@ -309,10 +340,12 @@ TEST_P(DequantizeNVFP4SwizzledTestSuite, TestDequantizeNVFP4Swizzled) const auto tensor_size = std::get<0>(GetParam()); const DType output_type = std::get<1>(GetParam()); const bool row_scaled_nvfp4 = std::get<2>(GetParam()); + const NVFP4DequantizeTestConfig config = std::get<3>(GetParam()); TRANSFORMER_ENGINE_TYPE_SWITCH_FP16_FP32_ONLY(output_type, OutputType, performTest_dequantize_nvfp4_swizzled( - tensor_size.first, tensor_size.second, row_scaled_nvfp4); + tensor_size.first, tensor_size.second, row_scaled_nvfp4, config.mode, + config.e4m3_max); ); } @@ -322,13 +355,20 @@ INSTANTIATE_TEST_SUITE_P( ::testing::Combine( ::testing::ValuesIn(nvfp4_tensor_dims), ::testing::Values(DType::kFloat32, DType::kBFloat16, DType::kFloat16), - ::testing::Bool()), + ::testing::Bool(), + ::testing::Values(NVFP4DequantizeTestConfig{}, + NVFP4DequantizeTestConfig{kNVTENVFP44Over6MinMAE, 448}, + NVFP4DequantizeTestConfig{kNVTENVFP44Over6MinMAE, 256})), [](const testing::TestParamInfo& info) { + const NVFP4DequantizeTestConfig config = std::get<3>(info.param); + const bool use_4over6 = config.mode != kNVTENVFP44Over6Disabled; std::string name = std::to_string(std::get<0>(info.param).first) + "X" + std::to_string(std::get<0>(info.param).second) + "X" + test::typeName(std::get<1>(info.param)) + "X" + (std::get<2>(info.param) ? "RowScaled" : "PerTensor") + "X" + + (use_4over6 ? "FourOverSix" : "Default") + "X" + + (config.e4m3_max == 256 ? "E4M3Max256" : "E4M3Max448") + "X" + "Swizzled"; return name; } diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index 4fd75bb92..e35f5e029 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -440,6 +440,18 @@ void Tensor::set_row_scaled_nvfp4(bool row_scaled_nvfp4) { } } +void Tensor::set_nvfp4_e4m3_max(int nvfp4_e4m3_max) { + NVTE_CHECK(tensor_.scaling_mode() == NVTE_NVFP4_1D_SCALING, + "NVFP4 E4M3 max is only supported for NVFP4 tensors."); + tensor_.set_nvfp4_e4m3_max(nvfp4_e4m3_max); +} + +int Tensor::nvfp4_e4m3_max() const { + NVTE_CHECK(tensor_.scaling_mode() == NVTE_NVFP4_1D_SCALING, + "NVFP4 E4M3 max is only supported for NVFP4 tensors."); + return tensor_.get_nvfp4_e4m3_max(); +} + void Tensor::to_cpu() { if (data_rowwise_) { data_rowwise_->to_cpu(); } if (data_columnwise_) { data_columnwise_->to_cpu(); } diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index 17f36a99d..fd03d283d 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -293,10 +293,13 @@ class Tensor { return columnwise_; } + int nvfp4_e4m3_max() const; + void set_tensor_amax_nullptr(); void set_with_gemm_swizzled_scales(bool with_gemm_swizzled_scales); void set_row_scaled_nvfp4(bool row_scaled_nvfp4); + void set_nvfp4_e4m3_max(int nvfp4_e4m3_max); void to_cpu(); void from_cpu(); diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py index a7ea4f089..bd4d02972 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py @@ -28,7 +28,12 @@ def check_nvfp4_gemm_versus_reference( x_columnwise: bool = False, w_columnwise: bool = False, row_scaled_nvfp4: bool = False, + use_4over6: bool = False, + nvfp4_e4m3_max: int = 448, + nvfp4_4over6_err_mode: str = "MAE", ): + if nvfp4_e4m3_max != 448 and not use_4over6: + pytest.skip("E4M3 max 256 is only meaningful for 4over6") te_dtype = tex.DType.kFloat4E2M1 # Setup device and random seed @@ -59,6 +64,9 @@ def check_nvfp4_gemm_versus_reference( with_rht=False, with_post_rht_amax=False, row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) w_quantizer = NVFP4Quantizer( fp4_dtype=te_dtype, @@ -68,6 +76,9 @@ def check_nvfp4_gemm_versus_reference( amax_reduction_group=None, with_rht=False, with_post_rht_amax=False, + nvfp4_use_4over6=use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) # Quantize x and w @@ -123,6 +134,9 @@ def check_nvfp4_gemm_versus_reference( eps=0.0, quant_tile_shape=(1, 16), row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) w_ref_quantizer = NVFP4QuantizerRef( dtype=utils.Fp4Formats.E2M1, @@ -131,6 +145,9 @@ def check_nvfp4_gemm_versus_reference( pow_2_scales=False, eps=0.0, quant_tile_shape=(1, 16), + nvfp4_use_4over6=use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) # Create reference quantized tensors needed by reference GEMM @@ -232,6 +249,8 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( *, use_bias: bool, single_output: bool, + use_4over6: bool = False, + nvfp4_4over6_err_mode: str = "MAE", ): te_dtype = tex.DType.kFloat4E2M1 device = "cuda" @@ -249,6 +268,8 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( with_rht=False, with_post_rht_amax=False, row_scaled_nvfp4=True, + nvfp4_use_4over6=use_4over6, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) w_quantizer = NVFP4Quantizer( fp4_dtype=te_dtype, @@ -258,6 +279,8 @@ def check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( amax_reduction_group=None, with_rht=False, with_post_rht_amax=False, + nvfp4_use_4over6=use_4over6, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) x_nvfp4 = [] @@ -321,6 +344,8 @@ def check_nvfp4_row_scaled_gemm_matches_emulated( M: int, K: int, N: int, + use_4over6: bool = False, + nvfp4_4over6_err_mode: str = "MAE", ): te_dtype = tex.DType.kFloat4E2M1 device = "cuda" @@ -339,6 +364,8 @@ def check_nvfp4_row_scaled_gemm_matches_emulated( with_rht=False, with_post_rht_amax=False, row_scaled_nvfp4=True, + nvfp4_use_4over6=use_4over6, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) x_tensorwise_quantizer = NVFP4Quantizer( fp4_dtype=te_dtype, @@ -348,6 +375,8 @@ def check_nvfp4_row_scaled_gemm_matches_emulated( amax_reduction_group=None, with_rht=False, with_post_rht_amax=False, + nvfp4_use_4over6=use_4over6, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) w_quantizer = NVFP4Quantizer( fp4_dtype=te_dtype, @@ -357,6 +386,8 @@ def check_nvfp4_row_scaled_gemm_matches_emulated( amax_reduction_group=None, with_rht=False, with_post_rht_amax=False, + nvfp4_use_4over6=use_4over6, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) x_row_scaled = x_row_scaled_quantizer.update_quantized( @@ -417,6 +448,9 @@ def check_nvfp4_row_scaled_gemm_matches_emulated( ids=["rowxrow", "colxrow", "colxcol"], ) @pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"]) +@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) +@pytest.mark.parametrize("nvfp4_e4m3_max", [448, 256], ids=["e4m3_448", "e4m3_256"]) +@pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"]) def test_nvfp4_gemm_versus_reference( M: int, K: int, @@ -428,6 +462,9 @@ def test_nvfp4_gemm_versus_reference( is_x_columnwise: bool, is_w_columnwise: bool, row_scaled_nvfp4: bool, + use_4over6: bool, + nvfp4_e4m3_max: int, + nvfp4_4over6_err_mode: str, ): if row_scaled_nvfp4: if accumulate: @@ -446,6 +483,9 @@ def test_nvfp4_gemm_versus_reference( x_columnwise=is_x_columnwise, w_columnwise=is_w_columnwise, row_scaled_nvfp4=row_scaled_nvfp4, + use_4over6=use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) @@ -471,6 +511,8 @@ def test_nvfp4_gemm_versus_reference( @pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("use_bias", [False, True], ids=["no_bias", "bias"]) @pytest.mark.parametrize("single_output", [False, True], ids=["list_output", "single_output"]) +@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) +@pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"]) def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( m_splits: list[int], k: int, @@ -480,6 +522,8 @@ def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( out_dtype: torch.dtype, use_bias: bool, single_output: bool, + use_4over6: bool, + nvfp4_4over6_err_mode: str, ): check_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( x_dtype=x_dtype, @@ -490,6 +534,8 @@ def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( n=n, use_bias=use_bias, single_output=single_output, + use_4over6=use_4over6, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) @@ -513,6 +559,8 @@ def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm( @pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("w_dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str) +@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) +@pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"]) def test_nvfp4_row_scaled_gemm_matches_emulated( M: int, K: int, @@ -520,6 +568,8 @@ def test_nvfp4_row_scaled_gemm_matches_emulated( x_dtype: torch.dtype, w_dtype: torch.dtype, out_dtype: torch.dtype, + use_4over6: bool, + nvfp4_4over6_err_mode: str, ): check_nvfp4_row_scaled_gemm_matches_emulated( x_dtype=x_dtype, @@ -528,4 +578,6 @@ def test_nvfp4_row_scaled_gemm_matches_emulated( M=M, K=K, N=N, + use_4over6=use_4over6, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py index 53569d90d..5bb92f70d 100644 --- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py +++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py @@ -20,7 +20,14 @@ def maybe_skip_row_scaled_unsupported_quantization( row_scaled_nvfp4: bool, return_transpose: bool, with_2d_quantization: bool = False, + use_4over6: bool = False, + x_dtype: torch.dtype | None = None, + M: int | None = None, + N: int | None = None, ) -> None: + if use_4over6 and with_2d_quantization: + if x_dtype != torch.bfloat16 or M is None or N is None or M % 32 != 0 or N % 32 != 0: + pytest.skip("NVFP4 2D 4over6 exact tests require the optimized BF16 kernel path") if not row_scaled_nvfp4: return if return_transpose: @@ -45,9 +52,14 @@ def check_quantization_nvfp4_versus_reference( use_cpp_allocator: bool, with_2d_quantization: bool, row_scaled_nvfp4: bool = False, + use_4over6: bool = False, + nvfp4_e4m3_max: int = 448, + nvfp4_4over6_err_mode: str = "MAE", ) -> None: + if nvfp4_e4m3_max != 448 and not use_4over6: + pytest.skip("E4M3 max 256 is only meaningful for 4over6") maybe_skip_row_scaled_unsupported_quantization( - row_scaled_nvfp4, return_transpose, with_2d_quantization + row_scaled_nvfp4, return_transpose, with_2d_quantization, use_4over6, x_dtype, M, N ) te_dtype = tex.DType.kFloat4E2M1 @@ -71,6 +83,9 @@ def check_quantization_nvfp4_versus_reference( with_post_rht_amax=False, with_2d_quantization=with_2d_quantization, row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) if use_cpp_allocator: x_nvfp4_sut = nvfp4_quantizer(x) @@ -104,6 +119,9 @@ def check_quantization_nvfp4_versus_reference( eps=0.0, quant_tile_shape=quant_tile_shape, row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) x_nvfp4_ref = ref_quantizer.quantize(x) @@ -179,6 +197,9 @@ def check_quantization_nvfp4_versus_reference( "with_2d_quantization", [True, False], ids=["2d_quantization", "1d_quantization"] ) @pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"]) +@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) +@pytest.mark.parametrize("nvfp4_e4m3_max", [448, 256], ids=["e4m3_448", "e4m3_256"]) +@pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"]) def test_quantization_block_tiling_versus_reference( x_dtype: torch.dtype, M: int, @@ -188,6 +209,9 @@ def test_quantization_block_tiling_versus_reference( use_cpp_allocator: bool, with_2d_quantization: bool, row_scaled_nvfp4: bool, + use_4over6: bool, + nvfp4_e4m3_max: int, + nvfp4_4over6_err_mode: str, ) -> None: check_quantization_nvfp4_versus_reference( x_dtype=x_dtype, @@ -198,6 +222,9 @@ def test_quantization_block_tiling_versus_reference( use_cpp_allocator=use_cpp_allocator, with_2d_quantization=with_2d_quantization, row_scaled_nvfp4=row_scaled_nvfp4, + use_4over6=use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) @@ -215,6 +242,8 @@ def test_quantization_block_tiling_versus_reference( "use_cpp_allocator", [True, False], ids=["cpp_allocator", "python_allocator"] ) @pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"]) +@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) +@pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"]) def test_nvfp4_quantization_extrema_versus_reference( x_dtype: torch.dtype, M: int, @@ -223,8 +252,12 @@ def test_nvfp4_quantization_extrema_versus_reference( return_transpose: bool, use_cpp_allocator: bool, row_scaled_nvfp4: bool, + use_4over6: bool, + nvfp4_4over6_err_mode: str, ): - maybe_skip_row_scaled_unsupported_quantization(row_scaled_nvfp4, return_transpose) + maybe_skip_row_scaled_unsupported_quantization( + row_scaled_nvfp4, return_transpose, use_4over6=use_4over6 + ) te_dtype = tex.DType.kFloat4E2M1 @@ -247,6 +280,8 @@ def test_nvfp4_quantization_extrema_versus_reference( with_rht=False, with_post_rht_amax=False, row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=use_4over6, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) if use_cpp_allocator: @@ -278,6 +313,8 @@ def test_nvfp4_quantization_extrema_versus_reference( eps=0.0, quant_tile_shape=(1, 16), row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=use_4over6, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) x_nvfp4_ref = ref_quantizer.quantize(x) @@ -322,6 +359,8 @@ def test_nvfp4_quantization_extrema_versus_reference( "use_cpp_allocator", [True, False], ids=["cpp_allocator", "python_allocator"] ) @pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"]) +@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) +@pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"]) def test_nvfp4_quantization_boundary_values( x_dtype: torch.dtype, M: int, @@ -329,13 +368,17 @@ def test_nvfp4_quantization_boundary_values( return_transpose: bool, use_cpp_allocator: bool, row_scaled_nvfp4: bool, + use_4over6: bool, + nvfp4_4over6_err_mode: str, ): """ Stress rounding/threshold behavior by placing values just below/above many potential bin edges within each 16-element microblock. Validates native vs reference byte-for-byte and scale parity. """ - maybe_skip_row_scaled_unsupported_quantization(row_scaled_nvfp4, return_transpose) + maybe_skip_row_scaled_unsupported_quantization( + row_scaled_nvfp4, return_transpose, use_4over6=use_4over6 + ) te_dtype = tex.DType.kFloat4E2M1 @@ -367,6 +410,8 @@ def test_nvfp4_quantization_boundary_values( with_rht=False, with_post_rht_amax=False, row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=use_4over6, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) if use_cpp_allocator: @@ -398,6 +443,8 @@ def test_nvfp4_quantization_boundary_values( eps=0.0, quant_tile_shape=(1, 16), row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=use_4over6, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) x_nvfp4_ref = ref_quantizer.quantize(x) @@ -442,6 +489,8 @@ def test_nvfp4_quantization_boundary_values( "use_cpp_allocator", [True, False], ids=["cpp_allocator", "python_allocator"] ) @pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"]) +@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) +@pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"]) def test_nvfp4_quantization_noncontiguous_inputs( x_dtype: torch.dtype, M: int, @@ -449,8 +498,12 @@ def test_nvfp4_quantization_noncontiguous_inputs( return_transpose: bool, use_cpp_allocator: bool, row_scaled_nvfp4: bool, + use_4over6: bool, + nvfp4_4over6_err_mode: str, ): - maybe_skip_row_scaled_unsupported_quantization(row_scaled_nvfp4, return_transpose) + maybe_skip_row_scaled_unsupported_quantization( + row_scaled_nvfp4, return_transpose, use_4over6=use_4over6 + ) te_dtype = tex.DType.kFloat4E2M1 @@ -473,6 +526,8 @@ def test_nvfp4_quantization_noncontiguous_inputs( with_rht=False, with_post_rht_amax=False, row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=use_4over6, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) if use_cpp_allocator: @@ -504,6 +559,8 @@ def test_nvfp4_quantization_noncontiguous_inputs( eps=0.0, quant_tile_shape=(1, 16), row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=use_4over6, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, ) x_nvfp4_ref = ref_quantizer.quantize(x_nc) diff --git a/tests/pytorch/test_backward_override.py b/tests/pytorch/test_backward_override.py index 43e9587d9..5e6f36e8b 100644 --- a/tests/pytorch/test_backward_override.py +++ b/tests/pytorch/test_backward_override.py @@ -83,6 +83,11 @@ marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4), id="NVFP4RowScaledBlockScaling", ), + pytest.param( + "nvfp4_4over6", + marks=pytest.mark.skipif(not nvfp4_available, reason=reason_for_no_nvfp4), + id="NVFP44Over6BlockScaling", + ), ] @@ -170,7 +175,7 @@ def _maybe_skip_recipe_dtype( ) -> None: if dtype == torch.bfloat16 and not bf16_available: pytest.skip(reason_for_no_bf16) - if recipe_name in ("nvfp4", "nvfp4_row_scaled"): + if recipe_name in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): if module_type in ("linear", "layernorm_linear") and dtype not in ( torch.bfloat16, torch.float32, @@ -185,6 +190,8 @@ def _maybe_skip_unsupported_recipe_module_combo(recipe_name: str, module_type: s pytest.skip("Fusible ops (te_ops.Linear) do not support Float8BlockScaling recipe") if module_type == "ops_linear" and recipe_name == "nvfp4_row_scaled": pytest.skip("Row-scaled NVFP4 currently does not support fused te_ops paths.") + if module_type == "grouped_linear" and recipe_name == "nvfp4_4over6": + pytest.skip("NVFP4 4over6 currently does not support grouped quantization.") def _make_quantized_forward_reference_recipe(recipe_name: str) -> recipe.Recipe: @@ -208,7 +215,7 @@ def _maybe_skip_unsupported_recipe_shape( " by 32." ) return - if recipe_name in ("nvfp4", "nvfp4_row_scaled") and ( + if recipe_name in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6") and ( flat_first_dim % 16 != 0 or last_dim % 16 != 0 ): pytest.skip( @@ -235,7 +242,7 @@ def _maybe_skip_unsupported_recipe_shape( pytest.skip( "te_ops.Linear + MXFP8 requires prod(shape[:-1]) and shape[-1] divisible by 32." ) - if recipe_name in ("nvfp4", "nvfp4_row_scaled") and ( + if recipe_name in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6") and ( flat_first_dim % 16 != 0 or last_dim % 16 != 0 ): pytest.skip( @@ -256,9 +263,13 @@ def _maybe_skip_unsupported_grouped_splits(recipe_name: str, m_splits: list[int] ) if recipe_name == "mxfp8" and any(m % 32 != 0 for m in non_empty_splits): pytest.skip("GroupedLinear + MXFP8 requires each non-empty m_split divisible by 32.") - if recipe_name in ("nvfp4", "nvfp4_row_scaled") and any(m % 16 != 0 for m in non_empty_splits): + if recipe_name in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6") and any( + m % 16 != 0 for m in non_empty_splits + ): pytest.skip("GroupedLinear + NVFP4 requires each non-empty m_split divisible by 16.") - if recipe_name in ("nvfp4", "nvfp4_row_scaled") and any(m % 64 != 0 for m in non_empty_splits): + if recipe_name in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6") and any( + m % 64 != 0 for m in non_empty_splits + ): pytest.skip( "GroupedLinear + NVFP4 grouped split_quantize currently requires each non-empty " "m_split divisible by 64 due to grouped amax kernel constraints." diff --git a/tests/pytorch/test_cpu_offloading.py b/tests/pytorch/test_cpu_offloading.py index 50196782f..35cc98a97 100644 --- a/tests/pytorch/test_cpu_offloading.py +++ b/tests/pytorch/test_cpu_offloading.py @@ -19,7 +19,7 @@ from transformer_engine.pytorch.fp8 import FP8GlobalStateManager import transformer_engine.pytorch as te from transformer_engine.common import recipe -from utils import ModelConfig, skip_unsupported_backward_override +from utils import ModelConfig, recipe_id, skip_unsupported_backward_override import transformer_engine_torch as tex # Check supported quantization schemes @@ -28,6 +28,33 @@ mxfp8_available, _ = FP8GlobalStateManager.is_mxfp8_available() nvfp4_available, _ = FP8GlobalStateManager.is_nvfp4_available() + +def nvfp4_row_scaled(): + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + row_scaled_activation=True, + backward_override="dequantized", + ) + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + return nvfp4_recipe + + +def nvfp4_4over6(): + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + nvfp4_4over6="all", + ) + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams(fp4_2d_quantization=True) + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + return nvfp4_recipe + + quantization_recipes: List[Optional[recipe.Recipe]] = [None] if fp8_available: quantization_recipes.extend((recipe.Float8CurrentScaling(), recipe.DelayedScaling())) @@ -37,6 +64,8 @@ quantization_recipes.append(recipe.MXFP8BlockScaling()) if nvfp4_available: quantization_recipes.append(recipe.NVFP4BlockScaling()) + quantization_recipes.append(nvfp4_4over6()) + quantization_recipes.append(nvfp4_row_scaled()) model_config = { @@ -176,7 +205,20 @@ def create_tensor(recipe: Optional[recipe.Recipe], requires_grad: bool = False) quantizer = te.tensor.mxfp8_tensor.MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) return quantizer(tensor) elif recipe.nvfp4(): - quantizer = te.tensor.nvfp4_tensor.NVFP4Quantizer() + qparams = recipe.fp4_quant_fwd_inp + use_4over6 = False + if recipe.nvfp4_4over6 in ("activations", "all"): + use_4over6 = True + quantizer = te.tensor.nvfp4_tensor.NVFP4Quantizer( + rowwise=True, + columnwise=not recipe.row_scaled_activation, + with_rht=qparams.random_hadamard_transform, + with_post_rht_amax=qparams.random_hadamard_transform, + with_2d_quantization=qparams.fp4_2d_quantization, + stochastic_rounding=qparams.stochastic_rounding, + row_scaled_nvfp4=recipe.row_scaled_activation, + nvfp4_use_4over6=use_4over6, + ) return quantizer(tensor) @staticmethod @@ -191,10 +233,24 @@ def get_tensor_size_mb(tensor): if tensor is None: return 0 if isinstance(tensor, te.quantized_tensor.QuantizedTensorStorage): - return sum(Utils.get_tensor_size_mb(t) for t in tensor.get_data_tensors()) + tensors = [ + value for value in tensor.get_metadata().values() if isinstance(value, torch.Tensor) + ] + return sum(Utils.get_tensor_size_mb(t) for t in tensors) else: return tensor.numel() * tensor.element_size() / (1024**2) + @staticmethod + def get_saved_tensor_gpu_size_mb(tensor): + if tensor is None or isinstance(tensor, int): + return 0 + if isinstance(tensor, tuple): + push_results, _ = tensor + return Utils.get_saved_tensor_gpu_size_mb(push_results) + if isinstance(tensor, list): + return sum(Utils.get_saved_tensor_gpu_size_mb(t) for t in tensor) + return Utils.get_tensor_size_mb(tensor) + @staticmethod def memory_leak_check(): # Should be called before each test. @@ -212,7 +268,7 @@ def memory_leak_check(): class TestsOffloadableLayerState: @pytest.mark.parametrize("random_num_tensors", [True, False]) - @pytest.mark.parametrize("recipe", quantization_recipes) + @pytest.mark.parametrize("recipe", quantization_recipes, ids=recipe_id) def test_general(self, random_num_tensors, recipe): """ Test general functionality of DefaultOffloadSynchronizer - offload NUM_LAYERS-1 out of NUM_LAYERS layers, @@ -289,7 +345,7 @@ def test_offload_base_tensor(self): class TestsDefaultOffloadSynchronizer: @pytest.mark.parametrize("random_num_tensors", [True, False]) - @pytest.mark.parametrize("recipe", quantization_recipes) + @pytest.mark.parametrize("recipe", quantization_recipes, ids=recipe_id) def test_general(self, random_num_tensors, recipe): """ Test general functionality of DefaultOffloadSynchronizer - offload NUM_LAYERS-1 out of NUM_LAYERS layers, @@ -335,7 +391,7 @@ def test_general(self, random_num_tensors, recipe): offload_synchronizer.finish_part_of_bwd() torch.cuda.synchronize() - @pytest.mark.parametrize("recipe", quantization_recipes) + @pytest.mark.parametrize("recipe", quantization_recipes, ids=recipe_id) def test_memory(self, recipe): torch.cuda.synchronize() Utils.memory_leak_check() @@ -363,11 +419,16 @@ def test_memory(self, recipe): del tensor, tensor_id torch.cuda.synchronize() + resident_gpu_size = sum( + Utils.get_saved_tensor_gpu_size_mb(tensor_id) for tensor_id in tensor_ids + ) if recipe is None: assert Utils.get_max_cuda_memory_mb() == pytest.approx( - init_cuda_memory + tensor_size, 0.1 + init_cuda_memory + resident_gpu_size, 0.1 ) - assert Utils.get_cuda_memory_mb() == pytest.approx(init_cuda_memory + tensor_size, 0.1) + assert Utils.get_cuda_memory_mb() == pytest.approx( + init_cuda_memory + resident_gpu_size, 0.1 + ) for i in range(NUM_LAYERS - 1, -1, -1): offload_synchronizer.bwd_step(i) @@ -385,7 +446,7 @@ def test_memory(self, recipe): ) assert Utils.get_cuda_memory_mb() == pytest.approx(init_cuda_memory, 0.1) - @pytest.mark.parametrize("recipe", quantization_recipes) + @pytest.mark.parametrize("recipe", quantization_recipes, ids=recipe_id) def test_multiple_tensor_offload(self, recipe): Utils.memory_leak_check() init_cpu_memory = Utils.get_cpu_memory_mb() @@ -416,7 +477,7 @@ def test_multiple_tensor_offload(self, recipe): class TestTELayers: @pytest.mark.parametrize("layer_type", Utils.get_layer_names()) - @pytest.mark.parametrize("recipe", quantization_recipes) + @pytest.mark.parametrize("recipe", quantization_recipes, ids=recipe_id) @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) def test_sanity(self, layer_type, recipe, backward_override): Utils.memory_leak_check() @@ -463,7 +524,7 @@ def test_sanity(self, layer_type, recipe, backward_override): del out, inp, layers @pytest.mark.parametrize("layer_type", Utils.get_layer_names()) - @pytest.mark.parametrize("recipe", quantization_recipes) + @pytest.mark.parametrize("recipe", quantization_recipes, ids=recipe_id) @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) def test_memory(self, layer_type, recipe, backward_override): Utils.memory_leak_check() @@ -536,7 +597,9 @@ def test_memory(self, layer_type, recipe, backward_override): out = out + 1 out = sync_function(out) del inp - if backward_override is None: + if recipe is not None and recipe.nvfp4() and recipe.row_scaled_activation: + assert Utils.get_cuda_memory_mb() <= cuda_memory_no_offload + elif backward_override is None: assert Utils.get_cuda_memory_mb() == pytest.approx(init_cuda_memory, 0.1) else: assert ( @@ -554,7 +617,7 @@ def test_memory(self, layer_type, recipe, backward_override): out.sum().backward() @pytest.mark.parametrize("layer_type", Utils.get_layer_names()) - @pytest.mark.parametrize("recipe", quantization_recipes) + @pytest.mark.parametrize("recipe", quantization_recipes, ids=recipe_id) @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) def test_manual_synchronization(self, recipe, layer_type, backward_override): Utils.memory_leak_check() @@ -623,7 +686,7 @@ def test_manual_synchronization(self, recipe, layer_type, backward_override): out_1.sum().backward() out_2.sum().backward() - @pytest.mark.parametrize("recipe", quantization_recipes) + @pytest.mark.parametrize("recipe", quantization_recipes, ids=recipe_id) @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) @pytest.mark.parametrize("layer_type", Utils.get_layer_names()) @pytest.mark.parametrize("use_cuda_graphs", [True, False]) diff --git a/tests/pytorch/test_cuda_graphs.py b/tests/pytorch/test_cuda_graphs.py index 33ba65e0d..bb4a4e385 100644 --- a/tests/pytorch/test_cuda_graphs.py +++ b/tests/pytorch/test_cuda_graphs.py @@ -65,13 +65,31 @@ def nvfp4_rht_and_2d_quantization(): def nvfp4_row_scaled(): - nvfp4_recipe = recipe.NVFP4BlockScaling(row_scaled_activation=True) + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + row_scaled_activation=True, + backward_override="dequantized", + ) nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() return nvfp4_recipe +def nvfp4_4over6(): + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + nvfp4_4over6="all", + ) + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams(fp4_2d_quantization=True) + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + return nvfp4_recipe + + def check_rht_usage(recipe: recipe.Recipe) -> bool: # if using RHT, we can only support bf16 # check fp4_quant_fwd_inp, fp4_quant_fwd_weight, fp4_quant_bwd_grad @@ -101,6 +119,7 @@ def get_nvfp4_inp_supported_dtypes(recipe: recipe.Recipe, dtype: torch.dtype) -> if nvfp4_available: fp8_recipes.append(nvfp4_rht_and_2d_quantization()) fp8_recipes.append(nvfp4_row_scaled()) + fp8_recipes.append(nvfp4_4over6()) if fp8_block_scaling_available: fp8_recipes.append(recipe.Float8BlockScaling()) if fp8_available: diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 3a3aa8be9..8e63caa98 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -40,6 +40,7 @@ Float8Quantizer, MXFP8Quantizer, NVFP4Quantizer, + QuantizerRole, is_bf16_available, ) from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor @@ -78,9 +79,10 @@ _quantization_list.append("mxfp8") if nvfp4_available: _quantization_list.append("nvfp4") + _quantization_list.append("nvfp4_4over6") -@pytest.fixture(autouse=True, scope="class") +@pytest.fixture(autouse=True, scope="function") def _reset_rng_states_per_test(): """Restore torch, CUDA, and Python ``random`` before each test in this module.""" reset_rng_states() @@ -107,7 +109,7 @@ def maybe_skip_quantization( pytest.skip(reason_for_no_fp8) if quantization == "mxfp8" and not mxfp8_available: pytest.skip(reason_for_no_mxfp8) - if quantization == "nvfp4" and not nvfp4_available: + if quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6") and not nvfp4_available: pytest.skip(reason_for_no_nvfp4) # Check dims @@ -120,13 +122,16 @@ def maybe_skip_quantization( elif quantization == "mxfp8": if math.prod(dims[:-1]) % 32 != 0 or dims[-1] % 32 != 0: pytest.skip("MXFP8 GEMMs require dims that are divisible by 32") - elif quantization == "nvfp4": + elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") # Check dtype if dtype is not None: - if quantization == "nvfp4" and dtype != torch.bfloat16: + if ( + quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6") + and dtype != torch.bfloat16 + ): pytest.skip("NVFP4 quantization is only supported with BF16 data") @@ -142,6 +147,7 @@ def make_reference_and_test_tensors( test_dtype: torch.dtype = torch.float32, test_device: torch.device = "cuda", test_is_quantized: bool = False, + quantizer_role: Optional[QuantizerRole] = None, requires_grad: bool = True, ) -> tuple[torch.Tensor, torch.Tensor]: """Construct tensors with the same values @@ -181,7 +187,7 @@ def make_reference_and_test_tensors( test = quantizer(test) elif quantization == "mxfp8": test = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3)(test) - elif quantization == "nvfp4": + elif quantization in ("nvfp4", "nvfp4_row_scaled"): test = NVFP4Quantizer( with_rht=False, with_post_rht_amax=False, @@ -189,6 +195,29 @@ def make_reference_and_test_tensors( stochastic_rounding=False, with_random_sign_mask=False, )(test) + elif quantization == "nvfp4_4over6": + tensor_type = "input" + if quantizer_role is not None: + tensor_type = quantizer_role.tensor_type + + nvfp4_use_4over6 = False + with_2d_quantization = False + nvfp4_e4m3_max = 448 + if tensor_type not in ("grad_output", "grad_input"): + nvfp4_use_4over6 = True + nvfp4_e4m3_max = 256 + if tensor_type == "weight": + with_2d_quantization = True + + test = NVFP4Quantizer( + with_rht=False, + with_post_rht_amax=False, + with_2d_quantization=with_2d_quantization, + stochastic_rounding=False, + with_random_sign_mask=False, + nvfp4_use_4over6=nvfp4_use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, + )(test) else: raise ValueError(f"Unsupported quantization scheme ({quantization})") if isinstance(test, QuantizedTensor) and not test_is_quantized: @@ -504,6 +533,7 @@ def test_dtype_cast( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="weight"), ) # Construct operation @@ -818,9 +848,13 @@ def test_quantize( test_device=device, requires_grad=True, ) + grad_quantization = quantization + if quantization == "nvfp4_4over6" and cast_backward: + # 4over6 is not applied to gradient quantizers. + grad_quantization = "nvfp4" dy_ref, dy_test = make_reference_and_test_tensors( in_shape, - quantization=quantization, + quantization=grad_quantization, test_dtype=dtype, test_device=device, requires_grad=False, @@ -911,6 +945,7 @@ def _test_basic_linear( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="weight"), ) dy_ref, dy_test = make_reference_and_test_tensors( out_shape, @@ -1083,6 +1118,7 @@ def test_linear( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="weight"), ) b_ref, b_test = None, None if bias: @@ -1513,7 +1549,7 @@ def test_add_extra_input( if in_place: if quantization in ("fp8_delayed_scaling", "fp8_current_scaling", "mxfp8"): tols = dtype_tols(x1_test._fp8_dtype) - elif quantization == "nvfp4": + elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): tols = dtype_tols(x1_test._fp4_dtype) y_test = y_test.to(dtype=torch.float64, device="cpu") dx1_test = x1_test.grad.to(dtype=torch.float64, device="cpu") @@ -1884,7 +1920,7 @@ def test_clamped_swiglu( # Expected numerical error tols = dtype_tols(dtype) - if quantized_compute and quantization == "nvfp4": + if quantized_compute and quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): tols = dtype_tols(tex.DType.kFloat4E2M1) elif quantized_compute: tols = dtype_tols(tex.DType.kFloat8E4M3) @@ -2077,6 +2113,8 @@ def test_grouped_linear( pytest.skip("Quantization scheme is not used") if quantization is not None and dtype not in (torch.bfloat16, torch.float16): pytest.skip("Quantized group GEMM is only supported with BF16/FP16") + if quantization == "nvfp4_4over6": + pytest.skip("NVFP4 4over6 grouped quantization is not supported") if single_grouped_bias and not bias: pytest.skip("single_grouped_bias requires bias=True") @@ -2113,6 +2151,7 @@ def test_grouped_linear( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="weight"), requires_grad=weight_requires_grad, ) b_ref, b_test = None, None @@ -2685,6 +2724,7 @@ def test_forward_linear_bias_activation( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="weight"), ) b_ref, b_test = None, None if bias: @@ -2790,6 +2830,7 @@ def test_forward_linear_bias_add( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="weight"), ) b_ref, b_test = None, None if bias: @@ -2903,6 +2944,7 @@ def test_forward_linear_scale_add( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="weight"), ) x2_ref, x2_test = make_reference_and_test_tensors( out_shape, @@ -3185,6 +3227,7 @@ def test_backward_linear_add( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="weight"), ) dy1_ref, dy1_test = make_reference_and_test_tensors( out_shape, @@ -3288,6 +3331,7 @@ def test_backward_linear_scale( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="weight"), ) dy_ref, dy_test = make_reference_and_test_tensors( out_shape, @@ -3504,55 +3548,62 @@ def test_layernorm_mlp( ) norm_w_ref, norm_w_test = make_reference_and_test_tensors( hidden_size, + min=-0.5, + max=0.5, test_dtype=dtype, test_device=device, ) norm_b_ref, norm_b_test = make_reference_and_test_tensors( hidden_size, + min=-0.5, + max=0.5, test_dtype=dtype, test_device=device, ) w1_ref, w1_test = make_reference_and_test_tensors( (ffn_hidden_size, hidden_size), quantization=quantization, + min=0, + max=1 / 64, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="weight"), ) w2_ref, w2_test = make_reference_and_test_tensors( (hidden_size, ffn_hidden_size // 2), + min=0, + max=1 / 64, quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="weight"), ) b1_ref, b1_test, b2_ref, b2_test = None, None, None, None if bias: b1_ref, b1_test = make_reference_and_test_tensors( ffn_hidden_size, + min=-0.5, + max=0.5, test_dtype=dtype, test_device=device, ) b2_ref, b2_test = make_reference_and_test_tensors( hidden_size, + min=-0.5, + max=0.5, test_dtype=dtype, test_device=device, ) dy_ref, dy_test = make_reference_and_test_tensors( in_shape, + min=-0.5, + max=0.5, quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="grad_output"), requires_grad=False, ) - with torch.no_grad(): - for t in (norm_w_ref, norm_w_test, norm_b_ref, norm_b_test): - t -= 0.5 - for t in (w1_ref, w1_test, w2_ref, w2_test): - t *= 1 / 64 - if bias: - for t in (b1_ref, b1_test, b2_ref, b2_test): - t -= 0.5 - for t in (dy_ref, dy_test): - t -= 0.5 # Reference implementation x = x_ref @@ -3686,7 +3737,14 @@ def test_grouped_mlp( pytest.skip("Scaled unary grouped MLP fusion is only supported with MXFP8") if not activation_is_glu and glu_interleave_size is not None: pytest.skip("Unary activations do not use GLU interleaving") - if quantization == "nvfp4" and activation == "scaled_clamped_qgeglu" and bias: + if quantization == "nvfp4_4over6": + pytest.skip("NVFP4 4over6 grouped quantization is not supported") + if ( + with_quantization + and quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6") + and activation == "scaled_clamped_qgeglu" + and bias + ): # TODO: ksivaman: Need to debug numerics for this case. pytest.skip("Bias/dbias not yet supported in NVFP4 fused grouped MLP with GeGLU") fc1_out_features = 2 * hidden_size if activation_is_glu else hidden_size @@ -3726,6 +3784,7 @@ def test_grouped_mlp( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="weight"), ) fc2_w_ref, fc2_w_test = make_reference_and_test_tensors( (hidden_size, hidden_size), @@ -3734,6 +3793,7 @@ def test_grouped_mlp( quantization=quantization, test_dtype=dtype, test_device=device, + quantizer_role=QuantizerRole(tensor_type="weight"), ) fc1_b_ref, fc1_b_test = None, None fc2_b_ref, fc2_b_test = None, None @@ -3918,7 +3978,7 @@ def test_grouped_mlp( # Loose tols for sanity checking tols = {"rtol": 0.125, "atol": 0.25} - if quantization == "nvfp4": + if quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): tols = {"rtol": 0.25, "atol": 0.5} # Check values diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index a718ea2a8..5f82bfcba 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -54,7 +54,7 @@ from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor from transformer_engine.common import recipe import transformer_engine_torch as tex -from utils import ModelConfig, reset_rng_states +from utils import ModelConfig, recipe_id, reset_rng_states, skip_unsupported_backward_override # Only run FP8 tests on supported devices. @@ -138,6 +138,32 @@ def nvfp4_rht_and_2d_quantization(): return nvfp4_recipe +def nvfp4_row_scaled(): + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + row_scaled_activation=True, + backward_override="high_precision", + ) + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + return nvfp4_recipe + + +def nvfp4_4over6(): + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + nvfp4_4over6="all", + ) + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams(fp4_2d_quantization=True) + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + return nvfp4_recipe + + def check_rht_usage(recipe: recipe.Recipe) -> bool: # if using RHT, we can only support bf16 # check fp4_quant_fwd_inp, fp4_quant_fwd_weight, fp4_quant_bwd_grad @@ -171,6 +197,8 @@ def get_nvfp4_inp_supported_dtypes(recipe: recipe.Recipe, dtype: torch.dtype) -> fp8_recipes.append(recipe.DelayedScaling()) if nvfp4_available: fp8_recipes.append(nvfp4_rht_and_2d_quantization()) + fp8_recipes.append(nvfp4_4over6()) + fp8_recipes.append(nvfp4_row_scaled()) use_cutlass_grouped_gemm = [False] # Only enable cutlass grouped gemm on Hopper @@ -627,11 +655,15 @@ def _test_e2e_selective_recompute( @pytest.mark.parametrize("bs", batch_sizes) @pytest.mark.parametrize("model", ["126m"]) @pytest.mark.parametrize("fp8", all_boolean) -@pytest.mark.parametrize("recipe", fp8_recipes) +@pytest.mark.parametrize("recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("fp8_model_params", all_boolean) def test_gpt_selective_activation_recompute(dtype, bs, model, fp8, recipe, fp8_model_params): if fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: pytest.skip("FP8 parameters are not supported in debug mode.") + if fp8 or fp8_model_params: + skip_unsupported_backward_override( + "transformer_layer", recipe, getattr(recipe, "backward_override", None) + ) if fp8 and recipe.nvfp4(): if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype): pytest.skip( @@ -739,7 +771,7 @@ def _test_e2e_full_recompute( @pytest.mark.parametrize("bs", batch_sizes) @pytest.mark.parametrize("model", ["126m"]) @pytest.mark.parametrize("fp8", all_boolean) -@pytest.mark.parametrize("recipe", fp8_recipes) +@pytest.mark.parametrize("recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("fp8_model_params", all_boolean) @pytest.mark.parametrize("use_reentrant", all_boolean) def test_gpt_full_activation_recompute( @@ -747,6 +779,10 @@ def test_gpt_full_activation_recompute( ): if fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: pytest.skip("FP8 parameters are not supported in debug mode.") + if fp8 or fp8_model_params: + skip_unsupported_backward_override( + "transformer_layer", recipe, getattr(recipe, "backward_override", None) + ) if fp8 and recipe.nvfp4(): if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype): pytest.skip( @@ -1324,7 +1360,7 @@ def test_linear_accuracy_delay_wgrad_compute(dtype, bs, model, bias, fuse_wgrad_ @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("model", ["small"]) -@pytest.mark.parametrize("recipe", fp8_recipes + [None]) +@pytest.mark.parametrize("recipe", fp8_recipes + [None], ids=recipe_id) def test_linear_accuracy_save_original_input(dtype, model, recipe): bs = 1 fuse_wgrad_accumulation = True @@ -1333,6 +1369,7 @@ def test_linear_accuracy_save_original_input(dtype, model, recipe): if fp8 and recipe.delayed(): pytest.skip("DelayedScaling recipe is not supported with save_original_input") + skip_unsupported_backward_override("linear", recipe, getattr(recipe, "backward_override", None)) config = model_configs[model] if config.max_seqlen_q % 16 != 0 and fp8: @@ -1894,7 +1931,7 @@ def _test_grouped_linear_accuracy( @pytest.mark.parametrize("num_gemms", [3, 6]) @pytest.mark.parametrize("bs", batch_sizes) @pytest.mark.parametrize("model", ["126m"]) -@pytest.mark.parametrize("recipe", fp8_recipes + [None]) +@pytest.mark.parametrize("recipe", fp8_recipes + [None], ids=recipe_id) @pytest.mark.parametrize("fp8_model_params", all_boolean) @pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean) @pytest.mark.parametrize("bias", all_boolean) @@ -1917,6 +1954,9 @@ def test_grouped_linear_accuracy( pytest.skip("FP8 parameters are not supported in debug mode.") if NVTE_TEST_NVINSPECT_ENABLED and delay_wgrad_compute: pytest.skip("Delayed wgrad compute is not supported in debug mode.") + skip_unsupported_backward_override( + "grouped_linear", recipe, getattr(recipe, "backward_override", None) + ) config = model_configs[model] if config.max_seqlen_q % 16 != 0 and fp8: @@ -2037,7 +2077,7 @@ def test_grouped_linear_accuracy_cutlass( @pytest.mark.parametrize("num_gemms", [3]) @pytest.mark.parametrize("bs", [1]) @pytest.mark.parametrize("model", ["126m"]) -@pytest.mark.parametrize("recipe", fp8_recipes + [None]) +@pytest.mark.parametrize("recipe", fp8_recipes + [None], ids=recipe_id) @pytest.mark.parametrize("fp8_model_params", [False]) @pytest.mark.parametrize("fuse_wgrad_accumulation", [True]) @pytest.mark.parametrize("bias", [False]) @@ -2061,6 +2101,9 @@ def test_grouped_linear_accuracy_save_original_input( pytest.skip("DelayedScaling recipe is not supported with save_original_input") if NVTE_TEST_NVINSPECT_ENABLED and delay_wgrad_compute: pytest.skip("Delayed wgrad compute is not supported in debug mode.") + skip_unsupported_backward_override( + "grouped_linear", recipe, getattr(recipe, "backward_override", None) + ) config = model_configs[model] if config.max_seqlen_q % 16 != 0 and fp8: @@ -2139,7 +2182,7 @@ def test_grouped_linear_accuracy_save_original_input( torch.testing.assert_close(o, o_ref, rtol=0, atol=0) -@pytest.mark.parametrize("recipe", fp8_recipes + [None]) +@pytest.mark.parametrize("recipe", fp8_recipes + [None], ids=recipe_id) def test_grouped_linear_accuracy_single_gemm(recipe): """Split the tests to save CI time""" test_grouped_linear_accuracy( @@ -2253,7 +2296,7 @@ def _generate_random_numbers(n, total_sum): @pytest.mark.parametrize("bs", batch_sizes) @pytest.mark.parametrize("model", ["126m"]) @pytest.mark.parametrize("fp8", [True]) -@pytest.mark.parametrize("recipe", fp8_recipes) +@pytest.mark.parametrize("recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("fp8_model_params", all_boolean) def test_padding_grouped_linear_accuracy( dtype, @@ -2267,6 +2310,9 @@ def test_padding_grouped_linear_accuracy( ): if fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: pytest.skip("FP8 parameters are not supported in debug mode.") + skip_unsupported_backward_override( + "grouped_linear", recipe, getattr(recipe, "backward_override", None) + ) config = model_configs[model] if config.max_seqlen_q % 16 != 0 and fp8: @@ -2328,7 +2374,7 @@ def test_padding_grouped_linear_accuracy( @pytest.mark.parametrize("bs", [1]) @pytest.mark.parametrize("model", ["126m"]) @pytest.mark.parametrize("fp8", [True]) -@pytest.mark.parametrize("recipe", fp8_recipes) +@pytest.mark.parametrize("recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("fp8_model_params", [False]) def test_padding_grouped_linear_accuracy_save_original_input( dtype, @@ -2344,6 +2390,9 @@ def test_padding_grouped_linear_accuracy_save_original_input( pytest.skip("FP8 parameters are not supported in debug mode.") if fp8 and recipe.delayed(): pytest.skip("DelayedScaling recipe is not supported with save_original_input") + skip_unsupported_backward_override( + "grouped_linear", recipe, getattr(recipe, "backward_override", None) + ) config = model_configs[model] if config.max_seqlen_q % 16 != 0 and fp8: @@ -2559,10 +2608,13 @@ def _test_gpt_fp8_parameters(bs, dtype, config, fp8_model_params, recipe): @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("bs", batch_sizes) @pytest.mark.parametrize("model", ["126m"]) -@pytest.mark.parametrize("recipe", fp8_recipes) +@pytest.mark.parametrize("recipe", fp8_recipes, ids=recipe_id) def test_gpt_fp8_parameters(dtype, bs, model, recipe): if NVTE_TEST_NVINSPECT_ENABLED: pytest.skip("FP8 parameters are not supported in debug mode.") + skip_unsupported_backward_override( + "transformer_layer", recipe, getattr(recipe, "backward_override", None) + ) if recipe.nvfp4(): if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype): diff --git a/tests/pytorch/test_quantized_tensor.py b/tests/pytorch/test_quantized_tensor.py index 119914fbc..c5161349e 100644 --- a/tests/pytorch/test_quantized_tensor.py +++ b/tests/pytorch/test_quantized_tensor.py @@ -28,7 +28,7 @@ import transformer_engine_torch as tex from references.ref_per_tensor_cs import ref_per_tensor_cs_cast -from utils import assert_close, quantization_tols +from utils import assert_close # PyTorch tensor dtypes _dtypes: List[torch.dtype] = [torch.float32, torch.float16, torch.bfloat16] @@ -69,6 +69,8 @@ def _to_list(x: Union[Iterable, Any]) -> List: _quantization_list.append("mxfp8") if nvfp4_available: _quantization_list.append("nvfp4") + _quantization_list.append("nvfp4_row_scaled") + _quantization_list.append("nvfp4_4over6") # delayed scaling @@ -163,13 +165,17 @@ def make_reference_and_test_tensors( test = quantizer(test) elif quantization == "mxfp8": test = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3)(test) - elif quantization == "nvfp4": + elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): + row_scaled_nvfp4 = quantization == "nvfp4_row_scaled" test = NVFP4Quantizer( + columnwise=not row_scaled_nvfp4, with_rht=False, with_post_rht_amax=False, with_2d_quantization=False, stochastic_rounding=False, + row_scaled_nvfp4=row_scaled_nvfp4, with_random_sign_mask=False, + nvfp4_use_4over6=(quantization == "nvfp4_4over6"), )(test) else: raise ValueError(f"Unsupported quantization scheme ({quantization})") @@ -785,13 +791,16 @@ def test_update_nd_tensor( ) elif quantization == "mxfp8": quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3) - elif quantization in ("nvfp4", "nvfp4_2d"): + elif quantization in ("nvfp4", "nvfp4_2d", "nvfp4_row_scaled", "nvfp4_4over6"): + row_scaled_nvfp4 = quantization == "nvfp4_row_scaled" quantizer = NVFP4Quantizer( rowwise=True, - columnwise=True, + columnwise=not row_scaled_nvfp4, with_rht=False, with_post_rht_amax=False, with_2d_quantization=(quantization == "nvfp4_2d"), + row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=(quantization == "nvfp4_4over6"), ) quantization = "nvfp4" else: @@ -806,9 +815,9 @@ def test_update_nd_tensor( q_x.copy_(x_new) # Check results + q_ref = quantizer(x_new) assert q_x.shape == torch.Size(shape) - tols = quantization_tols(quantization) - assert_close(q_x, x_new, **tols) + assert_close(q_x, q_ref, rtol=0, atol=0) @pytest.mark.skipif(not mxfp8_available, reason=reason_for_no_mxfp8) diff --git a/tests/pytorch/test_recipe.py b/tests/pytorch/test_recipe.py index 5f5221af7..9a14cee7f 100644 --- a/tests/pytorch/test_recipe.py +++ b/tests/pytorch/test_recipe.py @@ -26,6 +26,7 @@ from transformer_engine.pytorch.quantization import ( FP8GlobalStateManager, NVFP4BlockScalingRecipeState, + QuantizerRole, _amax_and_scale_update, ) import transformer_engine.pytorch.ops as te_ops @@ -514,8 +515,52 @@ def test_quantizer_update(self, module_class): @pytest.mark.skipif(not fp4_available, reason=reason_for_no_fp4) -def test_nvfp4_row_scaled_quantizer_roles(): - recipe = NVFP4BlockScaling(row_scaled_activation=True) +@pytest.mark.parametrize( + "nvfp4_4over6", + ["none", "weights", "activations", "all"], + ids=["disabled", "weights", "activations", "all"], +) +@pytest.mark.parametrize( + "nvfp4_4over6_e4m3_use_256", + ["none", "weights", "activations", "all"], + ids=["e4m3_448", "e4m3_256_weights", "e4m3_256_activations", "e4m3_256_all"], +) +@pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"]) +def test_nvfp4_row_scaled_quantizer_roles( + nvfp4_4over6, nvfp4_4over6_e4m3_use_256, nvfp4_4over6_err_mode +): + recipe = NVFP4BlockScaling( + disable_rht=True, + disable_2d_quantization=True, + nvfp4_4over6=nvfp4_4over6, + nvfp4_4over6_e4m3_use_256=nvfp4_4over6_e4m3_use_256, + nvfp4_4over6_err_mode=nvfp4_4over6_err_mode, + row_scaled_activation=True, + ) + + def expected_use_4over6(tensor_type): + if tensor_type in ("grad_output", "grad_input"): + return False + if nvfp4_4over6 == "all": + return True + if nvfp4_4over6 == "weights": + return tensor_type == "weight" + if nvfp4_4over6 == "activations": + return tensor_type != "weight" + return False + + def expected_e4m3_max(tensor_type): + if not expected_use_4over6(tensor_type): + return 448 + if nvfp4_4over6_e4m3_use_256 == "all": + return 256 + if nvfp4_4over6_e4m3_use_256 == "weights": + if tensor_type == "weight": + return 256 + if nvfp4_4over6_e4m3_use_256 == "activations": + if tensor_type != "weight": + return 256 + return 448 forward_quantizers = NVFP4BlockScalingRecipeState( recipe, @@ -523,20 +568,85 @@ def test_nvfp4_row_scaled_quantizer_roles(): num_quantizers=3, ).make_quantizers() assert [q.row_scaled_nvfp4 for q in forward_quantizers] == [True, False, True] + assert [q.stochastic_rounding for q in forward_quantizers] == [False, False, False] + assert [q.with_rht for q in forward_quantizers] == [False, False, False] + assert [q.nvfp4_use_4over6 for q in forward_quantizers] == [ + expected_use_4over6(tensor_type) for tensor_type in ("input", "weight", "output") + ] + assert [q.nvfp4_e4m3_max for q in forward_quantizers] == [ + expected_e4m3_max(tensor_type) for tensor_type in ("input", "weight", "output") + ] + assert [q.nvfp4_4over6_err_mode for q in forward_quantizers] == [nvfp4_4over6_err_mode] * 3 assert not forward_quantizers[0].is_quantizable(torch.empty(16, 16)) assert forward_quantizers[1].is_quantizable(torch.empty(16, 16)) + role_quantizers = NVFP4BlockScalingRecipeState( + recipe, + mode="forward", + num_quantizers=4, + roles=[ + QuantizerRole(module_type="linear", tensor_type="weight"), + QuantizerRole(module_type="linear", tensor_type="input"), + QuantizerRole(module_type="linear", tensor_type="output"), + None, + ], + ).make_quantizers() + assert [q.row_scaled_nvfp4 for q in role_quantizers] == [False, True, True, True] + assert [q.nvfp4_use_4over6 for q in role_quantizers] == [ + expected_use_4over6(tensor_type) for tensor_type in ("weight", "input", "output", "input") + ] + assert [q.nvfp4_e4m3_max for q in role_quantizers] == [ + expected_e4m3_max(tensor_type) for tensor_type in ("weight", "input", "output", "input") + ] + assert [q.nvfp4_4over6_err_mode for q in role_quantizers] == [nvfp4_4over6_err_mode] * 4 + backward_quantizers = NVFP4BlockScalingRecipeState( recipe, mode="backward", num_quantizers=2, + roles=[ + QuantizerRole(module_type="linear", tensor_type="grad_output"), + QuantizerRole(module_type="linear", tensor_type="grad_input"), + ], ).make_quantizers() assert [q.row_scaled_nvfp4 for q in backward_quantizers] == [False, False] + assert [q.nvfp4_use_4over6 for q in backward_quantizers] == [False, False] + assert [q.nvfp4_e4m3_max for q in backward_quantizers] == [448, 448] + assert [q.nvfp4_4over6_err_mode for q in backward_quantizers] == [nvfp4_4over6_err_mode] * 2 + assert [q.stochastic_rounding for q in backward_quantizers] == [True, True] + assert [q.with_rht for q in backward_quantizers] == [False, False] + + backward_operand_quantizers = NVFP4BlockScalingRecipeState( + recipe, + mode="backward", + num_quantizers=4, + roles=[ + QuantizerRole(module_type="linear", tensor_type="input"), + QuantizerRole(module_type="linear", tensor_type="weight"), + QuantizerRole(module_type="linear", tensor_type="grad_output"), + QuantizerRole(module_type="linear", tensor_type="grad_input"), + ], + ).make_quantizers() + assert [q.nvfp4_use_4over6 for q in backward_operand_quantizers] == [ + expected_use_4over6(tensor_type) + for tensor_type in ("input", "weight", "grad_output", "grad_input") + ] + assert [q.nvfp4_e4m3_max for q in backward_operand_quantizers] == [ + expected_e4m3_max(tensor_type) + for tensor_type in ("input", "weight", "grad_output", "grad_input") + ] + assert [q.stochastic_rounding for q in backward_operand_quantizers] == [ + False, + False, + True, + True, + ] @pytest.mark.skipif(not fp4_available, reason=reason_for_no_fp4) @pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16], ids=str) @pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"]) +@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"]) @pytest.mark.parametrize( "M, N", [ @@ -552,24 +662,30 @@ def test_nvfp4_row_scaled_quantizer_roles(): (8192, 8192), ], ) -def test_fp4_dequantize(dtype, row_scaled_nvfp4, M, N): +def test_fp4_dequantize(dtype, row_scaled_nvfp4, use_4over6, M, N): q = NVFP4Quantizer( columnwise=not row_scaled_nvfp4, row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=use_4over6, ) a = torch.rand((M, N)).cuda().to(dtype=dtype) starting_tensor = q(a) assert starting_tensor._row_scaled_nvfp4 == row_scaled_nvfp4 + assert starting_tensor._nvfp4_use_4over6 == use_4over6 assert starting_tensor._amax_rowwise.numel() == (M if row_scaled_nvfp4 else 1) dequantized_tensor = starting_tensor.dequantize() new_tensor = q(dequantized_tensor) assert new_tensor._row_scaled_nvfp4 == row_scaled_nvfp4 + assert new_tensor._nvfp4_use_4over6 == use_4over6 assert new_tensor._amax_rowwise.numel() == (M if row_scaled_nvfp4 else 1) - torch.testing.assert_close( - new_tensor._rowwise_data, - starting_tensor._rowwise_data, - rtol=0, - atol=0, - ) + # 4over6 can re-encode a dequantized block with the alternate 4/6 scale + # choice while preserving the dequantized values. + if not use_4over6: + torch.testing.assert_close( + new_tensor._rowwise_data, + starting_tensor._rowwise_data, + rtol=0, + atol=0, + ) new_dequantized_tensor = new_tensor.dequantize() torch.testing.assert_close(dequantized_tensor, new_dequantized_tensor) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index c811342df..27eafbecd 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -95,27 +95,43 @@ def nvfp4_vanilla(): def nvfp4_row_scaled(): - nvfp4_recipe = recipe.NVFP4BlockScaling(row_scaled_activation=True) + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + row_scaled_activation=True, + backward_override="dequantized", + ) nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() return nvfp4_recipe +def nvfp4_4over6(): + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + nvfp4_4over6="all", + ) + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams(fp4_2d_quantization=True) + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + return nvfp4_recipe + + fp8_recipes = [] if mxfp8_available: fp8_recipes.append(recipe.MXFP8BlockScaling()) if nvfp4_available: fp8_recipes.append(nvfp4_vanilla()) # TODO: fix check for this + fp8_recipes.append(nvfp4_4over6()) if fp8_block_scaling_available: fp8_recipes.append(recipe.Float8BlockScaling()) if fp8_available: fp8_recipes.append(recipe.Float8CurrentScaling()) fp8_recipes.append(recipe.DelayedScaling()) fp8_recipes.append(None) -fp8_recipes_with_row_scaled = fp8_recipes.copy() -if nvfp4_available: - fp8_recipes_with_row_scaled.insert(-1, nvfp4_row_scaled()) param_types = [torch.float32, torch.float16] if is_bf16_available(): # bf16 requires sm_80 or higher @@ -415,7 +431,11 @@ def test_sanity_normalization_amp(dtype, model, skip_wgrad, skip_dgrad, normaliz @pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes_with_row_scaled, ids=recipe_id) +@pytest.mark.parametrize( + "fp8_recipe", + fp8_recipes + ([nvfp4_row_scaled()] if nvfp4_available else []), + ids=recipe_id, +) @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) @pytest.mark.parametrize("model", ["small", "weird"]) @pytest.mark.parametrize("skip_wgrad", all_boolean) @@ -463,7 +483,11 @@ def test_sanity_layernorm_linear( @pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes_with_row_scaled, ids=recipe_id) +@pytest.mark.parametrize( + "fp8_recipe", + fp8_recipes + ([nvfp4_row_scaled()] if nvfp4_available else []), + ids=recipe_id, +) @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) @pytest.mark.parametrize("model", ["small", "weird"]) @pytest.mark.parametrize("skip_wgrad", all_boolean) @@ -501,7 +525,11 @@ def test_sanity_linear( @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("bs", batch_sizes_with_zero) @pytest.mark.parametrize("model", ["small", "weird"]) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes_with_row_scaled, ids=recipe_id) +@pytest.mark.parametrize( + "fp8_recipe", + fp8_recipes + ([nvfp4_row_scaled()] if nvfp4_available else []), + ids=recipe_id, +) @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) @pytest.mark.parametrize("fp8_model_params", all_boolean) @pytest.mark.parametrize("use_bias", all_boolean) @@ -542,7 +570,11 @@ def test_sanity_linear_with_zero_tokens( @pytest.mark.parametrize("dtype", param_types) @pytest.mark.parametrize("bs", batch_sizes_with_zero) @pytest.mark.parametrize("model", ["small", "weird"]) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes_with_row_scaled, ids=recipe_id) +@pytest.mark.parametrize( + "fp8_recipe", + fp8_recipes + ([nvfp4_row_scaled()] if nvfp4_available else []), + ids=recipe_id, +) @pytest.mark.parametrize("backward_override", [None, "high_precision", "dequantized"]) @pytest.mark.parametrize("fp8_model_params", all_boolean) @pytest.mark.parametrize("use_bias", all_boolean) @@ -621,7 +653,7 @@ def test_sanity_grouped_linear( @pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("model", ["small", "weird"]) @pytest.mark.parametrize("skip_wgrad", all_boolean) @pytest.mark.parametrize("zero_centered_gamma", all_boolean) @@ -671,7 +703,7 @@ def test_sanity_layernorm_mlp( @pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("model", ["small"]) @pytest.mark.parametrize("skip_wgrad", all_boolean) @pytest.mark.parametrize("bias", all_boolean) @@ -744,7 +776,7 @@ def test_sanity_gpt_126m(): @pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("model", ["small"]) @pytest.mark.parametrize("skip_wgrad", all_boolean) @pytest.mark.parametrize("normalization", all_normalizations) @@ -800,7 +832,7 @@ def test_sanity_bert_126m(): @pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("model", ["small"]) @pytest.mark.parametrize("skip_wgrad", all_boolean) @pytest.mark.parametrize("normalization", all_normalizations) @@ -856,7 +888,7 @@ def test_sanity_T5_126m(): @pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("model", ["small"]) @pytest.mark.parametrize("skip_wgrad", all_boolean) def test_sanity_amp_and_nvfuser(dtype, fp8_recipe, model, skip_wgrad): @@ -889,7 +921,7 @@ def test_sanity_amp_and_nvfuser(dtype, fp8_recipe, model, skip_wgrad): @pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("model", ["small"]) def test_sanity_drop_path(dtype, fp8_recipe, model): config = model_configs[model] @@ -924,7 +956,7 @@ def test_sanity_drop_path(dtype, fp8_recipe, model): @pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("model", ["small"]) @pytest.mark.parametrize("skip_wgrad", all_boolean) def test_sanity_fused_qkv_params(dtype, fp8_recipe, model, skip_wgrad): @@ -960,7 +992,7 @@ def test_sanity_fused_qkv_params(dtype, fp8_recipe, model, skip_wgrad): @pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("fp8_recipe", fp8_recipes) +@pytest.mark.parametrize("fp8_recipe", fp8_recipes, ids=recipe_id) @pytest.mark.parametrize("model", ["small"]) @pytest.mark.parametrize("skip_wgrad", all_boolean) def test_sanity_gradient_accumulation_fusion(dtype, fp8_recipe, model, skip_wgrad): diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index 51f72b1e5..137e5f5a7 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/tests/pytorch/test_torch_compile.py @@ -39,6 +39,33 @@ fp8_block_scaling_available = is_fp8_block_scaling_available() nvfp4_available = is_nvfp4_available() + +def nvfp4_row_scaled(): + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + row_scaled_activation=True, + backward_override="dequantized", + ) + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + return nvfp4_recipe + + +def nvfp4_4over6(): + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + nvfp4_4over6="all", + ) + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams(fp4_2d_quantization=True) + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + return nvfp4_recipe + + _all_recipes: list = [] if fp8_available: _all_recipes.append(recipe.Float8CurrentScaling()) @@ -48,7 +75,8 @@ _all_recipes.append(recipe.MXFP8BlockScaling()) if nvfp4_available: _all_recipes.append(recipe.NVFP4BlockScaling()) - _all_recipes.append(recipe.NVFP4BlockScaling(row_scaled_activation=True)) + _all_recipes.append(nvfp4_4over6()) + _all_recipes.append(nvfp4_row_scaled()) # --------------------------------------------------------------------------- @@ -97,8 +125,19 @@ def __fx_repr__(self): def _make_qfactory(tag: str): """Return a qfactory that produces ToyQuantizer instances tagged with *tag*.""" + quantizers = { + role: ToyQuantizer(tag=f"{tag}:{role}") + for role in ( + "linear_input", + "linear_weight", + "linear_output", + "linear_grad_output", + "linear_grad_input", + ) + } + def qfactory(role: str): - return ToyQuantizer(tag=f"{tag}:{role}") + return quantizers[role] return qfactory diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 2ee18aaf5..19cc118a9 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -118,7 +118,7 @@ def quantization_tols(name: str) -> dict[str, float]: "mxfp8_block_scaling", ): return dtype_tols(tex.DType.kFloat8E4M3) - if name in ("nvfp4", "nvfp4_row_scaled"): + if name in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): return dtype_tols(tex.DType.kFloat4E2M1) raise ValueError(f"Unsupported quantization scheme ({name})") @@ -145,21 +145,17 @@ def make_recipe(name: Optional[str], **recipe_kwargs: Any) -> Optional[Recipe]: ) if name == "fp8_block_scaling": return transformer_engine.common.recipe.Float8BlockScaling(**recipe_kwargs) - if name == "nvfp4": - return transformer_engine.common.recipe.NVFP4BlockScaling( - disable_rht=True, - disable_stochastic_rounding=True, - disable_2d_quantization=True, - **recipe_kwargs, - ) - if name == "nvfp4_row_scaled": - return transformer_engine.common.recipe.NVFP4BlockScaling( - disable_rht=True, - disable_stochastic_rounding=True, - disable_2d_quantization=True, - row_scaled_activation=True, - **recipe_kwargs, - ) + if name in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): + use_4over6 = name == "nvfp4_4over6" + kwargs = { + "disable_rht": True, + "disable_stochastic_rounding": True, + "disable_2d_quantization": not use_4over6, + "row_scaled_activation": name == "nvfp4_row_scaled", + "nvfp4_4over6": "all" if use_4over6 else "none", + } + kwargs.update(recipe_kwargs) + return transformer_engine.common.recipe.NVFP4BlockScaling(**kwargs) raise ValueError(f"Unsupported quantization scheme ({name})") @@ -167,6 +163,10 @@ def recipe_id(recipe: Optional[Recipe]) -> str: """Readable pytest id for a quantization recipe.""" if not isinstance(recipe, Recipe): return "None" + if recipe.nvfp4() and recipe.row_scaled_activation and recipe.nvfp4_4over6 != "none": + return "NVFP4RowScaled4Over6BlockScaling" + if recipe.nvfp4() and recipe.nvfp4_4over6 != "none": + return "NVFP44Over6BlockScaling" if recipe.nvfp4() and recipe.row_scaled_activation: return "NVFP4RowScaledBlockScaling" return type(recipe).__name__ @@ -185,6 +185,13 @@ def skip_unsupported_backward_override( and backward_override is None ): pytest.skip("Row-scaled NVFP4 does not support default quantized backward.") + if ( + quant_recipe is not None + and quant_recipe.nvfp4() + and quant_recipe.nvfp4_4over6 != "none" + and layer_type == "grouped_linear" + ): + pytest.skip("NVFP4 4over6 currently does not support grouped quantization.") if backward_override is None: return if quant_recipe is None and backward_override is not None: diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 123362ce1..316243c97 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -21,6 +21,7 @@ #include "../mxfp8/group_quantize_mxfp8.cuh" #include "../mxfp8/quantize_mxfp8.cuh" #include "../nvfp4/group_quantize_transpose_nvfp4.cuh" +#include "../nvfp4/quantize_4over6_nvfp4.cuh" #include "../nvfp4/quantize_transpose_nvfp4.cuh" namespace transformer_engine { @@ -101,6 +102,11 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, int32_t cols = input_tensor->flat_last_dim(); auto dtype = input_tensor->dtype(); const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; + const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; + NVTE_CHECK(nvfp4_use_4over6 || output_tensor->nvfp4_e4m3_max == 448, + "Non-4over6 NVFP4 quantization requires E4M3 max 448."); + NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, + "NVFP4 4over6 quantization does not support stochastic rounding."); if (row_scaled_nvfp4) { NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, "Row-scaled NVFP4 quantization does not support 2D quantization."); @@ -112,7 +118,15 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, (cols % 32 == 0) && output_tensor->has_data(); // Launch NVFP4 quantize kernel - if (use_optimized_kernel) { + if (nvfp4_use_4over6) { + if (quant_config_cpp.nvfp4_2d_quantization) { + nvfp4::quantize_4over6( + *input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + } else { + nvfp4::quantize_4over6( + *input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + } + } else if (use_optimized_kernel) { if (quant_config_cpp.nvfp4_2d_quantization) { nvfp4::quantize_transpose( *input_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); @@ -249,13 +263,26 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens int32_t rows = grad_tensor->flat_first_dim(); int32_t cols = grad_tensor->flat_last_dim(); auto dtype = grad_tensor->dtype(); + const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; + NVTE_CHECK(nvfp4_use_4over6 || output_tensor->nvfp4_e4m3_max == 448, + "Non-4over6 NVFP4 quantization requires E4M3 max 448."); + NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding, + "NVFP4 4over6 quantization does not support stochastic rounding."); NVTE_CHECK(!output_tensor->row_scaled_nvfp4, "Backward NVFP4 quantization does not support row-scaled outputs."); bool use_optimized_kernel = (dtype == DType::kBFloat16) && (rows % 32 == 0) && (cols % 32 == 0) && output_tensor->has_data(); // Launch NVFP4 quantize kernel - if (use_optimized_kernel) { + if (nvfp4_use_4over6) { + if (quant_config_cpp.nvfp4_2d_quantization) { + nvfp4::quantize_4over6( + *grad_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + } else { + nvfp4::quantize_4over6( + *grad_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); + } + } else if (use_optimized_kernel) { if (quant_config_cpp.nvfp4_2d_quantization) { nvfp4::quantize_transpose( *grad_tensor, noop_tensor, output_tensor, &quant_config_cpp, stream); @@ -277,7 +304,8 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens /*use_stochastic_rounding=*/quant_config_cpp.stochastic_rounding, /*rng_state=*/quant_config_cpp.rng_state, /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization, - /*row_scaled_nvfp4=*/false, /*noop_tensor=*/noop_tensor->data, + /*row_scaled_nvfp4=*/false, + /*noop_tensor=*/noop_tensor->data, /*stream=*/stream); } break; @@ -372,8 +400,15 @@ void group_quantize_fwd_host_aware_helper(const NVTETensor input, NVTETensor *ou int32_t cols = input_tensor->flat_last_dim(); auto dtype = input_tensor->dtype(); + const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; + for (const auto *output_tensor : output_tensors) { + NVTE_CHECK(nvfp4_use_4over6 || output_tensor->nvfp4_e4m3_max == 448, + "Non-4over6 NVFP4 quantization requires E4M3 max 448."); + } NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization, "2D quantization is not supported for group quantize."); + NVTE_CHECK(!nvfp4_use_4over6, + "NVFP4 4over6 quantization is not supported for group quantize."); // Launch NVFP4 group quantize kernel nvfp4::group_quantize_transpose( diff --git a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh index 792b068cb..3820430d5 100644 --- a/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/core_nvfp4.cuh @@ -75,10 +75,14 @@ namespace core { #if FP4_TYPE_SUPPORTED using namespace ptx; -// Compute the global encode scale factor for a given global amax +// Compute the global encode scale factor for a given global amax. +// NVFP4 uses the full E4M3 range by default. Some 4over6 tensors dispatch +// E4M3_MAX=256 to leave room for map-to-4 scale expansion. +template __device__ __forceinline__ float compute_global_encode_scaling_factor_FP4(const float global_amax) { using namespace detail; - constexpr float fp8_max = TypeExtrema::max; // 448.0f; + static_assert(E4M3_MAX == 448 || E4M3_MAX == 256, "Unsupported NVFP4 E4M3 max."); + constexpr float fp8_max = static_cast(E4M3_MAX); constexpr float fp4_max = TypeExtrema::max; // 6.0f; float global_encode_scale = fp8_max * fp4_max / global_amax; // If scale is infinity, return max value of float32 diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index d549a050e..faf3c58ad 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -31,12 +31,11 @@ namespace dispatch { namespace nvfp4 { namespace dequantize_kernel { #if FP4_TYPE_SUPPORTED -template +template __global__ void __launch_bounds__(512) dequantize_fp4_kernel(const void *const input, OType *output, const fp8e4m3 *const scales, - const float *const tensor_amax, const bool row_scaled_nvfp4, - const size_t N, const size_t M, const size_t scale_stride, - const size_t num_scale_tiles_X) { + const float *const tensor_amax, const size_t N, const size_t M, + const size_t scale_stride, const size_t num_scale_tiles_X) { const size_t thread_idx = blockIdx.x * blockDim.x + threadIdx.x; const size_t x = thread_idx % M; const size_t y = thread_idx / M; @@ -64,8 +63,9 @@ __global__ void __launch_bounds__(512) fp4vec value; value.vec = input_vectorized[my_index]; fp8e4m3 scale = scales[my_scale_index]; - float amax = row_scaled_nvfp4 ? tensor_amax[y] : tensor_amax[0]; - constexpr float factor_inv = 1.0 / (6.0 * 448.0); + float amax = ROW_SCALED_NVFP4 ? tensor_amax[y] : tensor_amax[0]; + static_assert(E4M3_MAX == 448 || E4M3_MAX == 256, "Unsupported NVFP4 E4M3 max."); + constexpr float factor_inv = 1.0f / (6.0f * static_cast(E4M3_MAX)); float final_scale = static_cast(scale) * amax * factor_inv; #pragma unroll for (int i = 0; i < 4; i++) { @@ -92,6 +92,7 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) const bool with_gemm_swizzled_scales = input.with_gemm_swizzled_scales; const bool row_scaled_nvfp4 = input.row_scaled_nvfp4; + const int e4m3_max = input.nvfp4_e4m3_max; constexpr int FP4_BLOCK_SIZE = 16; const size_t N = input.flat_first_dim(); @@ -112,14 +113,25 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) output->data.dtype, OType, TRANSFORMER_ENGINE_SWITCH_CONDITION( with_gemm_swizzled_scales, WITH_GEMM_SWIZZLED_SCALES, - - dequantize_fp4_kernel<<>>( - input.data.dptr, reinterpret_cast(output->data.dptr), - reinterpret_cast(input.scale_inv.dptr), - reinterpret_cast(input.amax.dptr), row_scaled_nvfp4, N, Mread, - input.scale_inv.shape.back(), - num_scale_tiles_X);); // NOLINT(*) - ); // NOLINT(*) + TRANSFORMER_ENGINE_SWITCH_CONDITION( + row_scaled_nvfp4, ROW_SCALED_NVFP4, + if (e4m3_max == 256) { + dequantize_fp4_kernel + <<>>( + input.data.dptr, reinterpret_cast(output->data.dptr), + reinterpret_cast(input.scale_inv.dptr), + reinterpret_cast(input.amax.dptr), N, Mread, + input.scale_inv.shape.back(), num_scale_tiles_X); + } else { + NVTE_CHECK(e4m3_max == 448, "Unsupported NVFP4 E4M3 max (got ", e4m3_max, ")"); + dequantize_fp4_kernel + <<>>( + input.data.dptr, reinterpret_cast(output->data.dptr), + reinterpret_cast(input.scale_inv.dptr), + reinterpret_cast(input.amax.dptr), N, Mread, + input.scale_inv.shape.back(), num_scale_tiles_X); + });); // NOLINT(*) + ); // NOLINT(*) NVTE_CHECK_CUDA(cudaGetLastError()); #else NVTE_ERROR("CUDA 12.8 or higher is needed for FP4 calculation!"); diff --git a/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh new file mode 100644 index 000000000..b6057370d --- /dev/null +++ b/transformer_engine/common/cast/nvfp4/quantize_4over6_nvfp4.cuh @@ -0,0 +1,668 @@ +/************************************************************************* + * Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * + * See LICENSE for license information. + ************************************************************************/ + +/*! \file quantize_4over6_nvfp4.cuh + * \brief Dedicated kernels for NVFP4 4over6 quantization. + * + * Four Over Six evaluates two TE-style NVFP4 encodings for every 1x16 + * quantization group. The map-to-6 candidate uses the normal scale. The + * map-to-4 candidate expands the E4M3 block scale by 1.5x so FP4 value 4 + * reaches the same range that FP4 value 6 reaches in the normal encoding. + * The selected candidate is the one with lower configured dequantization + * error; ties select map-to-6. The quantized candidates, dequantized values, + * and errors are kept in registers, matching the structure of the official + * Four Over Six implementation. + */ + +#ifndef TRANSFORMER_ENGINE_QUANTIZE_4OVER6_NVFP4_CUH_ +#define TRANSFORMER_ENGINE_QUANTIZE_4OVER6_NVFP4_CUH_ + +#include +#include +#include +#include +#include + +#include + +#include "../../common.h" +#include "../../util/math.h" +#include "../../utils.cuh" +#include "core_nvfp4.cuh" + +namespace transformer_engine { +namespace dispatch { +namespace nvfp4 { + +#if FP4_TYPE_SUPPORTED + +#define TRANSFORMER_ENGINE_NVFP4_4OVER6_MODE_SWITCH(MODE, MODE_CONST, ...) \ + switch (MODE) { \ + case kNVTENVFP44Over6MinMAE: { \ + constexpr NVTENVFP44Over6Mode MODE_CONST = kNVTENVFP44Over6MinMAE; \ + { __VA_ARGS__ } \ + } break; \ + case kNVTENVFP44Over6MinMSE: { \ + constexpr NVTENVFP44Over6Mode MODE_CONST = kNVTENVFP44Over6MinMSE; \ + { __VA_ARGS__ } \ + } break; \ + default: { \ + NVTE_ERROR("Unsupported NVFP4 4over6 mode."); \ + } \ + } + +#define TRANSFORMER_ENGINE_NVFP4_4OVER6_E4M3_MAX_SWITCH(E4M3_MAX_VALUE, E4M3_MAX_CONST, ...) \ + if ((E4M3_MAX_VALUE) == 256) { \ + constexpr int E4M3_MAX_CONST = 256; \ + { __VA_ARGS__ } \ + } else { \ + NVTE_CHECK((E4M3_MAX_VALUE) == 448, "Unsupported NVFP4 E4M3 max."); \ + constexpr int E4M3_MAX_CONST = 448; \ + { __VA_ARGS__ } \ + } + +namespace quantize_4over6_kernel { + +constexpr int kThreads = 128; +constexpr int kWarpThreads = 32; +constexpr int kGroupSize = 16; +constexpr int kTileRows = 128; +constexpr int kTileCols = 64; +constexpr int kTileColGroups = kTileCols / kGroupSize; +constexpr int kTileRowGroups = kTileRows / kGroupSize; +constexpr int kPipelineStages = 2; +constexpr int kStageRows = kTileRows / kPipelineStages; +constexpr int kStageRowGroups = kStageRows / kGroupSize; +constexpr int kElementsPerHalfGroup = 8; +constexpr int kPackedWordsPerGroup = 2; +static_assert(kTileRows == kPipelineStages * kStageRows); +static_assert(kStageRows % kGroupSize == 0); + +template +struct Config { + static constexpr NVTENVFP44Over6Mode mode = kMode; + static constexpr bool err_use_fast_math = kErrUseFastMath; +}; + +struct Candidate { + uint32_t packed[kPackedWordsPerGroup]; + float err; +}; + +struct CandidatePair { + Candidate map4; + Candidate map6; +}; + +struct ScalePair { + nvfp4_scale_t map4; + nvfp4_scale_t map6; + float inv_map4; + float inv_map6; +}; + +template +__device__ __forceinline__ float compute_error_rn(const float diff) { + if constexpr (kMode == kNVTENVFP44Over6MinMSE) { + return __fmul_rn(diff, diff); + } else if constexpr (kMode == kNVTENVFP44Over6MinMAE) { + return fabsf(diff); + } else { + NVTE_DEVICE_ERROR("Unsupported NVFP4 4over6 mode."); + return fabsf(diff); + } +} + +template +__device__ __forceinline__ float compute_error(const float diff) { + if constexpr (kMode == kNVTENVFP44Over6MinMSE) { + return diff * diff; + } else if constexpr (kMode == kNVTENVFP44Over6MinMAE) { + return fabsf(diff); + } else { + NVTE_DEVICE_ERROR("Unsupported NVFP4 4over6 mode."); + return fabsf(diff); + } +} + +template +__device__ __forceinline__ ScalePair compute_scale_pair(const float block_amax, + const float global_amax) { + static_assert(E4M3_MAX == 448 || E4M3_MAX == 256, "Unsupported NVFP4 E4M3 max."); + constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f + constexpr float fp8_max = detail::TypeExtrema::max; // 448.0f + constexpr float expand_to_map4 = 1.5f; + const float S_enc = core::compute_global_encode_scaling_factor_FP4(global_amax); + const float base = block_amax / fp4_max * S_enc; + + ScalePair scales; + scales.map4 = static_cast(fminf(base * expand_to_map4, fp8_max)); + scales.map6 = static_cast(fminf(base, fp8_max)); + + const float S_dec = 1.0f / S_enc; + scales.inv_map4 = + fminf(1.0f / (static_cast(scales.map4) * S_dec), detail::TypeExtrema::max); + scales.inv_map6 = + fminf(1.0f / (static_cast(scales.map6) * S_dec), detail::TypeExtrema::max); + return scales; +} + +template +__device__ __forceinline__ float load_input(const IType *ptr, const size_t idx) { + return static_cast(ptr[idx]); +} + +template +__device__ __forceinline__ void load_row_group(const IType *tile, const int row, + const int col_start, float (&x0)[8], float (&x1)[8], + float *amax) { + Vec x0_vec; + Vec x1_vec; + x0_vec.load_from(&tile[row * kTileCols + col_start]); + x1_vec.load_from(&tile[row * kTileCols + col_start + kElementsPerHalfGroup]); + + *amax = 0.0f; +#pragma unroll + for (int i = 0; i < kElementsPerHalfGroup; ++i) { + const float v0 = static_cast(x0_vec.data.elt[i]); + const float v1 = static_cast(x1_vec.data.elt[i]); + x0[i] = v0; + x1[i] = v1; + *amax = fmaxf(*amax, fabsf(v0)); + *amax = fmaxf(*amax, fabsf(v1)); + } +} + +template +__device__ __forceinline__ void load_col_group(const IType *tile, const int row_start, + const int col, float (&x0)[8], float (&x1)[8], + float *amax) { + *amax = 0.0f; +#pragma unroll + for (int i = 0; i < kElementsPerHalfGroup; ++i) { + const float v0 = load_input(tile, (row_start + i) * kTileCols + col); + const float v1 = load_input(tile, (row_start + i + kElementsPerHalfGroup) * kTileCols + col); + x0[i] = v0; + x1[i] = v1; + *amax = fmaxf(*amax, fabsf(v0)); + *amax = fmaxf(*amax, fabsf(v1)); + } +} + +template +__device__ __forceinline__ void accumulate_dequant_error(const uint32_t dequant_bits, const float x, + const float sf, const float global_amax, + float *err) { + constexpr float fp4_max = detail::TypeExtrema::max; // 6.0f + constexpr float fp8_max = static_cast(E4M3_MAX); + constexpr float err_denom = fp4_max * fp8_max; + const uint16_t half_bits = (dequant_bits >> SHIFT) & 0xFFFF; + + if constexpr (Cfg::err_use_fast_math) { + const float dequant = __half2float(__ushort_as_half(half_bits)); + const float val = dequant * sf * global_amax / err_denom; + const float diff = val - x; + *err += compute_error(diff); + } else { + const float dequant = __half2float(__ushort_as_half(half_bits)); + const float val = __fdiv_rn(__fmul_rn(__fmul_rn(dequant, sf), global_amax), err_denom); + const float diff = __fsub_rn(val, x); + *err = __fadd_rn(*err, compute_error_rn(diff)); + } +} + +template +__device__ __forceinline__ uint32_t cvt_fp32_to_fp4_8x_with_error(const float (&x)[8], + const float block_scale_inverse, + const nvfp4_scale_t sf, + const float global_amax, + float *err) { + uint32_t out = 0; + uint32_t out_dequant_1 = 0; + uint32_t out_dequant_2 = 0; + uint32_t out_dequant_3 = 0; + uint32_t out_dequant_4 = 0; + + constexpr bool is_blackwell = ARCH_BLACKWELL_FAMILY; + if constexpr (is_blackwell) { + asm volatile( + "{\n" + ".reg .b8 byte0, byte1, byte2, byte3;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte0, %6, %5;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte1, %8, %7;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte2, %10, %9;\n" + "cvt.rn.satfinite.e2m1x2.f32 byte3, %12, %11;\n" + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n" + "cvt.rn.f16x2.e2m1x2 %1, byte0;\n" + "cvt.rn.f16x2.e2m1x2 %2, byte1;\n" + "cvt.rn.f16x2.e2m1x2 %3, byte2;\n" + "cvt.rn.f16x2.e2m1x2 %4, byte3;\n" + "}" + : "=r"(out), "=r"(out_dequant_1), "=r"(out_dequant_2), "=r"(out_dequant_3), + "=r"(out_dequant_4) + : "f"(__fmul_rn(x[0], block_scale_inverse)), "f"(__fmul_rn(x[1], block_scale_inverse)), + "f"(__fmul_rn(x[2], block_scale_inverse)), "f"(__fmul_rn(x[3], block_scale_inverse)), + "f"(__fmul_rn(x[4], block_scale_inverse)), "f"(__fmul_rn(x[5], block_scale_inverse)), + "f"(__fmul_rn(x[6], block_scale_inverse)), "f"(__fmul_rn(x[7], block_scale_inverse))); + } else { + NVTE_DEVICE_ERROR( + "FP4 cvt PTX instructions are architecture-specific. " + "Try recompiling with sm_XXXa instead of sm_XXX."); + } + + const float sf_float = static_cast(sf); + accumulate_dequant_error(out_dequant_1, x[0], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_1, x[1], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_2, x[2], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_2, x[3], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_3, x[4], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_3, x[5], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_4, x[6], sf_float, global_amax, err); + accumulate_dequant_error(out_dequant_4, x[7], sf_float, global_amax, err); + return out; +} + +template +__device__ __forceinline__ CandidatePair make_candidates(const float (&x0)[8], const float (&x1)[8], + const ScalePair &scales, + const float global_amax) { + CandidatePair candidates; + candidates.map4.err = 0.0f; + candidates.map6.err = 0.0f; + candidates.map4.packed[0] = cvt_fp32_to_fp4_8x_with_error( + x0, scales.inv_map4, scales.map4, global_amax, &candidates.map4.err); + candidates.map6.packed[0] = cvt_fp32_to_fp4_8x_with_error( + x0, scales.inv_map6, scales.map6, global_amax, &candidates.map6.err); + candidates.map4.packed[1] = cvt_fp32_to_fp4_8x_with_error( + x1, scales.inv_map4, scales.map4, global_amax, &candidates.map4.err); + candidates.map6.packed[1] = cvt_fp32_to_fp4_8x_with_error( + x1, scales.inv_map6, scales.map6, global_amax, &candidates.map6.err); + return candidates; +} + +__device__ __forceinline__ float reduce_group_sum_16(float value) { + const int lane = threadIdx.x & (kWarpThreads - 1); + const int group_base = lane & ~(kGroupSize - 1); + const unsigned mask = 0xffffu << group_base; +#pragma unroll + for (int offset = kGroupSize / 2; offset > 0; offset /= 2) { + value += __shfl_down_sync(mask, value, offset, kGroupSize); + } + return __shfl_sync(mask, value, group_base, kWarpThreads); +} + +__device__ __forceinline__ float reduce_group_max_16(float value) { + const int lane = threadIdx.x & (kWarpThreads - 1); + const int group_base = lane & ~(kGroupSize - 1); + const unsigned mask = 0xffffu << group_base; +#pragma unroll + for (int offset = kGroupSize / 2; offset > 0; offset /= 2) { + value = fmaxf(value, __shfl_down_sync(mask, value, offset, kGroupSize)); + } + return __shfl_sync(mask, value, group_base, kWarpThreads); +} + +__device__ __forceinline__ void store_packed_group(const uint32_t *packed, fp4e2m1x2 *dst) { + const uint64_t packed64 = + static_cast(packed[0]) | (static_cast(packed[1]) << 32); + *reinterpret_cast(dst) = packed64; +} + +__device__ __forceinline__ const uint32_t *select_packed(const CandidatePair &candidates, + const bool pick_map4) { + if (pick_map4) { + return candidates.map4.packed; + } + return candidates.map6.packed; +} + +__device__ __forceinline__ nvfp4_scale_t select_scale(const ScalePair &scales, + const bool pick_map4) { + if (pick_map4) { + return scales.map4; + } + return scales.map6; +} + +__device__ __forceinline__ void cp_async_cg_16(void *dst, const void *src) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + const uint32_t dst_smem_ptr = __cvta_generic_to_shared(dst); + const uint64_t src_gmem_ptr = reinterpret_cast(src); + asm volatile("cp.async.cg.shared.global [%0], [%1], 16;\n" ::"r"(dst_smem_ptr), + "l"(src_gmem_ptr)); +#else + NVTE_DEVICE_ERROR("cp.async is only supported on SM 8.0+."); +#endif +} + +__device__ __forceinline__ void cp_async_commit_group() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + asm volatile("cp.async.commit_group;\n" ::); +#else + NVTE_DEVICE_ERROR("cp.async is only supported on SM 8.0+."); +#endif +} + +template +__device__ __forceinline__ void cp_async_wait_group() { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 800) + asm volatile("cp.async.wait_group %0;\n" ::"n"(N)); +#else + NVTE_DEVICE_ERROR("cp.async is only supported on SM 8.0+."); +#endif +} + +template +__device__ void load_stage_to_shared_async(const IType *input, IType *tile, const size_t rows, + const size_t cols, const size_t stage_row, + const size_t tile_col) { + constexpr int vec_elems = 16 / sizeof(IType); + constexpr int vecs_per_row = kTileCols / vec_elems; + constexpr int vecs = kStageRows * vecs_per_row; + using TileVec = Vec; + + for (int idx = threadIdx.x; idx < vecs; idx += blockDim.x) { + const int local_row = idx / vecs_per_row; + const int local_vec_col = idx - local_row * vecs_per_row; + const int local_col = local_vec_col * vec_elems; + const size_t global_row = stage_row + local_row; + const size_t global_col = tile_col + local_col; + IType *stage_ptr = &tile[local_row * kTileCols + local_col]; + + if (global_row < rows && global_col + vec_elems <= cols) { + cp_async_cg_16(stage_ptr, &input[global_row * cols + global_col]); + } else { + TileVec vec; + vec.clear(); +#pragma unroll + for (int i = 0; i < vec_elems; ++i) { + if (global_row < rows && global_col + i < cols) { + vec.data.elt[i] = input[global_row * cols + global_col + i]; + } + } + vec.store_to(stage_ptr); + } + } +} + +template +__device__ void quantize_stage_rowwise(const IType *tile, fp4e2m1x2 *output, nvfp4_scale_t *scales, + const float *amax, const size_t rows, const size_t cols, + const size_t stage_row, const size_t tile_col, + const size_t scale_stride) { + constexpr int groups = kStageRows * kTileColGroups; + for (int group = threadIdx.x; group < groups; group += blockDim.x) { + const int local_row = group % kStageRows; + const int local_col_group = group / kStageRows; + const int local_col = local_col_group * kGroupSize; + const size_t global_row = stage_row + local_row; + const size_t global_col = tile_col + local_col; + if (global_row >= rows || global_col >= cols) { + continue; + } + + float x0[8]; + float x1[8]; + float group_amax = 0.0f; + load_row_group(tile, local_row, local_col, x0, x1, &group_amax); + + float block_amax = group_amax; + if constexpr (USE_2D_QUANTIZATION) { + block_amax = reduce_group_max_16(group_amax); + } + + float global_amax = amax[0]; + if constexpr (ROW_SCALED_NVFP4) { + global_amax = amax[global_row]; + } + + const ScalePair scale_pair = compute_scale_pair(block_amax, global_amax); + CandidatePair candidates = make_candidates(x0, x1, scale_pair, global_amax); + + float err_map4 = candidates.map4.err; + float err_map6 = candidates.map6.err; + if constexpr (USE_2D_QUANTIZATION) { + err_map4 = reduce_group_sum_16(err_map4); + err_map6 = reduce_group_sum_16(err_map6); + } + + const bool pick_map4 = err_map4 < err_map6; + const nvfp4_scale_t selected_scale = select_scale(scale_pair, pick_map4); + const uint32_t *selected = select_packed(candidates, pick_map4); + + const size_t global_col_group = global_col / kGroupSize; + scales[global_row * scale_stride + global_col_group] = selected_scale; + store_packed_group(selected, &output[(global_row * cols + global_col) / 2]); + } +} + +template +__device__ void quantize_stage_colwise(const IType *tile, fp4e2m1x2 *output_t, + nvfp4_scale_t *scales_t, const float *amax, + const size_t rows, const size_t cols, const size_t stage_row, + const size_t tile_col, const size_t scale_stride_t) { + constexpr int groups = kStageRowGroups * kTileCols; + for (int group = threadIdx.x; group < groups; group += blockDim.x) { + const int local_row_group = group / kTileCols; + const int local_col = group - local_row_group * kTileCols; + const int local_row = local_row_group * kGroupSize; + const size_t global_row = stage_row + local_row; + const size_t global_col = tile_col + local_col; + if (global_row >= rows || global_col >= cols) { + continue; + } + + float x0[8]; + float x1[8]; + float group_amax = 0.0f; + load_col_group(tile, local_row, local_col, x0, x1, &group_amax); + + float block_amax = group_amax; + if constexpr (USE_2D_QUANTIZATION) { + block_amax = reduce_group_max_16(group_amax); + } + + const float global_amax = amax[0]; + const ScalePair scale_pair = compute_scale_pair(block_amax, global_amax); + CandidatePair candidates = make_candidates(x0, x1, scale_pair, global_amax); + + float err_map4 = candidates.map4.err; + float err_map6 = candidates.map6.err; + if constexpr (USE_2D_QUANTIZATION) { + err_map4 = reduce_group_sum_16(err_map4); + err_map6 = reduce_group_sum_16(err_map6); + } + + const bool pick_map4 = err_map4 < err_map6; + const nvfp4_scale_t selected_scale = select_scale(scale_pair, pick_map4); + const uint32_t *selected = select_packed(candidates, pick_map4); + + const size_t global_row_group = global_row / kGroupSize; + scales_t[global_col * scale_stride_t + global_row_group] = selected_scale; + store_packed_group(selected, &output_t[(global_col * rows + global_row) / 2]); + } +} + +template +__global__ void __launch_bounds__(kThreads) + quantize_4over6_kernel(const IType *input, fp4e2m1x2 *output, fp4e2m1x2 *output_t, + nvfp4_scale_t *scales, nvfp4_scale_t *scales_t, + const float *amax_rowwise, const float *amax_colwise, const size_t rows, + const size_t cols, const size_t scale_stride, + const size_t scale_stride_t, const float *noop) { +#if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000) + if (noop != nullptr && noop[0] == 1.0f) { + return; + } + + extern __shared__ char dynamic_shmem[]; + auto *tiles = reinterpret_cast(dynamic_shmem); + const size_t tile_col = blockIdx.x * kTileCols; + const size_t tile_row = blockIdx.y * kTileRows; + + IType *stage_tiles[kPipelineStages]; +#pragma unroll + for (int stage = 0; stage < kPipelineStages; ++stage) { + stage_tiles[stage] = &tiles[stage * kStageRows * kTileCols]; + } + + load_stage_to_shared_async(input, stage_tiles[0], rows, cols, tile_row, tile_col); + cp_async_commit_group(); + cp_async_wait_group<0>(); + __syncthreads(); + + for (int stage = 0; stage < kPipelineStages; ++stage) { + const int next_stage = stage + 1; + if (next_stage < kPipelineStages) { + const size_t next_stage_row = tile_row + next_stage * kStageRows; + load_stage_to_shared_async(input, stage_tiles[next_stage], rows, cols, next_stage_row, + tile_col); + cp_async_commit_group(); + } + + const size_t stage_row = tile_row + stage * kStageRows; + IType *stage_tile = stage_tiles[stage]; + + if constexpr (RETURN_IDENTITY) { + quantize_stage_rowwise( + stage_tile, output, scales, amax_rowwise, rows, cols, stage_row, tile_col, scale_stride); + } + + if constexpr (RETURN_TRANSPOSE) { + const float *columnwise_amax = amax_colwise; + if (columnwise_amax == nullptr) { + columnwise_amax = amax_rowwise; + } + quantize_stage_colwise( + stage_tile, output_t, scales_t, columnwise_amax, rows, cols, stage_row, tile_col, + scale_stride_t); + } + + if (next_stage < kPipelineStages) { + cp_async_wait_group<0>(); + __syncthreads(); + } + } +#else + NVTE_DEVICE_ERROR("sm_100 or higher is required."); +#endif +} + +template +void launch_quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, + cudaStream_t stream) { + const size_t rows = input.flat_first_dim(); + const size_t cols = input.flat_last_dim(); + const bool row_scaled_nvfp4 = output->row_scaled_nvfp4; + const bool return_identity = output->has_data(); + const bool return_transpose = output->has_columnwise_data(); + + const auto *input_ptr = reinterpret_cast(input.data.dptr); + auto *output_ptr = reinterpret_cast(output->data.dptr); + auto *output_t_ptr = reinterpret_cast(output->columnwise_data.dptr); + auto *scales_ptr = reinterpret_cast(output->scale_inv.dptr); + auto *scales_t_ptr = reinterpret_cast(output->columnwise_scale_inv.dptr); + const auto *amax_rowwise_ptr = reinterpret_cast(output->amax.dptr); + const auto *amax_colwise_ptr = reinterpret_cast(output->columnwise_amax.dptr); + const auto *noop_ptr = reinterpret_cast(noop->data.dptr); + + const dim3 grid(DIVUP(cols, static_cast(kTileCols)), + DIVUP(rows, static_cast(kTileRows))); + const dim3 block(kThreads); + const size_t shmem = kPipelineStages * kStageRows * kTileCols * sizeof(IType); + const size_t scale_stride = return_identity ? output->scale_inv.shape[1] : 0; + const size_t scale_stride_t = return_transpose ? output->columnwise_scale_inv.shape[1] : 0; + + TRANSFORMER_ENGINE_SWITCH_CONDITION(return_identity, RETURN_IDENTITY, { + TRANSFORMER_ENGINE_SWITCH_CONDITION(return_transpose, RETURN_TRANSPOSE, { + TRANSFORMER_ENGINE_SWITCH_CONDITION(row_scaled_nvfp4, ROW_SCALED_NVFP4, { + auto kernel = quantize_4over6_kernel; + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, shmem); + kernel<<>>(input_ptr, output_ptr, output_t_ptr, scales_ptr, + scales_t_ptr, amax_rowwise_ptr, amax_colwise_ptr, + rows, cols, scale_stride, scale_stride_t, noop_ptr); + }); + }); + }); +} + +} // namespace quantize_4over6_kernel + +#endif // FP4_TYPE_SUPPORTED + +template +void quantize_4over6(const Tensor &input, const Tensor *noop, Tensor *output, + const QuantizationConfig *quant_config, cudaStream_t stream) { +#if FP4_TYPE_SUPPORTED + using namespace quantize_4over6_kernel; + + checkCuDriverContext(stream); + CheckNoopTensor(*noop, "cast_noop"); + CheckInputTensor(input, "input"); + CheckOutputTensor(*output, "output", false); + + NVTE_CHECK(quant_config != nullptr && quant_config->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled, + "NVFP4 4over6 quantization requires a non-disabled 4over6 mode."); + NVTE_CHECK(!quant_config->stochastic_rounding, + "NVFP4 4over6 quantization does not support stochastic rounding."); + NVTE_CHECK(output->has_data() || output->has_columnwise_data(), + "NVFP4 4over6 output tensor must have rowwise or columnwise data."); + NVTE_CHECK(!output->with_gemm_swizzled_scales, "Output must have scales in compact format."); + NVTE_CHECK(input.flat_last_dim() % kGroupSize == 0, + "NVFP4 4over6 quantization requires columns divisible by ", kGroupSize, "."); + NVTE_CHECK(!(output->has_columnwise_data() || use_2d_quantization) || + input.flat_first_dim() % kGroupSize == 0, + "NVFP4 4over6 columnwise or 2D quantization requires rows divisible by ", kGroupSize, + "."); + NVTE_CHECK(!output->row_scaled_nvfp4 || !use_2d_quantization, + "Row-scaled NVFP4 quantization does not support 2D quantization."); + NVTE_CHECK(!output->row_scaled_nvfp4 || !output->has_columnwise_data(), + "Row-scaled NVFP4 quantization does not produce columnwise output."); + NVTE_CHECK(!use_2d_quantization || output->has_data(), + "NVFP4 4over6 2D quantization requires rowwise output."); + + if (output->has_data()) { + NVTE_CHECK(output->scale_inv.dptr != nullptr, "Scaling tensor must be allocated."); + NVTE_CHECK(output->amax.dptr != nullptr, "Rowwise amax tensor must be allocated."); + NVTE_CHECK(is_fp4_dtype(output->data.dtype), "Output must have FP4 type."); + } + if (output->has_columnwise_data()) { + NVTE_CHECK(output->columnwise_scale_inv.dptr != nullptr, + "Transposed scaling tensor must be allocated."); + NVTE_CHECK(is_fp4_dtype(output->columnwise_data.dtype), + "Transposed output must have FP4 type."); + NVTE_CHECK(output->columnwise_amax.dptr != nullptr || output->amax.dptr != nullptr, + "NVFP4 4over6 columnwise quantization requires columnwise amax or rowwise amax."); + } + + TRANSFORMER_ENGINE_NVFP4_4OVER6_E4M3_MAX_SWITCH( + output->nvfp4_e4m3_max, E4M3_MAX, + TRANSFORMER_ENGINE_NVFP4_4OVER6_MODE_SWITCH( + quant_config->nvfp4_4over6_mode, MODE, + TRANSFORMER_ENGINE_SWITCH_CONDITION( + quant_config->nvfp4_4over6_err_use_fast_math, ERR_USE_FAST_MATH, { + using Cfg = quantize_4over6_kernel::Config; + TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( + input.dtype(), IType, + quantize_4over6_kernel::launch_quantize_4over6( + input, noop, output, stream);); + }););); + + NVTE_CHECK_CUDA(cudaGetLastError()); +#else + NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); +#endif // FP4_TYPE_SUPPORTED +} + +} // namespace nvfp4 +} // namespace dispatch +} // namespace transformer_engine + +#endif // TRANSFORMER_ENGINE_QUANTIZE_4OVER6_NVFP4_CUH_ diff --git a/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp b/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp index 28218e2b4..a1a0dd9d0 100644 --- a/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp +++ b/transformer_engine/common/comm_gemm_overlap/comm_gemm_overlap.cpp @@ -230,6 +230,10 @@ TensorWrapper CommOverlapCore::get_tensor_chunk(const TensorWrapper &source, siz chunk.set_row_scaled_nvfp4(source.get_row_scaled_nvfp4()); continue; } + if (param_type == NVTETensorParam::kNVTENVFP4E4M3Max) { + chunk.set_nvfp4_e4m3_max(source.get_nvfp4_e4m3_max()); + continue; + } auto param = source.get_parameter(param_type); auto param_dptr = reinterpret_cast(param.data_ptr); auto param_dtype = static_cast(param.dtype); diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 12479f2a9..5b6a9bf41 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -178,6 +178,12 @@ struct Tensor { * Only meaningful for NVFP4 tensors. */ bool row_scaled_nvfp4 = false; + /*! \brief Global E4M3 scale bound used by NVFP4. + * + * Standard NVFP4 uses 448. Some 4over6 tensors use 256 to leave room for + * map-to-4 local scale expansion. + */ + int nvfp4_e4m3_max = 448; /*! Map from NVTETensorParam to parameter sizes */ static constexpr size_t attr_sizes[] = { @@ -189,7 +195,8 @@ struct Tensor { sizeof(NVTEBasicTensor), // kNVTEColumnwiseScaleInv sizeof(NVTEBasicTensor), // kNVTEColumnwiseAmax sizeof(uint8_t), // kNVTEWithGEMMSwizzledScales - sizeof(uint8_t) // kNVTERowScaledNVFP4 + sizeof(uint8_t), // kNVTERowScaledNVFP4 + sizeof(int) // kNVTENVFP4E4M3Max }; Tensor() : scaling_mode{NVTE_DELAYED_TENSOR_SCALING}, nvte_tensor{0} {} @@ -206,6 +213,7 @@ struct Tensor { scaling_mode = NVTE_DELAYED_TENSOR_SCALING; with_gemm_swizzled_scales = false; row_scaled_nvfp4 = false; + nvfp4_e4m3_max = 448; } explicit operator NVTETensor() const noexcept { return nvte_tensor; } @@ -477,6 +485,8 @@ struct QuantizationConfig { bool nvfp4_2d_quantization = false; bool stochastic_rounding = false; bool use_fast_math = false; + NVTENVFP44Over6Mode nvfp4_4over6_mode = kNVTENVFP44Over6Disabled; + bool nvfp4_4over6_err_use_fast_math = false; static constexpr size_t attr_sizes[] = { sizeof(uint8_t), // force_pow_2_scales @@ -486,7 +496,9 @@ struct QuantizationConfig { sizeof(NVTETensor), // rng_seed and offset sizeof(uint8_t), // nvfp4_2d_quantization sizeof(uint8_t), // stochastic_rounding - sizeof(uint8_t) // use_fast_math + sizeof(uint8_t), // use_fast_math + sizeof(uint8_t), // nvfp4_4over6_mode + sizeof(uint8_t) // nvfp4_4over6_err_use_fast_math }; }; diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index 045ae8889..ffb324315 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -83,6 +83,13 @@ enum NVTETensorParam { * its values are populated during quantization. */ kNVTERowScaledNVFP4 = 8, + /*! Global E4M3 scale bound used by an NVFP4 tensor. + * + * This is part of the tensor data contract. Downstream dequantization and + * GEMM scale consumers must use the same bound used during quantization. + * Standard NVFP4 uses 448; 4over6 may use 256 for map-to-4 headroom. + */ + kNVTENVFP4E4M3Max = 9, kNVTENumTensorParams }; @@ -111,6 +118,15 @@ enum NVTEScalingMode { NVTE_INVALID_SCALING = 100 }; +/*! \enum NVTENVFP44Over6Mode + * \brief Method for NVFP4 4over6 quantization. + */ +enum NVTENVFP44Over6Mode { + kNVTENVFP44Over6Disabled = 0, /*!< 4over6 is not applied */ + kNVTENVFP44Over6MinMAE = 1, /*!< Select the candidate with lower mean absolute error */ + kNVTENVFP44Over6MinMSE = 2, /*!< Select the candidate with lower mean squared error */ +}; + /*! \brief TE Tensor type * * NVTETensor is a contiguous tensor type storing a pointer @@ -381,6 +397,20 @@ enum NVTEQuantizationConfigAttribute { * inconsistently between kernels. */ kNVTEQuantizationConfigUseFastMath = 7, + /*! Method for NVFP4 4over6 block scale selection. + * + * Non-disabled modes evaluate map-to-4 and map-to-6 candidates for each + * 1x16 block and store the lower-error candidate. The value is an + * NVTENVFP44Over6Mode encoded as uint8_t. + */ + kNVTEQuantizationConfigNVFP44Over6Mode = 8, + /*! Whether the NVFP4 4over6 candidate error computation may use fast math. + * + * This is intentionally separate from kNVTEQuantizationConfigUseFastMath so + * callers can keep candidate selection bitwise deterministic independent + * of ordinary NVFP4 fast-math settings. + */ + kNVTEQuantizationConfigNVFP44Over6ErrUseFastMath = 9, kNVTEQuantizationConfigNumAttributes }; @@ -781,6 +811,11 @@ class TensorWrapper { nvte_set_tensor_param_v2(tensor_, kNVTERowScaledNVFP4, &val, sizeof(val)); } + void set_nvfp4_e4m3_max(int nvfp4_e4m3_max) { + const auto val = nvfp4_e4m3_max; + nvte_set_tensor_param_v2(tensor_, kNVTENVFP4E4M3Max, &val, sizeof(val)); + } + // Parameter getters NVTEBasicTensor get_parameter(const NVTETensorParam param) const noexcept { @@ -823,6 +858,12 @@ class TensorWrapper { return static_cast(val); } + int get_nvfp4_e4m3_max() const { + int val = 448; + nvte_get_tensor_param_v2(tensor_, kNVTENVFP4E4M3Max, &val, sizeof(val), nullptr); + return val; + } + /*! \brief Get an underlying NVTETensor. * * \return NVTETensor held by this TensorWrapper. @@ -1318,6 +1359,20 @@ class QuantizationConfigWrapper { sizeof(val)); } + /*! \brief Set NVFP4 4over6 candidate-selection mode */ + void set_nvfp4_4over6_mode(NVTENVFP44Over6Mode mode) { + const auto val = static_cast(mode); + nvte_set_quantization_config_attribute(config_, kNVTEQuantizationConfigNVFP44Over6Mode, &val, + sizeof(val)); + } + + /*! \brief Set whether NVFP4 4over6 candidate error computation uses fast math */ + void set_nvfp4_4over6_err_use_fast_math(bool use_fast_math) { + const auto val = static_cast(use_fast_math); + nvte_set_quantization_config_attribute( + config_, kNVTEQuantizationConfigNVFP44Over6ErrUseFastMath, &val, sizeof(val)); + } + private: /*! \brief Wrapped NVTEQuantizationConfig. */ NVTEQuantizationConfig config_ = nullptr; diff --git a/transformer_engine/common/recipe/__init__.py b/transformer_engine/common/recipe/__init__.py index b773a81d1..8a03f2f51 100644 --- a/transformer_engine/common/recipe/__init__.py +++ b/transformer_engine/common/recipe/__init__.py @@ -13,6 +13,8 @@ _BACKWARD_OVERRIDES = (None, "high_precision", "dequantized") +_NVFP4_4OVER6_SCOPES = ("none", "weights", "activations", "all") +_NVFP4_4OVER6_ERR_MODES = ("MAE", "MSE") class _FormatHelper(NamedTuple): @@ -522,6 +524,19 @@ class NVFP4BlockScaling(Recipe): If set to `True`, forward activation quantizers emit row-scaled NVFP4 tensors. In this mode, rowwise ``amax`` metadata is stored as a vector with one FP32 value per tensor row. + nvfp4_4over6 : {'none', 'weights', 'activations', 'all'}, default = 'none' + Enable 4over6 adaptive NVFP4 block scaling for selected tensor + scopes. For each selected FP4 block, quantization compares + map-to-4 and map-to-6 candidates and stores the candidate with + lower configured error. Current 4over6 support targets RL and + post-training scenarios; pre-training paths that combine 4over6 + with RHT are not yet implemented. + nvfp4_4over6_e4m3_use_256 : {'none', 'weights', 'activations', 'all'}, default = 'all' + Select 4over6 tensors that use 256 as the global E4M3 scale + bound. By default, all 4over6 tensors use 256. Use ``'none'`` + to keep the standard NVFP4 448 bound for 4over6 tensors. + nvfp4_4over6_err_mode : {'MAE', 'MSE'}, default = 'MAE' + Error metric used by NVFP4 4over6 candidate selection. backward_override : {None, 'high_precision', 'dequantized'}, default = None Backward precision mode. None does not modify backward behavior, `high_precision` keeps original high-precision operands for backward, @@ -536,6 +551,9 @@ class NVFP4BlockScaling(Recipe): ) disable_2d_quantization: bool = os.getenv("NVTE_NVFP4_DISABLE_2D_QUANTIZATION", "0") == "1" row_scaled_activation: bool = os.getenv("NVTE_NVFP4_ROW_SCALED_ACTIVATION", "0") == "1" + nvfp4_4over6: str = os.getenv("NVTE_NVFP4_4OVER6", "none") + nvfp4_4over6_e4m3_use_256: str = os.getenv("NVTE_NVFP4_4OVER6_E4M3_USE_256", "all") + nvfp4_4over6_err_mode: str = os.getenv("NVTE_NVFP4_4OVER6_ERR_MODE", "MAE").upper() fp4_format: Format = Format.E2M1 fp8_format: Format = Format.E4M3 @@ -551,6 +569,15 @@ def __post_init__(self) -> None: assert ( self.backward_override in _BACKWARD_OVERRIDES ), "NVTE_BACKWARD_OVERRIDE must be unset or one of: 'high_precision', 'dequantized'." + assert ( + self.nvfp4_4over6 in _NVFP4_4OVER6_SCOPES + ), "NVTE_NVFP4_4OVER6 must be one of: 'none', 'weights', 'activations', 'all'." + assert ( + self.nvfp4_4over6_e4m3_use_256 in _NVFP4_4OVER6_SCOPES + ), "NVTE_NVFP4_4OVER6_E4M3_USE_256 must be one of: 'none', 'weights', 'activations', 'all'." + assert ( + self.nvfp4_4over6_err_mode in _NVFP4_4OVER6_ERR_MODES + ), "NVTE_NVFP4_4OVER6_ERR_MODE must be one of: 'MAE', 'MSE'." # Quantization params # Note: RHT is currently only applied to column-wise usage so that @@ -580,6 +607,9 @@ def _make_repr(self) -> str: f"fp8_mha={self.fp8_mha}, " f"backward_override={self.backward_override}, " f"row_scaled_activation={self.row_scaled_activation}, " + f"nvfp4_4over6={self.nvfp4_4over6}, " + f"nvfp4_4over6_e4m3_use_256={self.nvfp4_4over6_e4m3_use_256}, " + f"nvfp4_4over6_err_mode={self.nvfp4_4over6_err_mode}, " f"fp4_quant_fwd_inp={self.fp4_quant_fwd_inp}, " f"fp4_quant_fwd_weight={self.fp4_quant_fwd_weight}, " f"fp4_quant_bwd_grad={self.fp4_quant_bwd_grad}, " diff --git a/transformer_engine/common/recipe/nvfp4.cu b/transformer_engine/common/recipe/nvfp4.cu index 1c419d4f8..576e6139c 100644 --- a/transformer_engine/common/recipe/nvfp4.cu +++ b/transformer_engine/common/recipe/nvfp4.cu @@ -65,15 +65,15 @@ namespace nvfp4_recipe { * --------------------------------------------------------------------------- */ -// constexpr float factor = 6.0 * 6.0 * 448.0 * 448.0; -constexpr float factor_inv = 1.0 / (6.0 * 6.0 * 448.0 * 448.0); constexpr int kTileDim = 16; constexpr int kThreadsPerBlock = 256; // Kernel to compute alpha *= amax_A * amax_B / factor __global__ void compute_nvfp4_per_tensor_scale_kernel(float alpha_in, const float *amax_A, - const float *amax_B, float *alpha_out) { - // factor is defined in the enclosing namespace + const float *amax_B, float fp8_max_A, + float fp8_max_B, float *alpha_out) { + constexpr float fp4_max = 6.0f; + const float factor_inv = 1.0f / (fp4_max * fp4_max * fp8_max_A * fp8_max_B); *alpha_out = alpha_in * (*amax_A) * (*amax_B) * factor_inv; } @@ -924,6 +924,8 @@ void nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_r void *amax_A_ptr = use_rowwise_amax_A ? tA->amax.dptr : tA->columnwise_amax.dptr; void *amax_B_ptr = use_rowwise_amax_B ? tB->amax.dptr : tB->columnwise_amax.dptr; void *alpha_ptr = tOut->data.dptr; + const float fp8_max_A = static_cast(tA->nvfp4_e4m3_max); + const float fp8_max_B = static_cast(tB->nvfp4_e4m3_max); // check for not null pointers NVTE_CHECK(amax_A_ptr != nullptr, "amax_A_ptr is null"); @@ -932,7 +934,8 @@ void nvte_nvfp4_compute_per_tensor_scale(const NVTETensor inpA, const bool use_r nvfp4_recipe::compute_nvfp4_per_tensor_scale_kernel<<<1, 1, 0, stream>>>( alpha_in, reinterpret_cast(amax_A_ptr), - reinterpret_cast(amax_B_ptr), reinterpret_cast(alpha_ptr)); + reinterpret_cast(amax_B_ptr), fp8_max_A, fp8_max_B, + reinterpret_cast(alpha_ptr)); NVTE_CHECK_CUDA(cudaGetLastError()); #else NVTE_ERROR("FP4 support requires CUDA 12.8+, but compile-time CUDA version is ", CUDA_VERSION); diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 1a52d7601..561f64d59 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -855,6 +855,11 @@ void nvte_set_tensor_param_v2(NVTETensor tensor, NVTETensorParam param, const vo case kNVTERowScaledNVFP4: t.row_scaled_nvfp4 = static_cast(*reinterpret_cast(buf)); break; + case kNVTENVFP4E4M3Max: + std::memcpy(&t.nvfp4_e4m3_max, buf, attr_size); + NVTE_CHECK(t.nvfp4_e4m3_max == 448 || t.nvfp4_e4m3_max == 256, + "Unsupported NVFP4 E4M3 max (got ", t.nvfp4_e4m3_max, ")"); + break; default: NVTE_ERROR("Unsupported tensor parameter (", static_cast(param), ")"); } @@ -938,6 +943,9 @@ void nvte_get_tensor_param_v2(const NVTETensor tensor, NVTETensorParam param, vo case kNVTERowScaledNVFP4: *reinterpret_cast(buf) = static_cast(t->row_scaled_nvfp4); break; + case kNVTENVFP4E4M3Max: + std::memcpy(buf, &t->nvfp4_e4m3_max, attr_size); + break; default: NVTE_ERROR("Unsupported tensor parameter (", static_cast(param), ")"); } @@ -1049,6 +1057,14 @@ void nvte_get_quantization_config_attribute(NVTEQuantizationConfig config, case kNVTEQuantizationConfigUseFastMath: bool_to_uint8(config_.use_fast_math, buf); break; + case kNVTEQuantizationConfigNVFP44Over6Mode: { + const auto val = static_cast(config_.nvfp4_4over6_mode); + std::memcpy(buf, &val, attr_size); + break; + } + case kNVTEQuantizationConfigNVFP44Over6ErrUseFastMath: + bool_to_uint8(config_.nvfp4_4over6_err_use_fast_math, buf); + break; default: NVTE_ERROR("Unsupported NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); } @@ -1104,6 +1120,18 @@ void nvte_set_quantization_config_attribute(NVTEQuantizationConfig config, case kNVTEQuantizationConfigUseFastMath: uint8_to_bool(buf, config_.use_fast_math); break; + case kNVTEQuantizationConfigNVFP44Over6Mode: { + const auto val = *reinterpret_cast(buf); + NVTE_CHECK(val == static_cast(kNVTENVFP44Over6Disabled) || + val == static_cast(kNVTENVFP44Over6MinMAE) || + val == static_cast(kNVTENVFP44Over6MinMSE), + "Invalid NVFP4 4over6 mode (got ", static_cast(val), ")"); + config_.nvfp4_4over6_mode = static_cast(val); + break; + } + case kNVTEQuantizationConfigNVFP44Over6ErrUseFastMath: + uint8_to_bool(buf, config_.nvfp4_4over6_err_use_fast_math); + break; default: NVTE_ERROR("Unsupported NVTEQuantizationConfigAttribute (got ", static_cast(attr), ")"); } diff --git a/transformer_engine/pytorch/csrc/common.h b/transformer_engine/pytorch/csrc/common.h index 94350da1e..b376b3022 100644 --- a/transformer_engine/pytorch/csrc/common.h +++ b/transformer_engine/pytorch/csrc/common.h @@ -327,6 +327,10 @@ class NVFP4Quantizer : public Quantizer { // 2D block scaling bool with_2d_quantization; bool stochastic_rounding; + // 4over6 candidate-selection mode used when quantizing emitted NVFP4 tensors. + NVTENVFP44Over6Mode nvfp4_4over6_mode; + // Global E4M3 scale bound used by emitted NVFP4 tensors. + int nvfp4_e4m3_max; // Whether tensors emitted by this quantizer use row-scaled NVFP4 metadata. bool row_scaled_nvfp4; diff --git a/transformer_engine/pytorch/csrc/extensions/cast.cpp b/transformer_engine/pytorch/csrc/extensions/cast.cpp index 2b38339d6..d1a9cd858 100644 --- a/transformer_engine/pytorch/csrc/extensions/cast.cpp +++ b/transformer_engine/pytorch/csrc/extensions/cast.cpp @@ -84,6 +84,8 @@ void group_quantize_nvfp4_impl(const GroupedTensorWrapper &grouped_input_tensor, // assert the 2D scaling case, since 2D scaling grouped quant kernel is not ready yet NVTE_CHECK(!nvfp4_quantizer_cpp->with_2d_quantization, "2D scaling grouped quant kernel is not ready yet"); + NVTE_CHECK(nvfp4_quantizer_cpp->nvfp4_4over6_mode == kNVTENVFP44Over6Disabled, + "NVFP4 4over6 quantization is not supported for grouped quantization."); auto quant_config_cpp = QuantizationConfigWrapper(); @@ -722,6 +724,9 @@ std::tuple, std::vector, bool> bulk_alloc // Quantization parameters const auto rowwise_usage = quantizer_cpp_list[0]->rowwise_usage; const bool row_scaled_nvfp4 = quantizer_cpp_list[0]->row_scaled_nvfp4; + const bool nvfp4_use_4over6 = + quantizer_cpp_list[0]->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; + const int nvfp4_e4m3_max = quantizer_cpp_list[0]->nvfp4_e4m3_max; const auto columnwise_usage = quantizer_cpp_list[0]->columnwise_usage; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 bulk allocation requires rowwise usage."); @@ -866,10 +871,12 @@ std::tuple, std::vector, bool> bulk_alloc py::object amax_columnwise = columnwise_usage ? py::cast(amax_columnwise_list[i]) : py::none(); // Construct Python tensor - tensor_py_list.emplace_back(NVFP4TensorClass(rowwise_data, rowwise_scale, columnwise_data, - columnwise_scale, amax_rowwise, amax_columnwise, - fp4_dtype, quantizer_py_list[i], - with_gemm_swizzled_scales, row_scaled_nvfp4)); + tensor_py_list.emplace_back( + NVFP4TensorClass(rowwise_data, rowwise_scale, columnwise_data, columnwise_scale, + amax_rowwise, amax_columnwise, fp4_dtype, quantizer_py_list[i], + with_gemm_swizzled_scales, py::arg("row_scaled_nvfp4") = row_scaled_nvfp4, + py::arg("nvfp4_use_4over6") = nvfp4_use_4over6, + py::arg("nvfp4_e4m3_max") = nvfp4_e4m3_max)); // Construct C++ tensor // Use a TensorWrapper variable to hold the output of makeTransformerEngineTensor, @@ -887,6 +894,7 @@ std::tuple, std::vector, bool> bulk_alloc columnwise_usage ? columnwise_scale_shapes[i] : std::vector{0}, scaling_mode); tensor_wrapper.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); tensor_wrapper.set_row_scaled_nvfp4(row_scaled_nvfp4); + tensor_wrapper.set_nvfp4_e4m3_max(nvfp4_e4m3_max); // Set the amax rowwise and amax columnwise if available if (rowwise_usage) { @@ -997,6 +1005,9 @@ void split_quantize_nvfp4_impl_with_rht_helper(const TensorWrapper &input, cudaStream_t stream) { const size_t num_tensors = split_sections.size(); const auto &quantizer = *quantizers.front(); + const bool nvfp4_use_4over6 = quantizer.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; + NVTE_CHECK(!nvfp4_use_4over6, + "NVFP4 4over6 quantization is not supported with RHT split quantization."); std::vector nvte_tensor_input_list; std::vector nvte_tensor_output_list; @@ -1032,6 +1043,13 @@ void split_quantize_nvfp4_impl_with_rht_helper(const TensorWrapper &input, num_tensors, need_stochastic_rounding, with_bulk_generate_rng_states, need_separate_rng_states, quant_config_list, quant_config_list_colwise); + for (auto &config : quant_config_list) { + config.set_nvfp4_4over6_mode(quantizer.nvfp4_4over6_mode); + } + for (auto &config : quant_config_list_colwise) { + config.set_nvfp4_4over6_mode(quantizer.nvfp4_4over6_mode); + } + // Enable NVFP4 kernels to use math operations that sacrifice // accuracy for performance. These optimizations are experimental // and inconsistently implemented. @@ -1039,8 +1057,10 @@ void split_quantize_nvfp4_impl_with_rht_helper(const TensorWrapper &input, // 1. replace 1 / x by reciprocal_approximate_ftz(x) // 2. when RHT cast fusion is available, fusion allows cast to be performed on FP32 data, // this will essentially remove a round trip between FP32 to BF16 then FP32 + // NVFP4 4over6 candidate error math is controlled separately by + // NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH. const auto use_fast_math = transformer_engine::getenv("NVTE_USE_FAST_MATH"); - if (use_fast_math) { + if (use_fast_math && !nvfp4_use_4over6) { for (auto &config : quant_config_list) { config.set_use_fast_math(true); } @@ -1049,6 +1069,17 @@ void split_quantize_nvfp4_impl_with_rht_helper(const TensorWrapper &input, } } + const auto use_4over6_err_use_fast_math = + transformer_engine::getenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH"); + if (use_4over6_err_use_fast_math) { + for (auto &config : quant_config_list) { + config.set_nvfp4_4over6_err_use_fast_math(true); + } + for (auto &config : quant_config_list_colwise) { + config.set_nvfp4_4over6_err_use_fast_math(true); + } + } + auto &quant_config_list_colwise_to_use = need_separate_rng_states ? quant_config_list_colwise : quant_config_list; @@ -1157,6 +1188,9 @@ void split_quantize_nvfp4_impl_helper(const TensorWrapper &input, cudaStream_t stream) { const size_t num_tensors = input_list.size(); const auto &quantizer = *quantizers.front(); + const bool nvfp4_use_4over6 = quantizer.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; + NVTE_CHECK(!nvfp4_use_4over6 || !quantizer.stochastic_rounding, + "NVFP4 4over6 quantization does not support stochastic rounding."); std::vector nvte_tensor_input_list; std::vector nvte_tensor_output_list; @@ -1189,6 +1223,27 @@ void split_quantize_nvfp4_impl_helper(const TensorWrapper &input, need_separate_rng_states, quant_config_list, dummy_quant_config_list_colwise); // colwise rng states are not needed in this case + for (auto &config : quant_config_list) { + config.set_nvfp4_4over6_mode(quantizer.nvfp4_4over6_mode); + } + + // NVFP4 4over6 candidate error math is controlled separately by + // NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH. + const auto use_fast_math = transformer_engine::getenv("NVTE_USE_FAST_MATH"); + if (use_fast_math && !nvfp4_use_4over6) { + for (auto &config : quant_config_list) { + config.set_use_fast_math(true); + } + } + + const auto use_4over6_err_use_fast_math = + transformer_engine::getenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH"); + if (use_4over6_err_use_fast_math) { + for (auto &config : quant_config_list) { + config.set_nvfp4_4over6_err_use_fast_math(true); + } + } + // We need: // 1. Rowwise amax = amax for input // 2. Columnwise amax = amax for input too @@ -1259,6 +1314,11 @@ void split_quantize_nvfp4_impl(const TensorWrapper &input, "NVFP4 split-quantize does not support 2D quantization"); NVTE_CHECK(!quantizer.with_amax_reduction, "NVFP4 split-quantize does not support amax reduction"); + if (quantizer.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled) { + NVTE_CHECK(!quantizer.with_rht, "NVFP4 4over6 quantization does not support RHT."); + NVTE_CHECK(!quantizer.stochastic_rounding, + "NVFP4 4over6 quantization does not support stochastic rounding."); + } // Check input tensor shape const size_t input_last_dim = input.ndim() > 0 ? input.size(input.ndim() - 1) : 1; diff --git a/transformer_engine/pytorch/csrc/quantizer.cpp b/transformer_engine/pytorch/csrc/quantizer.cpp index 7045995dd..bc87b54ba 100644 --- a/transformer_engine/pytorch/csrc/quantizer.cpp +++ b/transformer_engine/pytorch/csrc/quantizer.cpp @@ -1729,6 +1729,20 @@ NVFP4Quantizer::NVFP4Quantizer(const py::handle& quantizer) : Quantizer(quantize this->with_post_rht_amax = quantizer.attr("with_post_rht_amax").cast(); this->with_2d_quantization = quantizer.attr("with_2d_quantization").cast(); this->stochastic_rounding = quantizer.attr("stochastic_rounding").cast(); + const bool nvfp4_use_4over6 = quantizer.attr("nvfp4_use_4over6").cast(); + this->nvfp4_e4m3_max = quantizer.attr("nvfp4_e4m3_max").cast(); + NVTE_CHECK(this->nvfp4_e4m3_max == 448 || this->nvfp4_e4m3_max == 256, + "Unsupported NVFP4 E4M3 max: ", this->nvfp4_e4m3_max); + const auto nvfp4_4over6_err_mode = quantizer.attr("nvfp4_4over6_err_mode").cast(); + if (!nvfp4_use_4over6) { + this->nvfp4_4over6_mode = kNVTENVFP44Over6Disabled; + } else if (nvfp4_4over6_err_mode == "MAE") { + this->nvfp4_4over6_mode = kNVTENVFP44Over6MinMAE; + } else if (nvfp4_4over6_err_mode == "MSE") { + this->nvfp4_4over6_mode = kNVTENVFP44Over6MinMSE; + } else { + NVTE_ERROR("Unsupported NVFP4 4over6 error mode: ", nvfp4_4over6_err_mode); + } this->row_scaled_nvfp4 = quantizer.attr("row_scaled_nvfp4").cast(); // Get amax reduction group if needed for NVFP4 AG @@ -1778,6 +1792,8 @@ std::pair NVFP4Quantizer::create_tensor( "NVFP4 requires tensor dims that are divisible by ", NVFP4_BLOCK_SIZE, " (got shape=", shape, ")"); const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; + const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; + const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 quantization requires rowwise usage."); NVTE_CHECK(!columnwise_usage, @@ -1845,6 +1861,8 @@ std::pair NVFP4Quantizer::create_tensor( kwargs["quantizer"] = this->quantizer; kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); + kwargs["nvfp4_use_4over6"] = py::cast(nvfp4_use_4over6); + kwargs["nvfp4_e4m3_max"] = py::cast(nvfp4_e4m3_max); kwargs["fake_dtype"] = GetATenDType(dtype); py::tuple args(0); @@ -1875,6 +1893,8 @@ std::pair NVFP4Quantizer::create_tensor( kwargs["with_gemm_swizzled_scales"] = py::cast(with_gemm_swizzled_scales); kwargs["device"] = py::cast(device); kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); + kwargs["nvfp4_use_4over6"] = py::cast(nvfp4_use_4over6); + kwargs["nvfp4_e4m3_max"] = py::cast(nvfp4_e4m3_max); py::tuple args(0); PyObject* result = PyObject_Call(reinterpret_cast(NVFP4TensorPythonClass), args.ptr(), kwargs.ptr()); @@ -1908,6 +1928,7 @@ std::pair NVFP4Quantizer::create_tensor( } out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); + out_cpp.set_nvfp4_e4m3_max(nvfp4_e4m3_max); this->set_quantization_params(&out_cpp); return {std::move(out_cpp), std::move(out_py)}; @@ -1936,6 +1957,8 @@ std::pair NVFP4Quantizer::create_grouped_tenso std::optional columnwise_amax; const std::vector logical_shape_vec = {logical_first_dim, logical_last_dim}; const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; + const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; + const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 grouped quantization requires rowwise usage."); NVTE_CHECK(!columnwise_usage, @@ -2010,6 +2033,8 @@ std::pair NVFP4Quantizer::create_grouped_tenso kwargs["tensor_offsets"] = tensor_offsets.has_value() ? py::cast(*tensor_offsets) : py::none(); kwargs["with_gemm_swizzled_scales"] = this->optimize_for_gemm; kwargs["row_scaled_nvfp4"] = py::cast(row_scaled_nvfp4); + kwargs["nvfp4_use_4over6"] = py::cast(nvfp4_use_4over6); + kwargs["nvfp4_e4m3_max"] = py::cast(nvfp4_e4m3_max); PyObject* result = PyObject_Call(GroupedTensorClass.ptr(), args.ptr(), kwargs.ptr()); if (result == nullptr) { PyErr_Print(); @@ -2085,12 +2110,16 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( const auto [flat_first_dim, flat_last_dim] = get_2d_dims(shape); const bool row_scaled_nvfp4 = this->row_scaled_nvfp4; + const bool nvfp4_use_4over6 = this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; + const int nvfp4_e4m3_max = this->nvfp4_e4m3_max; if (row_scaled_nvfp4) { NVTE_CHECK(rowwise_usage, "Row-scaled NVFP4 quantization requires rowwise usage."); NVTE_CHECK(!columnwise_usage, "Row-scaled NVFP4 quantization does not support columnwise usage."); } tensor.attr("_row_scaled_nvfp4") = py::cast(row_scaled_nvfp4); + tensor.attr("_nvfp4_use_4over6") = py::cast(nvfp4_use_4over6); + tensor.attr("_nvfp4_e4m3_max") = py::cast(nvfp4_e4m3_max); // Coerce row-wise data if (rowwise_usage) { @@ -2195,6 +2224,7 @@ std::pair NVFP4Quantizer::convert_and_update_tensor( } out_cpp.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); out_cpp.set_row_scaled_nvfp4(row_scaled_nvfp4); + out_cpp.set_nvfp4_e4m3_max(nvfp4_e4m3_max); this->set_quantization_params(&out_cpp); return {std::move(out_cpp), std::move(tensor)}; @@ -2285,6 +2315,14 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou } quant_config.set_nvfp4_2d_quantization(this->with_2d_quantization); quant_config.set_stochastic_rounding(this->stochastic_rounding); + quant_config.set_nvfp4_4over6_mode(this->nvfp4_4over6_mode); + quant_config_columnwise.set_nvfp4_4over6_mode(this->nvfp4_4over6_mode); + + if (this->nvfp4_4over6_mode != kNVTENVFP44Over6Disabled) { + NVTE_CHECK(!this->with_rht, "NVFP4 4over6 quantization does not support RHT."); + NVTE_CHECK(!this->stochastic_rounding, + "NVFP4 4over6 quantization does not support stochastic rounding."); + } // We only need RHT for columnwise usage. // flat first dim and last dim for multi dimensional input @@ -2425,12 +2463,21 @@ void NVFP4Quantizer::quantize_impl(const TensorWrapper& input, TensorWrapper& ou // 1. replace 1 / x by reciprocal_approximate_ftz(x) // 2. when RHT cast fusion is available, fusion allows cast to be performed on FP32 data, // this will essentially remove a round trip between FP32 to BF16 then FP32 + // NVFP4 4over6 candidate error math is controlled separately by + // NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH. const auto use_fast_math = transformer_engine::getenv("NVTE_USE_FAST_MATH"); - if (use_fast_math) { + if (use_fast_math && this->nvfp4_4over6_mode == kNVTENVFP44Over6Disabled) { quant_config.set_use_fast_math(true); quant_config_columnwise.set_use_fast_math(true); } + const auto use_4over6_err_use_fast_math = + transformer_engine::getenv("NVTE_NVFP4_4OVER6_ERR_USE_FAST_MATH"); + if (use_4over6_err_use_fast_math) { + quant_config.set_nvfp4_4over6_err_use_fast_math(true); + quant_config_columnwise.set_nvfp4_4over6_err_use_fast_math(true); + } + if (this->with_rht) { if (eligible_for_rht_cast_fusion) { // fusion kernel requires passing in RHT matrix directly for maximum performance diff --git a/transformer_engine/pytorch/csrc/type_converters.cpp b/transformer_engine/pytorch/csrc/type_converters.cpp index 37ab0b053..ddb85808a 100644 --- a/transformer_engine/pytorch/csrc/type_converters.cpp +++ b/transformer_engine/pytorch/csrc/type_converters.cpp @@ -135,6 +135,7 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) const bool columnwise_usage = !(tensor.attr("_columnwise_data").is_none()); const bool with_gemm_swizzled_scales = tensor.attr("_with_gemm_swizzled_scales").cast(); const bool row_scaled_nvfp4 = tensor.attr("_row_scaled_nvfp4").cast(); + const int nvfp4_e4m3_max = tensor.attr("_nvfp4_e4m3_max").cast(); NVTE_CHECK(rowwise_usage || columnwise_usage, "No data found for NVFP4 Tensor."); @@ -165,6 +166,7 @@ TensorWrapper NVTETensorFromNVFP4Tensor(py::handle tensor, Quantizer *quantizer) // Scale layout ret.set_with_gemm_swizzled_scales(with_gemm_swizzled_scales); ret.set_row_scaled_nvfp4(row_scaled_nvfp4); + ret.set_nvfp4_e4m3_max(nvfp4_e4m3_max); // Quantizer state quantizer->set_quantization_params(&ret); diff --git a/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py b/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py index acb7abefd..5c23c7670 100644 --- a/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py +++ b/transformer_engine/pytorch/custom_recipes/quantization_ref_nvfp4.py @@ -221,6 +221,8 @@ class NVFP4TensorRef(QuantizedTensorStorage): scale_t: Optional[torch.Tensor] = None global_amax_row: Optional[torch.Tensor] = None global_amax_col: Optional[torch.Tensor] = None + nvfp4_use_4over6: bool = False + nvfp4_e4m3_max: int = 448 dtype: Optional[torch.dtype] = None device: Optional[torch.device] = None @@ -350,9 +352,15 @@ def __init__( eps: float = 0.0, quant_tile_shape: Tuple[int, int] = (1, 16), row_scaled_nvfp4: bool = False, + nvfp4_use_4over6: bool = False, + nvfp4_e4m3_max: int = 448, + nvfp4_4over6_err_mode: str = "MAE", with_rht: bool = False, with_random_sign_mask: bool = True, ): + nvfp4_4over6_err_mode = nvfp4_4over6_err_mode.upper() + if nvfp4_4over6_err_mode not in ("MAE", "MSE"): + raise ValueError("nvfp4_4over6_err_mode must be 'MAE' or 'MSE'.") if row_scaled_nvfp4: if not rowwise: raise ValueError("Row-scaled NVFP4 reference quantization requires rowwise usage.") @@ -360,6 +368,11 @@ def __init__( raise ValueError( "Row-scaled NVFP4 reference quantization does not support columnwise usage." ) + if nvfp4_use_4over6: + if pow_2_scales: + raise ValueError("4over6 is only supported for NVFP4 (non-pow2) mode.") + if quant_tile_shape not in ((1, 16), (16, 16)): + raise ValueError("4over6 reference quantization only supports 1x16 or 16x16 tiles.") super().__init__(rowwise=rowwise, columnwise=columnwise) self.internal = True @@ -368,6 +381,11 @@ def __init__( self.eps = eps self.quant_tile_shape = quant_tile_shape self.row_scaled_nvfp4 = row_scaled_nvfp4 + self.nvfp4_use_4over6 = nvfp4_use_4over6 + self.nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_use_4over6 else 448 + if self.nvfp4_e4m3_max not in (448, 256): + raise ValueError("nvfp4_e4m3_max must be 448 or 256.") + self.nvfp4_4over6_err_mode = nvfp4_4over6_err_mode self.with_rht = with_rht self.with_random_sign_mask = with_random_sign_mask @@ -446,6 +464,113 @@ def _recover_swizzled_scales( result = torch.reshape(tmp, (rounded_m, rounded_n)) return result[:m, :scale_n] + @staticmethod + def _quantize_blockwise_4over6_reference( + x: torch.Tensor, + vec_max: torch.Tensor, + global_amax: torch.Tensor, + global_encode_scale: torch.Tensor, + global_decode_scale: torch.Tensor, + row_scaled_nvfp4: bool, + tile_len_y: int, + nvfp4_4over6_err_mode: str, + nvfp4_e4m3_max: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Quantize NVFP4 with 4over6 candidate selection. + + This mirrors the CUDA path: map-to-4 uses a 1.5x expanded E4M3 block scale, + the configured error is computed in the original input domain with the + selected global E4M3 denominator, and ties choose map-to-6. + """ + m, num_blocks, tile_len_x = x.shape + n = num_blocks * tile_len_x + FLOAT4_E2M1_MAX = torch.tensor(6.0, device=x.device, dtype=torch.float32) + FLOAT8_E4M3_MAX = torch.tensor(448.0, device=x.device, dtype=torch.float32) + GLOBAL_SCALE_E4M3_MAX = torch.tensor( + float(nvfp4_e4m3_max), device=x.device, dtype=torch.float32 + ) + + decode_scale_base = torch.div(vec_max, FLOAT4_E2M1_MAX) * global_encode_scale + decode_scale_map4 = decode_scale_base * 1.5 + decode_scale_map6 = decode_scale_base + decode_scale_map4 = torch.clamp( + decode_scale_map4, min=-FLOAT8_E4M3_MAX, max=FLOAT8_E4M3_MAX + ).to(torch.float8_e4m3fn) + decode_scale_map6 = torch.clamp( + decode_scale_map6, min=-FLOAT8_E4M3_MAX, max=FLOAT8_E4M3_MAX + ).to(torch.float8_e4m3fn) + + fp32_max = torch.tensor( + torch.finfo(torch.float32).max, + device=decode_scale_map4.device, + dtype=torch.float32, + ) + encode_scale_map4 = torch.min( + torch.div(1.0, decode_scale_map4.to(torch.float32) * global_decode_scale), + fp32_max, + ) + encode_scale_map6 = torch.min( + torch.div(1.0, decode_scale_map6.to(torch.float32) * global_decode_scale), + fp32_max, + ) + + clipped_x_map4 = torch.clamp( + x.to(torch.float32) * encode_scale_map4, + -FLOAT4_E2M1_MAX, + FLOAT4_E2M1_MAX, + ).reshape(m, n) + clipped_x_map6 = torch.clamp( + x.to(torch.float32) * encode_scale_map6, + -FLOAT4_E2M1_MAX, + FLOAT4_E2M1_MAX, + ).reshape(m, n) + qx_map4 = cast_to_fp4x2(clipped_x_map4) + qx_map6 = cast_to_fp4x2(clipped_x_map6) + + fp4_map4 = cast_from_fp4x2(qx_map4, torch.float32).view(m, num_blocks, tile_len_x) + fp4_map6 = cast_from_fp4x2(qx_map6, torch.float32).view(m, num_blocks, tile_len_x) + denom = FLOAT4_E2M1_MAX * GLOBAL_SCALE_E4M3_MAX + sf_map4 = decode_scale_map4.to(torch.float32).squeeze(-1) + sf_map6 = decode_scale_map6.to(torch.float32).squeeze(-1) + if row_scaled_nvfp4: + error_global_amax = global_amax.squeeze(-1) + else: + error_global_amax = global_amax + x_float = x.to(torch.float32) + err_map4 = torch.zeros_like(vec_max) + err_map6 = torch.zeros_like(vec_max) + for idx in range(tile_len_x): + val_map4 = fp4_map4[:, :, idx] * sf_map4 + val_map4 = val_map4 * error_global_amax + val_map4 = val_map4 / denom + diff_map4 = val_map4 - x_float[:, :, idx] + if nvfp4_4over6_err_mode == "MSE": + err_map4 = err_map4 + (diff_map4 * diff_map4).unsqueeze(-1) + else: + err_map4 = err_map4 + torch.abs(diff_map4).unsqueeze(-1) + + val_map6 = fp4_map6[:, :, idx] * sf_map6 + val_map6 = val_map6 * error_global_amax + val_map6 = val_map6 / denom + diff_map6 = val_map6 - x_float[:, :, idx] + if nvfp4_4over6_err_mode == "MSE": + err_map6 = err_map6 + (diff_map6 * diff_map6).unsqueeze(-1) + else: + err_map6 = err_map6 + torch.abs(diff_map6).unsqueeze(-1) + if tile_len_y == 1: + pick_map4 = err_map4 < err_map6 + else: + err_map4_blocks = err_map4.view(m // tile_len_y, tile_len_y, num_blocks, 1).sum(dim=1) + err_map6_blocks = err_map6.view(m // tile_len_y, tile_len_y, num_blocks, 1).sum(dim=1) + pick_map4 = (err_map4_blocks < err_map6_blocks).repeat_interleave(tile_len_y, dim=0) + qx = torch.where( + pick_map4.expand(-1, -1, tile_len_x // 2), + qx_map4.view(m, num_blocks, tile_len_x // 2), + qx_map6.view(m, num_blocks, tile_len_x // 2), + ).reshape(m, n // 2) + decode_scale = torch.where(pick_map4, decode_scale_map4, decode_scale_map6).squeeze(-1) + return qx, decode_scale + @classmethod def _quantize_blockwise_reference( cls, @@ -456,6 +581,9 @@ def _quantize_blockwise_reference( *, pow_2_scales: bool, row_scaled_nvfp4: bool = False, + nvfp4_use_4over6: bool = False, + nvfp4_e4m3_max: int = 448, + nvfp4_4over6_err_mode: str = "MAE", eps: float, # pylint: disable=unused-argument ) -> Tuple[torch.Tensor, torch.Tensor]: @@ -488,6 +616,10 @@ def _quantize_blockwise_reference( x = x.view(m, n // tile_len_x, tile_len_x) FLOAT4_E2M1_MAX = torch.tensor(6.0, device=x.device, dtype=torch.float32) FLOAT8_E4M3_MAX = torch.tensor(448.0, device=x.device, dtype=torch.float32) + global_scale_e4m3_max = float(nvfp4_e4m3_max if nvfp4_use_4over6 else 448) + GLOBAL_SCALE_E4M3_MAX = torch.tensor( + global_scale_e4m3_max, device=x.device, dtype=torch.float32 + ) decode_scale = torch.div(vec_max, FLOAT4_E2M1_MAX) if pow_2_scales: @@ -500,7 +632,7 @@ def _quantize_blockwise_reference( if row_scaled_nvfp4: global_amax = global_amax.to(torch.float32).view(m, 1, 1) - global_encode_scale = torch.div(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX, global_amax) + global_encode_scale = torch.div(GLOBAL_SCALE_E4M3_MAX * FLOAT4_E2M1_MAX, global_amax) global_encode_scale = torch.min( global_encode_scale, torch.tensor( @@ -519,6 +651,22 @@ def _quantize_blockwise_reference( global_encode_scale, ) global_decode_scale = torch.div(1.0, global_encode_scale) + if nvfp4_use_4over6: + # FourOverSix compares map-to-4 and map-to-6 candidates using + # the configured original input-domain error, while keeping TE-style FP4 + # quantization for each candidate. + return cls._quantize_blockwise_4over6_reference( + x, + vec_max, + global_amax, + global_encode_scale, + global_decode_scale, + row_scaled_nvfp4, + tile_len_y, + nvfp4_4over6_err_mode, + nvfp4_e4m3_max, + ) + global_encode_scale_multiplier = global_encode_scale * torch.reciprocal(FLOAT4_E2M1_MAX) # Match the kernel's default path: fold the FP4 reciprocal into the @@ -679,6 +827,9 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ self.quant_tile_shape[0], pow_2_scales=self.pow_2_scales, row_scaled_nvfp4=self.row_scaled_nvfp4, + nvfp4_use_4over6=self.nvfp4_use_4over6, + nvfp4_e4m3_max=self.nvfp4_e4m3_max, + nvfp4_4over6_err_mode=self.nvfp4_4over6_err_mode, eps=self.eps, ) if transpose_scales: @@ -702,6 +853,9 @@ def _quantize(self, tensor: torch.Tensor) -> Tuple[ self.quant_tile_shape[1], self.quant_tile_shape[0], pow_2_scales=self.pow_2_scales, + nvfp4_use_4over6=self.nvfp4_use_4over6, + nvfp4_e4m3_max=self.nvfp4_e4m3_max, + nvfp4_4over6_err_mode=self.nvfp4_4over6_err_mode, eps=self.eps, ) @@ -741,6 +895,8 @@ def quantize( scale_t=sx_t, global_amax_row=global_amax_row, global_amax_col=global_amax_col, + nvfp4_use_4over6=self.nvfp4_use_4over6, + nvfp4_e4m3_max=self.nvfp4_e4m3_max, dtype=tensor.dtype, device=tensor.device, quant_dtype=self.dtype, @@ -788,6 +944,8 @@ def update_quantized( dst.scale_t = sx_t dst.global_amax_row = global_amax_row dst.global_amax_col = global_amax_col + dst.nvfp4_use_4over6 = self.nvfp4_use_4over6 + dst.nvfp4_e4m3_max = self.nvfp4_e4m3_max dst.dtype = src.dtype dst.quant_dtype = self.dtype dst.original_shape = original_shape @@ -893,7 +1051,35 @@ def qgemm( sx = sx.to(torch.float32) sw = sw.to(torch.float32) - factor = 6.0 * 6.0 * 448.0 * 448.0 + qresult_x_nvfp4_use_4over6 = getattr( + qresult_x, + "nvfp4_use_4over6", + getattr(qresult_x, "_nvfp4_use_4over6", self.nvfp4_use_4over6), + ) + qresult_w_nvfp4_use_4over6 = getattr( + qresult_w, + "nvfp4_use_4over6", + getattr(qresult_w, "_nvfp4_use_4over6", self.nvfp4_use_4over6), + ) + qresult_x_e4m3_max = getattr( + qresult_x, + "nvfp4_e4m3_max", + getattr(qresult_x, "_nvfp4_e4m3_max", self.nvfp4_e4m3_max), + ) + qresult_w_e4m3_max = getattr( + qresult_w, + "nvfp4_e4m3_max", + getattr(qresult_w, "_nvfp4_e4m3_max", self.nvfp4_e4m3_max), + ) + if qresult_x_nvfp4_use_4over6: + fp8_max_x = float(qresult_x_e4m3_max) + else: + fp8_max_x = 448.0 + if qresult_w_nvfp4_use_4over6: + fp8_max_w = float(qresult_w_e4m3_max) + else: + fp8_max_w = 448.0 + factor = 6.0 * 6.0 * fp8_max_x * fp8_max_w if gemm_type == quantization.GEMMType.WGRAD: partial_alpha = qresult_x.global_amax_col * qresult_w.global_amax_col diff --git a/transformer_engine/pytorch/quantization.py b/transformer_engine/pytorch/quantization.py index 0c4072351..e503b4b56 100644 --- a/transformer_engine/pytorch/quantization.py +++ b/transformer_engine/pytorch/quantization.py @@ -1635,7 +1635,11 @@ def make_quantizers(self) -> list: * Forward, ``"weight"`` -> ``recipe.fp4_quant_fwd_weight``. * Forward, ``"input"`` / ``"output"`` (and any unknown forward type) -> ``recipe.fp4_quant_fwd_inp``. - * Backward, any slot -> ``recipe.fp4_quant_bwd_grad``. + * ``"grad_output"`` / ``"grad_input"`` -> ``recipe.fp4_quant_bwd_grad``. + * NVFP4 4over6 is applied to non-gradient slots selected by + ``recipe.nvfp4_4over6``. Gradient slots always use standard NVFP4, + which lets gradient RHT and stochastic rounding follow the base + recipe. When the owning module/op provides a role list via ``get_quantizer_roles``, the per-slot ``tensor_type`` drives dispatch. @@ -1647,7 +1651,7 @@ def make_quantizers(self) -> list: from .tensor.nvfp4_tensor import NVFP4Quantizer def _qparams(tensor_type: str): - if self.mode == "backward": + if tensor_type in ("grad_output", "grad_input"): return self.recipe.fp4_quant_bwd_grad if tensor_type == "weight": return self.recipe.fp4_quant_fwd_weight @@ -1655,6 +1659,34 @@ def _qparams(tensor_type: str): def _make(tensor_type: str) -> NVFP4Quantizer: qparams = _qparams(tensor_type) + nvfp4_use_4over6 = False + if tensor_type not in ("grad_output", "grad_input"): + if self.recipe.nvfp4_4over6 == "all": + nvfp4_use_4over6 = True + elif self.recipe.nvfp4_4over6 == "weights": + nvfp4_use_4over6 = tensor_type == "weight" + elif self.recipe.nvfp4_4over6 == "activations": + nvfp4_use_4over6 = tensor_type != "weight" + nvfp4_e4m3_max = 448 + if nvfp4_use_4over6: + # Current 4over6 kernels target RL and post-training quantization paths. + # Pre-training usage still needs a fused RHT + 4over6 quantization kernel. + if qparams.random_hadamard_transform: + raise ValueError("NVFP4 4over6 quantization does not support RHT.") + if qparams.stochastic_rounding: + raise ValueError( + "NVFP4 4over6 quantization does not support stochastic rounding." + ) + if self.recipe.nvfp4_4over6_e4m3_use_256 == "all": + nvfp4_e4m3_max = 256 + elif self.recipe.nvfp4_4over6_e4m3_use_256 == "weights": + if tensor_type == "weight": + nvfp4_e4m3_max = 256 + elif self.recipe.nvfp4_4over6_e4m3_use_256 == "activations": + if tensor_type != "weight": + nvfp4_e4m3_max = 256 + elif self.recipe.nvfp4_4over6_e4m3_use_256 == "none": + nvfp4_e4m3_max = 448 return NVFP4Quantizer( fp4_dtype=self.dtype, rowwise=True, @@ -1668,6 +1700,9 @@ def _make(tensor_type: str) -> NVFP4Quantizer: and tensor_type != "weight" and self.recipe.row_scaled_activation ), + nvfp4_use_4over6=nvfp4_use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, + nvfp4_4over6_err_mode=self.recipe.nvfp4_4over6_err_mode, ) if self.mode not in ("forward", "backward"): diff --git a/transformer_engine/pytorch/tensor/grouped_tensor.py b/transformer_engine/pytorch/tensor/grouped_tensor.py index f28f972b5..0cc03602a 100644 --- a/transformer_engine/pytorch/tensor/grouped_tensor.py +++ b/transformer_engine/pytorch/tensor/grouped_tensor.py @@ -93,6 +93,8 @@ def __new__( stride: Optional[List[int]] = None, with_gemm_swizzled_scales: bool = False, row_scaled_nvfp4: bool = False, + nvfp4_use_4over6: bool = False, + nvfp4_e4m3_max: int = 448, ): if ( shapes is not None @@ -166,6 +168,8 @@ def __new__( columnwise_scale_inv_offsets=columnwise_scale_inv_offsets, with_gemm_swizzled_scales=with_gemm_swizzled_scales, row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=nvfp4_use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, ) return instance @@ -198,6 +202,8 @@ def copy_grouped_storage_metadata(dst: GroupedTensor, src: GroupedTensor) -> Non dst.quantized_tensors = src.quantized_tensors dst._with_gemm_swizzled_scales = src._with_gemm_swizzled_scales dst.row_scaled_nvfp4 = src.row_scaled_nvfp4 + dst.nvfp4_use_4over6 = src.nvfp4_use_4over6 + dst.nvfp4_e4m3_max = src.nvfp4_e4m3_max def make_wrapper_like(src: GroupedTensor, requires_grad: bool) -> GroupedTensor: """Create a wrapper of the same type and tensor metadata as src.""" diff --git a/transformer_engine/pytorch/tensor/nvfp4_tensor.py b/transformer_engine/pytorch/tensor/nvfp4_tensor.py index 2ebefefaa..24962d67f 100644 --- a/transformer_engine/pytorch/tensor/nvfp4_tensor.py +++ b/transformer_engine/pytorch/tensor/nvfp4_tensor.py @@ -130,6 +130,12 @@ class NVFP4Quantizer(Quantizer): """Whether emitted NVFP4 tensors store one FP32 amax per row.""" row_scaled_nvfp4: bool + """Whether to use NVFP4 4over6 map-to-4/map-to-6 block selection.""" + nvfp4_use_4over6: bool + """Global E4M3 scale bound used by emitted NVFP4 tensors.""" + nvfp4_e4m3_max: int + """NVFP4 4over6 candidate-selection error mode.""" + nvfp4_4over6_err_mode: str """RHT matrix random sign mask""" rht_matrix_random_sign_mask_t: int @@ -147,6 +153,9 @@ def __init__( with_2d_quantization: bool = False, stochastic_rounding: bool = False, row_scaled_nvfp4: bool = False, + nvfp4_use_4over6: bool = False, + nvfp4_e4m3_max: int = 448, + nvfp4_4over6_err_mode: str = "MAE", with_random_sign_mask: bool = True, ) -> None: super().__init__(rowwise=rowwise, columnwise=columnwise) @@ -158,6 +167,13 @@ def __init__( self.with_2d_quantization = with_2d_quantization self.stochastic_rounding = stochastic_rounding self.row_scaled_nvfp4 = row_scaled_nvfp4 + self.nvfp4_use_4over6 = nvfp4_use_4over6 + self.nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_use_4over6 else 448 + if self.nvfp4_e4m3_max not in (448, 256): + raise ValueError("nvfp4_e4m3_max must be 448 or 256.") + self.nvfp4_4over6_err_mode = nvfp4_4over6_err_mode.upper() + if self.nvfp4_4over6_err_mode not in ("MAE", "MSE"): + raise ValueError("nvfp4_4over6_err_mode must be 'MAE' or 'MSE'.") self.rht_matrix_random_sign_mask_t = get_random_sign_mask_for_rht( with_random_sign_mask, torch.cuda.current_device() ) @@ -204,6 +220,9 @@ def copy(self) -> NVFP4Quantizer: with_2d_quantization=self.with_2d_quantization, stochastic_rounding=self.stochastic_rounding, row_scaled_nvfp4=self.row_scaled_nvfp4, + nvfp4_use_4over6=self.nvfp4_use_4over6, + nvfp4_e4m3_max=self.nvfp4_e4m3_max, + nvfp4_4over6_err_mode=self.nvfp4_4over6_err_mode, ) quantizer.internal = self.internal quantizer.optimize_for_gemm = self.optimize_for_gemm @@ -356,6 +375,8 @@ def __new__( quantizer: Quantizer, with_gemm_swizzled_scales: bool, row_scaled_nvfp4: bool = False, + nvfp4_use_4over6: bool = False, + nvfp4_e4m3_max: int = 448, **kwargs, ): instance = super().__new__( @@ -371,6 +392,8 @@ def __new__( with_gemm_swizzled_scales, *args, row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=nvfp4_use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, **kwargs, ) return instance @@ -528,6 +551,9 @@ def fsdp_pre_all_gather(self, mesh, orig_size, contiguous_orig_stride, module, m columnwise_usage, self._amax_rowwise, self._amax_columnwise, + self._row_scaled_nvfp4, + self._nvfp4_use_4over6, + self._nvfp4_e4m3_max, self.shape[-1], ) return sharded_tensors, metadata @@ -546,7 +572,16 @@ def fsdp_post_all_gather( all-gathered rowwise data. Columnwise data is derived locally via _create_columnwise() instead of being all-gathered. """ - fp4_dtype, columnwise_usage, amax_rowwise, amax_columnwise, K = metadata + ( + fp4_dtype, + columnwise_usage, + amax_rowwise, + amax_columnwise, + row_scaled_nvfp4, + nvfp4_use_4over6, + nvfp4_e4m3_max, + K, + ) = metadata # Only rowwise data+scales were all-gathered rowwise_data, rowwise_scale_inv = all_gather_outputs[:2] @@ -569,6 +604,9 @@ def fsdp_post_all_gather( out._rowwise_scale_inv = rowwise_scale_inv out._amax_rowwise = amax_rowwise out._amax_columnwise = amax_columnwise + out._row_scaled_nvfp4 = row_scaled_nvfp4 + out._nvfp4_use_4over6 = nvfp4_use_4over6 + out._nvfp4_e4m3_max = nvfp4_e4m3_max else: # Construct new tensor (first iteration) out = NVFP4Tensor( @@ -585,6 +623,9 @@ def fsdp_post_all_gather( requires_grad=False, with_gemm_swizzled_scales=False, device=rowwise_data.device, + row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=nvfp4_use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, ) # Derive columnwise data locally via transpose instead of all-gathering it @@ -724,6 +765,9 @@ def __torch_dispatch__(cls, func, types, args, kwargs=None): requires_grad=tensor.requires_grad, with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, device=tensor.device, + row_scaled_nvfp4=tensor._row_scaled_nvfp4, + nvfp4_use_4over6=tensor._nvfp4_use_4over6, + nvfp4_e4m3_max=tensor._nvfp4_e4m3_max, ) # Default case @@ -745,6 +789,9 @@ def __reduce_ex__(self, protocol: int) -> tuple: self.dtype, self._quantizer, self._with_gemm_swizzled_scales, + self._row_scaled_nvfp4, + self._nvfp4_use_4over6, + self._nvfp4_e4m3_max, ), ) @@ -837,6 +884,9 @@ def _set_data(self, tensor: torch.Tensor) -> None: self._amax_rowwise = tensor._amax_rowwise self._amax_columnwise = tensor._amax_columnwise self._with_gemm_swizzled_scales = tensor._with_gemm_swizzled_scales + self._row_scaled_nvfp4 = tensor._row_scaled_nvfp4 + self._nvfp4_use_4over6 = tensor._nvfp4_use_4over6 + self._nvfp4_e4m3_max = tensor._nvfp4_e4m3_max return # Quantize to FP8 @@ -889,7 +939,10 @@ def _make_nvfp4_tensor_in_reduce_ex( fp4_dtype: TE_DType, dtype: torch.dtype, quantizer: Quantizer, - with_gemm_swizzled_scales: bool = False, + with_gemm_swizzled_scales: bool, + row_scaled_nvfp4: bool = False, + nvfp4_use_4over6: bool = False, + nvfp4_e4m3_max: int = 448, ) -> NVFP4Tensor: """Reconstruct an ``NVFP4Tensor`` from its ``__reduce_ex__`` payload.""" # Infer device from whichever inner buffer is populated so the wrapper @@ -914,6 +967,9 @@ def _make_nvfp4_tensor_in_reduce_ex( requires_grad=False, with_gemm_swizzled_scales=with_gemm_swizzled_scales, device=device, + row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=nvfp4_use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, ) @@ -997,6 +1053,9 @@ def forward( requires_grad=tensor.requires_grad, with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, device=tensor.device, + row_scaled_nvfp4=tensor._row_scaled_nvfp4, + nvfp4_use_4over6=tensor._nvfp4_use_4over6, + nvfp4_e4m3_max=tensor._nvfp4_e4m3_max, ) @staticmethod @@ -1040,6 +1099,9 @@ def backward( requires_grad=grad.requires_grad, with_gemm_swizzled_scales=grad._with_gemm_swizzled_scales, device=grad.device, + row_scaled_nvfp4=grad._row_scaled_nvfp4, + nvfp4_use_4over6=grad._nvfp4_use_4over6, + nvfp4_e4m3_max=grad._nvfp4_e4m3_max, ) return dgrad, None return grad.view(ctx.shape), None @@ -1125,6 +1187,9 @@ def forward( requires_grad=tensor.requires_grad, with_gemm_swizzled_scales=tensor._with_gemm_swizzled_scales, device=tensor.device, + row_scaled_nvfp4=tensor._row_scaled_nvfp4, + nvfp4_use_4over6=tensor._nvfp4_use_4over6, + nvfp4_e4m3_max=tensor._nvfp4_e4m3_max, ) @staticmethod @@ -1168,6 +1233,9 @@ def backward( requires_grad=grad.requires_grad, with_gemm_swizzled_scales=grad._with_gemm_swizzled_scales, device=grad.device, + row_scaled_nvfp4=grad._row_scaled_nvfp4, + nvfp4_use_4over6=grad._nvfp4_use_4over6, + nvfp4_e4m3_max=grad._nvfp4_e4m3_max, ) return dgrad, None return grad.view(ctx.shape), None diff --git a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py index ac56d334b..438e12402 100644 --- a/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/grouped_tensor_storage.py @@ -73,6 +73,8 @@ def _initialize_storage_fields( stride: Optional[List[int]] = None, with_gemm_swizzled_scales: bool = False, row_scaled_nvfp4: bool = False, + nvfp4_use_4over6: bool = False, + nvfp4_e4m3_max: int = 448, ) -> None: """ Initialize a GroupedTensor. @@ -149,6 +151,8 @@ def _initialize_storage_fields( instance.quantized_tensors = None instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales instance.row_scaled_nvfp4 = row_scaled_nvfp4 + instance.nvfp4_use_4over6 = nvfp4_use_4over6 + instance.nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_use_4over6 else 448 def __new__( cls, @@ -175,6 +179,8 @@ def __new__( stride: Optional[List[int]] = None, with_gemm_swizzled_scales: bool = False, row_scaled_nvfp4: bool = False, + nvfp4_use_4over6: bool = False, + nvfp4_e4m3_max: int = 448, ): instance = object.__new__(cls) cls._initialize_storage_fields( @@ -201,6 +207,8 @@ def __new__( stride=stride, with_gemm_swizzled_scales=with_gemm_swizzled_scales, row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=nvfp4_use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, ) return instance @@ -307,6 +315,33 @@ def get_dtype(self) -> torch.dtype: return self.fake_dtype + @property + def row_scaled_nvfp4(self) -> bool: + """Whether grouped NVFP4 tensors use row-scaled amax metadata.""" + return self._row_scaled_nvfp4 + + @row_scaled_nvfp4.setter + def row_scaled_nvfp4(self, row_scaled_nvfp4: bool) -> None: + self._row_scaled_nvfp4 = row_scaled_nvfp4 + + @property + def nvfp4_use_4over6(self) -> bool: + """Whether grouped NVFP4 tensors carry 4over6 metadata.""" + return self._nvfp4_use_4over6 + + @nvfp4_use_4over6.setter + def nvfp4_use_4over6(self, nvfp4_use_4over6: bool) -> None: + self._nvfp4_use_4over6 = nvfp4_use_4over6 + + @property + def nvfp4_e4m3_max(self) -> int: + """Global E4M3 scale bound used by grouped NVFP4 tensors.""" + return self._nvfp4_e4m3_max + + @nvfp4_e4m3_max.setter + def nvfp4_e4m3_max(self, nvfp4_e4m3_max: int) -> None: + self._nvfp4_e4m3_max = nvfp4_e4m3_max + def prepare_for_saving( self, ) -> Tuple[list[Optional[torch.Tensor]], "GroupedTensorStorage"]: @@ -376,6 +411,8 @@ def clear(self) -> None: self.tensor_shapes = [] self.fake_dtype = torch.float32 self.row_scaled_nvfp4 = False + self.nvfp4_use_4over6 = False + self.nvfp4_e4m3_max = 448 def __repr__(self) -> str: """String representation of the GroupedTensorStorage.""" @@ -545,6 +582,8 @@ def copy(self) -> "GroupedTensorStorage": columnwise_scale_inv_offsets=self.columnwise_scale_inv_offsets, with_gemm_swizzled_scales=self._with_gemm_swizzled_scales, row_scaled_nvfp4=self.row_scaled_nvfp4, + nvfp4_use_4over6=self.nvfp4_use_4over6, + nvfp4_e4m3_max=self.nvfp4_e4m3_max, ) @staticmethod @@ -656,6 +695,8 @@ def make_grouped_tensor( scale_inv_offsets = None columnwise_scale_inv_offsets = None row_scaled_nvfp4 = False + nvfp4_use_4over6 = False + nvfp4_e4m3_max = 448 if no_quantization: assert dtype is not None, "dtype must be provided for unquantized GroupedTensor" if rowwise_usage: @@ -715,6 +756,8 @@ def make_grouped_tensor( amax = torch.empty(num_tensors, dtype=torch.float32, device=device) elif quantizer._get_compatible_recipe().nvfp4(): row_scaled_nvfp4 = quantizer.row_scaled_nvfp4 + nvfp4_use_4over6 = quantizer.nvfp4_use_4over6 + nvfp4_e4m3_max = quantizer.nvfp4_e4m3_max if row_scaled_nvfp4: if not rowwise_usage: raise ValueError( @@ -843,6 +886,8 @@ def make_grouped_tensor( quantizer.optimize_for_gemm if quantizer is not None else False ), row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=nvfp4_use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, ) grouped_tensor.quantized_tensors = grouped_tensor.split_into_quantized_tensors() return grouped_tensor @@ -957,6 +1002,8 @@ def split_into_quantized_tensors( self.columnwise_scale_inv_offsets = columnwise_scale_inv_offsets nvfp4_rowwise_amax_offsets = None row_scaled_nvfp4 = self.row_scaled_nvfp4 + nvfp4_use_4over6 = self.nvfp4_use_4over6 + nvfp4_e4m3_max = self.nvfp4_e4m3_max if recipe.nvfp4() and row_scaled_nvfp4: cum = 0 nvfp4_rowwise_amax_offsets = [0] @@ -1184,6 +1231,8 @@ def split_into_quantized_tensors( quantizer=quantizer, with_gemm_swizzled_scales=quantizer.optimize_for_gemm, row_scaled_nvfp4=row_scaled_nvfp4, + nvfp4_use_4over6=nvfp4_use_4over6, + nvfp4_e4m3_max=nvfp4_e4m3_max, ) result.append(tensor) diff --git a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py index 490184e5f..250fa6bdb 100644 --- a/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py +++ b/transformer_engine/pytorch/tensor/storage/nvfp4_tensor_storage.py @@ -104,6 +104,10 @@ class NVFP4TensorStorage(QuantizedTensorStorage): _with_gemm_swizzled_scales: bool # Whether this NVFP4 tensor uses row-scaled amax metadata _row_scaled_nvfp4: bool + # Whether this NVFP4 tensor uses 4over6 map-to-4/map-to-6 block selection + _nvfp4_use_4over6: bool + # Global E4M3 scale bound used by this NVFP4 tensor + _nvfp4_e4m3_max: int def __new__( cls, @@ -119,6 +123,8 @@ def __new__( *args, fake_dtype: Optional[torch.dtype] = None, row_scaled_nvfp4: bool = False, + nvfp4_use_4over6: bool = False, + nvfp4_e4m3_max: int = 448, **kwargs, ): if cls is NVFP4TensorStorage: @@ -137,6 +143,8 @@ def __new__( instance._amax_columnwise = amax_columnwise instance._with_gemm_swizzled_scales = with_gemm_swizzled_scales instance._row_scaled_nvfp4 = row_scaled_nvfp4 + instance._nvfp4_use_4over6 = nvfp4_use_4over6 + instance._nvfp4_e4m3_max = nvfp4_e4m3_max if nvfp4_use_4over6 else 448 return instance @@ -163,6 +171,10 @@ def copy_from_storage(self, src: QuantizedTensorStorage) -> None: raise RuntimeError("Scale layout mismatch in copy_from_storage") if self._row_scaled_nvfp4 != src._row_scaled_nvfp4: raise RuntimeError("Rowwise amax scaling mode mismatch in copy_from_storage") + if self._nvfp4_use_4over6 != src._nvfp4_use_4over6: + raise RuntimeError("NVFP4 4over6 mode mismatch in copy_from_storage") + if self._nvfp4_e4m3_max != src._nvfp4_e4m3_max: + raise RuntimeError("NVFP4 4over6 E4M3 scale bound mismatch in copy_from_storage") def _copy_optional(dst: Optional[torch.Tensor], src_tensor: Optional[torch.Tensor]): if dst is not None and src_tensor is not None: @@ -188,6 +200,8 @@ def get_metadata(self) -> Dict[str, Any]: "quantizer": self._quantizer, "with_gemm_swizzled_scales": self._with_gemm_swizzled_scales, "row_scaled_nvfp4": self._row_scaled_nvfp4, + "nvfp4_use_4over6": self._nvfp4_use_4over6, + "nvfp4_e4m3_max": self._nvfp4_e4m3_max, "fake_dtype": self._dtype, } @@ -321,6 +335,8 @@ def view(self, shape: torch.Size): fp4_dtype=self._fp4_dtype, with_gemm_swizzled_scales=self._with_gemm_swizzled_scales, row_scaled_nvfp4=self._row_scaled_nvfp4, + nvfp4_use_4over6=self._nvfp4_use_4over6, + nvfp4_e4m3_max=self._nvfp4_e4m3_max, fake_dtype=self._dtype, ) From 80ea3133efa4c3a3679845b8ee46dfc06e2792f0 Mon Sep 17 00:00:00 2001 From: Sudhakar Singh Date: Fri, 22 May 2026 18:45:01 -0700 Subject: [PATCH 080/180] [PyTorch] Add `pad_between_seqs` support for non-CP and CP (A2A and P2P) with FA3 + THD (varlen) (#2596) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * [PyTorch] Add pad_between_seqs support for FlashAttention 3 with CP Add support for padding between sequences (pad_between_seqs) in the FlashAttention 3 backend when used with context parallelism (CP). Key changes: - backends.py: Pass fa_pad_between_seqs through to FA3 forward/backward - context_parallel.py: Handle pad_between_seqs in A2A and P2P CP paths, zero FA3 padding garbage in CP forward, fix a2a backward alignment - dot_product_attention.py: Auto-detect pad_between_seqs from cu_seqlens - utils.py: Gate FA3 deterministic backward for hdim>=256, fix flash_attn_supported override for cross-attention and large head_dim, disable UnfusedDotProductAttention for pad_between_seqs, add SM100+ FA3 skip Signed-off-by: Sudhakar Singh * [PyTorch] Add pad_between_seqs tests for CP and non-CP FlashAttention Add test parametrization for pad_between_seqs in flash attention tests. Update run_attention_with_cp.py to support the new parameter and fix batch boundary alignment in the non-CP FA3 path. Run tests in parallel when multiple GPUs are available. Signed-off-by: Sudhakar Singh * [QA] Add CP deterministic tests to L3 and support TE_PATH in FA test Add deterministic CP test runs to L3 FA versions test. Support TE_PATH positional arg and fix GPU threshold for parallel test execution. Signed-off-by: Sudhakar Singh * [PyTorch] Fix FA3 deterministic gate to match upstream backward constraint The previous check disabled FA3 for deterministic mode whenever head_dim_qk > 128, which was overly conservative — FA3 forward supports deterministic execution at any head dim. The actual constraint from flash_api.cpp is that the backward pass does not support deterministic mode when max(head_size, head_size_v) >= 256. Narrow the gate to only disable FA3 during training (backward) and raise the threshold to >= 256, checking both head_dim_qk and head_dim_v to handle MLA configs with asymmetric head dimensions. Ref: https://github.com/Dao-AILab/flash-attention/blob/ac6f2eb5/hopper/flash_api.cpp#L1370 Signed-off-by: Sudhakar Singh * [PyTorch] Disable FlashAttention 4 for pad_between_seqs with THD The pad_between_seqs gate in get_attention_backend only disabled FlashAttention 2, letting FA4 leak through to the test-time fused-vs-flash comparison. On B200 runners that install flash-attn-4, this caused test_dpa_qkv_layout_thd to compare FusedAttention against an FA4 output whose padded positions contain garbage, producing 48 numerics failures in L3_pytorch_FA_versions_test--B200_1GPU. The log message already claimed FA4 would be disabled — this change makes the code match the message: set use_flash_attention_4 = False alongside use_flash_attention_2 when pad_between_seqs is True. FA3 continues to support pad_between_seqs via seqused_q/seqused_k. Signed-off-by: Sudhakar Singh * [QA] Fix cutlass-dsl utils shadow in FA versions test FA4 install brings in nvidia-cutlass-dsl, whose `import cutlass` adds cutlass/base_dsl/ to sys.path. That directory contains a utils/ package that shadows tests/pytorch/utils.py, breaking collection of test_attention_with_cp.py with: ImportError: cannot import name 'ModelConfig' from 'utils' Prepend $TE_PATH/tests/pytorch to PYTHONPATH so the local utils.py is always resolved first, regardless of what FA4 dependencies install. Signed-off-by: Sudhakar Singh * skip tests which OOM in deterministic+backward+hopper+large_configs as its a known cudnn issue Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * make cp det and nondet tests run in parallel whenever possible Signed-off-by: Sudhakar Singh * [QA] L3: gate CP tests per-arch to avoid CI timeout PR 2596 added deterministic CP runs to the L3 FA-versions matrix, multiplying CP wall time across every FA version and causing CI timeouts (pipeline 50243000). Run CP tests once per arch instead, picking the FA version each arch's CP code path actually supports: - sm90 (H100): FA3 3.0.0b1 - context_parallel.py is FA3-only on Hopper (use_flash_attn_3 threaded throughout, FA4 not wired in; pad_between_seqs gated on use_flash_attn_3 at lines 1038, 1366) - sm>90 (B200): latest FA4 - FA3 is not built/installed for sm>90 Non-CP test_attention.py still runs for every FA version in the array. Also drop FA 2.7.3 from the sm90 list (no longer maintained as a target) and bump the FA4 pin from 4.0.0b8 to 4.0.0b11. b8 has an SM90 backward kernel bug fixed by upstream PR #2513 in b11 (get_smem_store_C() got multiple values for argument 'transpose'). Signed-off-by: Sudhakar Singh * [QA] L3: skip pre-installed FA3 build, per-FA junit XMLs Three follow-ups on top of 13ba0046 (L3 per-arch CP gating): 1. Skip the inline FA3 source build when flash_attn_interface is already importable. This makes the script a no-op on FA3 install when the base image has FA3 baked in (companion to TE !573 on te_ci, which auto-sets INSTALL_FA3=${RUN_L3_TESTS} so FA3 is preinstalled for L3 pipelines). Saves ~20 min of L3 H100 wall time once both land. Falls back to the existing inline build when FA3 is not pre-installed. 2. Suffix junit XMLs with the FA version (pytest_test_attention_fa2_8_3.xml etc.) so per-iteration results are preserved instead of overwritten. Pipeline 50348672 had no per-FA timing visibility because pytest.xml was clobbered by each loop iteration. 3. Include FA version in test_fail messages so CI dashboards show which FA iteration caused a failure (was "test_attention.py", now "test_attention.py (FA 2.8.3)"). Also fold the CP_FA_VERSION assignment into the same if-block as FA_versions (was a separate if-block immediately after) since the two are arch-keyed in lockstep. Signed-off-by: Sudhakar Singh * b200 shouldnt run FA3 even if present Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * L3: drop stale RUN_L3_TESTS=1 note; use flash_attn_3 for FA3 check Address two pending review comments: 1. The "auto-set when RUN_L3_TESTS=1" annotation on the base-image FA3 preinstall is no longer accurate; drop it so readers don't grep for a coupling that doesn't exist. 2. `flash_attn_interface` reads like a generic FA API even though the top-level shim is only created by the FA3 install. Switching to `import flash_attn_3` makes the FA3-specific intent unambiguous and matches the FA3 package layout produced by the source build. Local validation on H100 (sm90) with FA3 active, TE worktree resolving to the editable install (verified via three-layer import check from /tmp): test_attention_with_cp.py parallel det+nondet — 45 passed / 0 failed nondet (3:52), 33 passed / 0 failed det (2:55). 33 pad-True nondet passes + 21 pad-True det passes confirm the FA3+THD+CP path is exercised; 5 det OOM cases skip cleanly via the existing inline guard. Same test scope is exercised by L1_pytorch_distributed_unittest (parallel det+nondet) and the FA3 iteration of L3_pytorch_FA_versions_test; the changes here are L3-only documentation/detection tweaks and do not alter the Python test code, but the L1+L3 CP execution was re-run on the cleaned PR head end-to-end as proof. Signed-off-by: Sudhakar Singh * Address review nits: bHSS-gated OOM skip; drop Dockerfile.base specifics 1. Det FusedAttention backward THD/sm90 OOM skip: gate on the actual memory pressure (b*H*S*S) instead of num_heads >= 20. The cuDNN workspace is proportional to bHSS, so a future config with H >= 20 but small b or S would be needlessly skipped under the old guard, while a config with H < 20 but large b*S that hit the same OOM wouldn't be caught. Threshold 1e9 empirically matches the existing 5-case skip set on the test_essential fused subset (cp_2_0, cp_2_2, cp_3_1, cp_4_2, cp_4_3 — bHSS in 1.07B–4.29B) and lets cp_1_0/ cp_2_1/cp_2_4/cp_3_2/cp_3_4 (bHSS ~0.40B) keep running. 2. L3 FA3 install comment: drop the "Dockerfile.base INSTALL_FA3=1" reference. The detection check is the contract; mentioning a specific image variable couples this script to an out-of-tree provisioning detail that may evolve independently. Local validation on H100 (sm90) with FA3 active and TE worktree resolving to editable (verified via /tmp-cwd three-layer import check after reinstall — the /usr/local TE shadow had reappeared between sessions): test_attention_with_cp.py parallel det+nondet — 45 passed / 0 failed nondet (4:09), 33 passed / 0 failed det (3:14). 33 pad-True nondet passes + 21 pad-True det passes; 5 det OOM cases skip via the new bHSS gate — same cases as the old num_heads-only gate. Signed-off-by: Sudhakar Singh * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Name the OOM-skip threshold and explain the 128*bHSS workspace observation Address review nits on the deterministic THD-backward OOM guard: 1. Replace the magic number 1_000_000_000 with the named constant SM90_DET_FUSED_THD_BWD_MAX_BHSS = 1 << 30, so the value is searchable and labeled. 2. Replace the prefatory comment with a short note tying the number to cuDNN's actual workspace request (~128 * bHSS bytes, measured on cuDNN 9.21.0 sm90 — see local sweep). At bHSS = 1<<30 the request is 128 GiB, which doesn't fit on H100's 80 GB. 3. Flag the b>=3 caveat for future readers: cuDNN rounds the batch up internally so workspace grows super-linearly past b=2 (b=4 asks for 4x the b=2 workspace, not 2x). The current fused-essential matrix is all b=2, so the threshold stays correct for what the test exercises; the note is there so the next person doesn't have to rediscover it. Skip set is unchanged — cp_2_0, cp_2_1, cp_3_1, cp_4_2, cp_4_3. Signed-off-by: Sudhakar Singh * Reword OOM-skip comment as observations, not cuDNN-internal claims We measured the workspace request from outside cuDNN, so the comment should say "observed" rather than asserting what cuDNN does. Reframes the ~128 * bHSS bytes formula and the super-linear b>=3 behavior as empirical observations from our sweep. No code change. Signed-off-by: Sudhakar Singh --------- Signed-off-by: Sudhakar Singh Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- qa/L1_pytorch_distributed_unittest/test.sh | 19 ++- qa/L3_pytorch_FA_versions_test/test.sh | 85 ++++++++++- .../attention/run_attention_with_cp.py | 96 +++++++----- tests/pytorch/attention/test_attention.py | 34 ++--- .../attention/test_attention_with_cp.py | 30 +++- .../dot_product_attention/backends.py | 32 +++- .../dot_product_attention/context_parallel.py | 141 +++++++++++++++--- .../dot_product_attention.py | 3 + .../attention/dot_product_attention/utils.py | 38 +++-- 9 files changed, 371 insertions(+), 107 deletions(-) diff --git a/qa/L1_pytorch_distributed_unittest/test.sh b/qa/L1_pytorch_distributed_unittest/test.sh index db13e9f1e..7eb34a62e 100644 --- a/qa/L1_pytorch_distributed_unittest/test.sh +++ b/qa/L1_pytorch_distributed_unittest/test.sh @@ -22,6 +22,24 @@ mkdir -p "$XML_LOG_DIR" pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" +# Run CP tests (deterministic + non-deterministic) first so they can be parallelized. +# Each needs 4 GPUs, so >=8 GPUs allows them to run concurrently on disjoint GPU sets. +NUM_GPUS=$(python3 -c "import torch; print(torch.cuda.device_count())") +echo "Detected $NUM_GPUS GPU(s)" +if [ "$NUM_GPUS" -ge 8 ]; then + echo "Running CP tests in parallel: non-deterministic on GPUs 0-3, deterministic on GPUs 4-7" + CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & + PID_CP_NONDET=$! + CUDA_VISIBLE_DEVICES=4,5,6,7 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & + PID_CP_DET=$! + wait $PID_CP_NONDET || test_fail "test_attention_with_cp.py" + wait $PID_CP_DET || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py" +else + echo "Running CP tests sequentially: need >=8 GPUs for parallel execution" + python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py" + NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_deterministic_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 test_attention_with_cp.py" +fi + python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_sanity.xml $TE_PATH/tests/pytorch/distributed/test_sanity.py || test_fail "test_sanity.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/distributed/test_numerics.py || test_fail "test_numerics.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_numerics_exact.xml $TE_PATH/tests/pytorch/distributed/test_numerics_exact.py || test_fail "test_numerics_exact.py" @@ -29,7 +47,6 @@ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops.xml $TE_ python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_torch_fsdp2.xml $TE_PATH/tests/pytorch/distributed/test_torch_fsdp2.py || test_fail "test_torch_fsdp2.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_comm_gemm_overlap.xml $TE_PATH/tests/pytorch/distributed/test_comm_gemm_overlap.py || test_fail "test_comm_gemm_overlap.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_fusible_ops_with_userbuffers.xml $TE_PATH/tests/pytorch/distributed/test_fusible_ops_with_userbuffers.py || test_fail "test_fusible_ops_with_userbuffers.py" -python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_attention_with_cp.xml $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cp_utils.xml $TE_PATH/tests/pytorch/attention/test_cp_utils.py || test_fail "test_cp_utils.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_cast_master_weights_to_fp8.xml $TE_PATH/tests/pytorch/distributed/test_cast_master_weights_to_fp8.py || test_fail "test_cast_master_weights_to_fp8.py" python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest_test_newton_schulz.xml $TE_PATH/tests/pytorch/distributed/test_newton_schulz.py || test_fail "test_newton_schulz.py" diff --git a/qa/L3_pytorch_FA_versions_test/test.sh b/qa/L3_pytorch_FA_versions_test/test.sh index 642eb93b0..30f1fc38c 100644 --- a/qa/L3_pytorch_FA_versions_test/test.sh +++ b/qa/L3_pytorch_FA_versions_test/test.sh @@ -2,13 +2,25 @@ # # See LICENSE for license information. -set -e +function error_exit() { + echo "Error: $1" + exit 1 +} + +function test_fail() { + RET=1 + FAILED_CASES="$FAILED_CASES $1" + echo "Error: sub-test failed: $1" +} + +RET=0 +FAILED_CASES="" : ${TE_PATH:=/opt/transformerengine} : ${XML_LOG_DIR:=/logs} mkdir -p "$XML_LOG_DIR" -pip3 install pytest==8.2.1 +pip3 install pytest==8.2.1 || error_exit "Failed to install pytest" # Limit parallel build jobs to avoid overwhelming system resources export MAX_JOBS=32 @@ -16,12 +28,18 @@ export MAX_JOBS=32 # Iterate over Flash Attention versions sm_arch=`python3 -c "import torch; sm = torch.cuda.get_device_capability(0); print(sm[0]*10+sm[1])"` export FLASH_ATTN_CUDA_ARCHS=$sm_arch +# CP tests are expensive and run only once per arch: +# - sm90 (H100): FA3 (3.0.0b1) - context_parallel.py only supports FA3 on Hopper +# - sm>90 (B200): latest FA4 - FA3 is not built/installed for sm>90 +# Non-CP tests still run for every FA version in the array. if [ $sm_arch -gt 90 ] then - FA_versions=(2.8.3 4.0.0b8) + FA_versions=(2.8.3 4.0.0b11) + CP_FA_VERSION="${FA_versions[-1]}" elif [ $sm_arch -eq 90 ] then - FA_versions=(2.7.3 2.8.3 3.0.0b1 4.0.0b8) + FA_versions=(2.8.3 3.0.0b1 4.0.0b11) + CP_FA_VERSION="3.0.0b1" fi for fa_version in "${FA_versions[@]}" @@ -35,12 +53,63 @@ do then pip3 install flash-attn-4==${fa_version} nvidia-cutlass-dsl[cu13]==4.4.2 --no-build-isolation else - git clone https://github.com/Dao-AILab/flash-attention.git - cd flash-attention/hopper && python setup.py install - cd ../../ + # FA3 source build (~20 min). Skip if FA3 is already installed. + if python3 -c "import flash_attn_3" 2>/dev/null; then + echo "FA3 already installed (from base image); skipping source build" + else + git clone https://github.com/Dao-AILab/flash-attention.git + cd flash-attention/hopper && python setup.py install + cd ../../ + fi fi + # Ensure local test utils is found before nvidia-cutlass-dsl's utils package + export PYTHONPATH=$TE_PATH/tests/pytorch:${PYTHONPATH:-} + # Run tests - NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s --junitxml=$XML_LOG_DIR/pytest.xml $TE_PATH/tests/pytorch/attention/test_attention.py + NUM_GPUS=$(nvidia-smi -L | wc -l) + echo "Detected $NUM_GPUS GPU(s)" + + # Suffix junit XMLs with the FA version so per-iteration results are preserved + # (otherwise pytest.xml is overwritten on each loop iteration and we lose timing + # data for all but the last FA version). + fa_tag="${fa_version//./_}" + XML_ATTN="$XML_LOG_DIR/pytest_test_attention_fa${fa_tag}.xml" + XML_CP="$XML_LOG_DIR/pytest_test_attention_with_cp_fa${fa_tag}.xml" + + if [ "$fa_version" = "$CP_FA_VERSION" ]; then + echo "Running CP tests with FA $fa_version (CP version for sm$sm_arch)" + if [ "$NUM_GPUS" -ge 5 ]; then + CP_NUM_GPUS=$(( NUM_GPUS - 1 > 4 ? 4 : NUM_GPUS - 1 )) + CP_GPUS=$(seq -s, 1 $CP_NUM_GPUS) + echo "Running tests in parallel: test_attention.py on GPU 0, test_attention_with_cp.py on GPUs $CP_GPUS ($CP_NUM_GPUS GPUs)" + + CUDA_VISIBLE_DEVICES=0 NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s \ + --junitxml=$XML_ATTN \ + $TE_PATH/tests/pytorch/attention/test_attention.py & + PID_ATTN=$! + CUDA_VISIBLE_DEVICES=$CP_GPUS NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s \ + --junitxml=$XML_CP \ + $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py & + PID_CP=$! + + wait $PID_ATTN || test_fail "test_attention.py (FA $fa_version)" + wait $PID_CP || test_fail "test_attention_with_cp.py (FA $fa_version)" + else + echo "Running tests sequentially: need >=5 GPUs for parallel execution (1 for test_attention + 4 for test_attention_with_cp)" + NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s --junitxml=$XML_ATTN $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py (FA $fa_version)" + NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s --junitxml=$XML_CP $TE_PATH/tests/pytorch/attention/test_attention_with_cp.py || test_fail "test_attention_with_cp.py (FA $fa_version)" + fi + else + echo "Skipping CP tests for FA $fa_version (CP only runs with FA $CP_FA_VERSION on sm$sm_arch)" + NVTE_TORCH_COMPILE=0 python3 -m pytest -v -s --junitxml=$XML_ATTN $TE_PATH/tests/pytorch/attention/test_attention.py || test_fail "test_attention.py (FA $fa_version)" + fi done + +if [ "$RET" -ne 0 ]; then + echo "Error in the following test cases:$FAILED_CASES" + exit 1 +fi +echo "All tests passed" +exit 0 diff --git a/tests/pytorch/attention/run_attention_with_cp.py b/tests/pytorch/attention/run_attention_with_cp.py index 9f6b4944e..6fca61d3c 100644 --- a/tests/pytorch/attention/run_attention_with_cp.py +++ b/tests/pytorch/attention/run_attention_with_cp.py @@ -47,6 +47,7 @@ def generate_input_shapes( config: ModelConfig, world_size: int, kernel_backend: str, + fa_pad_between_seqs: str = "False", ): if qkv_format == "bshd": q_input_shape = ( @@ -115,9 +116,12 @@ def generate_input_shapes( ).cuda() cu_seqlens_q = torch.clone(cu_seqlens_q_padded) - # Since FlashAttention doesn't support pad b/w sequences, and FusedAttention does, - # cu_seqlens_q is updated to reflect non-padded lengths for FusedAttention only. - if kernel_backend == "FusedAttention": + # Generate padded data (cu_seqlens_q reflects non-padded lengths, so it + # differs from cu_seqlens_q_padded) for FusedAttention always, and for + # FlashAttention only when its test param requests it. DPA auto-detects + # pad_between_seqs downstream from the cu_seqlens_q vs cu_seqlens_q_padded + # mismatch. + if kernel_backend == "FusedAttention" or fa_pad_between_seqs == "True": cu_seqlens_q[1:] = seqlens_q.cumsum(0, dtype=torch.int32).cuda() # NOTE: In case of Cross-Attention, `cu_seqlens_kv` and `cu_seqlens_kv_padded` @@ -196,6 +200,7 @@ def run_dpa_with_cp( scaling_mode="delayed", f16_O="False", is_training="True", + fa_pad_between_seqs="False", deterministic="False", log_level=logging.WARNING, ): @@ -314,7 +319,7 @@ def run_dpa_with_cp( cu_seqlens_kv, cu_seqlens_q_padded, cu_seqlens_kv_padded, - ) = generate_input_shapes(qkv_format, config, world_size, kernel_backend) + ) = generate_input_shapes(qkv_format, config, world_size, kernel_backend, fa_pad_between_seqs) q_orig = torch.clamp(torch.randn(q_input_shape, dtype=dtypes[dtype]), min=-1, max=1).cuda() k_orig = torch.clamp(torch.randn(k_input_shape, dtype=dtypes[dtype]), min=-1, max=1).cuda() v_orig = torch.clamp(torch.randn(v_input_shape, dtype=dtypes[dtype]), min=-1, max=1).cuda() @@ -557,11 +562,11 @@ def run_dpa_with_cp( tensors_to_deq[i] = tensor.dequantize() if not fp8_bwd: tensors[0], tensors[5] = tensors_to_deq - for i, tensor in enumerate(tensors): + for tensor, name in zip(tensors, names): # dbias/dbias_ could be None, so skip check for it if tensor is not None: - assert torch.all(~torch.isnan(tensor)), f"{names[i]} contains NaN" - assert torch.all(~torch.isinf(tensor)), f"{names[i]} contains Inf" + assert torch.all(~torch.isnan(tensor)), f"{name} has nan values" + assert torch.all(~torch.isinf(tensor)), f"{name} has inf values" out, dq, dk, dv, dbias, out_, dq_, dk_, dv_, dbias_ = tensors ############ compare results between CP and no-CP ############ @@ -617,49 +622,60 @@ def run_dpa_with_cp( if is_training: dq, out = [x.index_select(0, seq_idx_q).contiguous() for x in [dq, out]] dk, dv = [x.index_select(0, seq_idx_kv).contiguous() for x in [dk, dv]] - dq_, dk_, dv_, out_ = [dq_, dk_, dv_, out_] cu_seqlens_q_padded = cu_seqlens_q_padded // world_size cu_seqlens_q = get_cu_seqlens_on_cp_rank( cu_seqlens_q, cu_seqlens_q_padded, world_size, rank, True, True ) - cu_pads_q = cu_seqlens_q_padded - cu_seqlens_q - num_pads_q = cu_pads_q[1:] - cu_pads_q[:-1] - for x in [dq, out, dq_, out_]: - assert torch.count_nonzero(x[cu_seqlens_q_padded[-1] :]).item() == 0 - for b in range(config.batch_size): - assert ( - num_pads_q[b] == 0 - or torch.count_nonzero( - x[ - (cu_seqlens_q_padded[b + 1] - num_pads_q[b]) : cu_seqlens_q_padded[ - b + 1 - ] - ] - ).item() - == 0 - ) + num_pads_q = (cu_seqlens_q_padded - cu_seqlens_q)[1:] - ( + cu_seqlens_q_padded - cu_seqlens_q + )[:-1] cu_seqlens_kv_padded = cu_seqlens_kv_padded // world_size cu_seqlens_kv = get_cu_seqlens_on_cp_rank( cu_seqlens_kv, cu_seqlens_kv_padded, world_size, rank, True, True ) - cu_pads_kv = cu_seqlens_kv_padded - cu_seqlens_kv - num_pads_kv = cu_pads_kv[1:] - cu_pads_kv[:-1] - for x in [dk, dv, dk_, dv_]: - assert torch.count_nonzero(x[cu_seqlens_kv_padded[-1] :]).item() == 0 - for b in range(config.batch_size): - assert ( - num_pads_kv[b] == 0 - or torch.count_nonzero( - x[ - ( - cu_seqlens_kv_padded[b + 1] - num_pads_kv[b] - ) : cu_seqlens_kv_padded[b + 1] - ] - ).item() - == 0 + num_pads_kv = (cu_seqlens_kv_padded - cu_seqlens_kv)[1:] - ( + cu_seqlens_kv_padded - cu_seqlens_kv + )[:-1] + # FA3 leaves garbage at padding positions despite seqused_q/k (tile spillover). + # Forward out_ can't be pre-zeroed because FA3's custom op returns out_ as an + # output rather than mutating it in-place, triggering PyTorch's aliasing constraint. + # Backward dq/dk/dv CAN be pre-zeroed because FA3 marks them as mutated inputs. + if fa_pad_between_seqs == "True": + # out_ is a view inside the CP custom autograd Function, so in-place + # zeroing is blocked by PyTorch. Clone to break the view relationship. + out_ = out_.clone() + for x in [out, out_, dq]: + for b in range(config.batch_size): + x[ + cu_seqlens_q_padded[b + 1] - num_pads_q[b] : cu_seqlens_q_padded[b + 1] + ] = 0.0 + x[cu_seqlens_q_padded[-1] :] = 0.0 + for x in [dk, dv]: + for b in range(config.batch_size): + x[ + cu_seqlens_kv_padded[b + 1] + - num_pads_kv[b] : cu_seqlens_kv_padded[b + 1] + ] = 0.0 + x[cu_seqlens_kv_padded[-1] :] = 0.0 + # Verify CP backward tensors have clean padding (pre-zeroed in context_parallel.py). + for xname, x, cu, np_ in [ + ("dq_", dq_, cu_seqlens_q_padded, num_pads_q), + ("dk_", dk_, cu_seqlens_kv_padded, num_pads_kv), + ("dv_", dv_, cu_seqlens_kv_padded, num_pads_kv), + ]: + nnz = torch.count_nonzero(x[cu[-1] :]).item() + assert nnz == 0, ( + f"{xname} has {nnz} nonzero values in tail padding — " + "context_parallel.py should zero padding positions" ) + for b in range(config.batch_size): + if np_[b] > 0: + nnz = torch.count_nonzero(x[cu[b + 1] - np_[b] : cu[b + 1]]).item() + assert nnz == 0, ( + f"{xname} has {nnz} nonzero values in batch {b} padding — " + "context_parallel.py should zero padding positions" + ) else: - # Forward-only: reshape only out/out_ for comparison out = out.index_select(0, seq_idx_q).contiguous() out_ = out_ diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 32ea1694e..5c46949f6 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -124,7 +124,7 @@ def reset_global_fp8_state(): @pytest.mark.parametrize("workspace_opt", [True, False]) @pytest.mark.parametrize("qkv_layout", [None]) @pytest.mark.parametrize("swa", [False]) -@pytest.mark.parametrize("pad_between_seqs", [False]) +@pytest.mark.parametrize("pad_between_seqs", [False, True]) def test_dot_product_attention( dtype, model_configs, @@ -157,6 +157,8 @@ def test_dot_product_attention( config.window_size = check_set_window_size(config.attn_mask_type, config.window_size) qkv_format = qkv_layout.replace("3", "").replace("2", "").split("_")[0] + if pad_between_seqs and qkv_format != "thd": + pytest.skip("pad_between_seqs only applies to THD format!") if qkv_format == "thd" and "padding" not in config.attn_mask_type: config.attn_mask_type = ( "padding_" + config.attn_mask_type if config.attn_mask_type != "no_mask" else "padding" @@ -195,19 +197,6 @@ def test_dot_product_attention( ) flash_attn_supported, fused_attn_supported, unfused_attn_supported = available_backends - # FlashAttention does not support pad_between_seqs, but _run_dot_product_attention - # mannually pads and unpads the input and output of FlashAttention for testing purposes - if ( - pad_between_seqs - and FlashAttentionUtils.is_installed - and not ( - config.max_seqlen_q != config.max_seqlen_kv - and config.attn_mask_type in ["causal", "padding_causal"] - ) - and (config.window_size[0] == -1 or FlashAttentionUtils.v2_3_plus) - ): - flash_attn_supported = True - # Skip if only unfused backend is supported if (len(fused_attn_backends) + flash_attn_supported + unfused_attn_supported) < 2: pytest.skip("Less than two backends to compare.") @@ -1301,12 +1290,12 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: block.softmax_offset.requires_grad = True # Run a forward and backward pass - if backend in ["FlashAttention", "UnfusedDotProductAttention"]: + if backend in ["UnfusedDotProductAttention"]: q = inp_orig[0] k = inp_orig[1] v = inp_orig[2] d_out = out_grad_orig - if backend == "FusedAttention": + if backend in ["FusedAttention", "FlashAttention"]: q = inp[0] k = inp[1] v = inp[2] @@ -1322,14 +1311,19 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: max_seqlen_kv=config.max_seqlen_kv, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, - cu_seqlens_q_padded=cu_seqlens_q_after_pad if backend == "FusedAttention" else None, - cu_seqlens_kv_padded=cu_seqlens_kv_after_pad if backend == "FusedAttention" else None, + cu_seqlens_q_padded=( + cu_seqlens_q_after_pad if backend in ["FusedAttention", "FlashAttention"] else None + ), + cu_seqlens_kv_padded=( + cu_seqlens_kv_after_pad if backend in ["FusedAttention", "FlashAttention"] else None + ), attn_mask_type=config.attn_mask_type, checkpoint_core_attention=ckpt_attn, core_attention_bias_type=config.attn_bias_type, core_attention_bias=bias, alibi_slopes=alibi_slopes, fast_zero_fill=True, + pad_between_seqs=pad_between_seqs, # Only pass num_splits when exercising the FlashAttention path num_splits=config.num_splits if backend == "FlashAttention" else 1, ) @@ -1343,12 +1337,12 @@ def get_dummy_cuda_rng_tracker() -> CudaRNGStatesTracker: if is_training and config.softmax_type != "vanilla": d_softmax_offset = block.softmax_offset.grad - if backend in ["FlashAttention", "UnfusedDotProductAttention"]: + if backend in ["UnfusedDotProductAttention"]: if is_training: return out, max_logit, (q.grad, k.grad, v.grad, d_softmax_offset) else: return out, max_logit, (None, None, None, d_softmax_offset) - if backend == "FusedAttention": + if backend in ["FusedAttention", "FlashAttention"]: if qkv_format == "thd" and pad_between_seqs: out_orig = torch.Tensor([]).to(device="cuda", dtype=dtype) if is_training: diff --git a/tests/pytorch/attention/test_attention_with_cp.py b/tests/pytorch/attention/test_attention_with_cp.py index f0d2c27c1..a03f51f6c 100644 --- a/tests/pytorch/attention/test_attention_with_cp.py +++ b/tests/pytorch/attention/test_attention_with_cp.py @@ -306,10 +306,19 @@ def _submit(pool: PoolWorker, **kwargs) -> None: @pytest.mark.parametrize("model", model_configs_flash_attn.keys()) @pytest.mark.parametrize("qkv_format", qkv_formats) @pytest.mark.parametrize("cp_comm_type", cp_comm_types) -def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type): +@pytest.mark.parametrize("pad_between_seqs", [False, True]) +def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type, pad_between_seqs): num_gpus = 4 if cp_comm_type == "a2a+p2p" else 2 pool = cp_pool(num_gpus) + if pad_between_seqs: + if qkv_format != "thd": + pytest.skip("pad_between_seqs only applies to THD format!") + if not FlashAttentionUtils.v3_is_installed or get_device_compute_capability() > (9, 0): + pytest.skip("pad_between_seqs with CP requires Flash Attention v3 on Hopper (sm90)!") + if cp_comm_type == "a2a+p2p": + pytest.skip("pad_between_seqs is not yet supported with A2A+P2P CP comm type!") + config = model_configs_flash_attn[model] config.context_parallel = True config.cp_comm_type = cp_comm_type @@ -361,6 +370,7 @@ def test_cp_with_flash_attention(cp_pool, dtype, model, qkv_format, cp_comm_type qkv_format=qkv_format, kernel_backend="FlashAttention", cp_comm_type=cp_comm_type, + fa_pad_between_seqs=pad_between_seqs, log_level=pytest_logging_level, ) @@ -606,6 +616,7 @@ def test_cp_with_fused_attention( is_training=is_training, deterministic=_deterministic, ) + _, fused_attn_supported, _ = available_backends if fused_attn_supported and config.attn_mask_type in ["causal", "padding_causal"]: config_copy = copy.deepcopy(config) @@ -628,6 +639,23 @@ def test_cp_with_fused_attention( pytest.skip("Deterministic mode does not support non-vanilla softmax with FusedAttention") if _deterministic and config.attn_bias_type == "post_scale_bias" and is_training: pytest.skip("Deterministic mode does not support post_scale_bias with requires_grad") + # Observed: cuDNN det THD backward asks for ~128 * bHSS bytes of workspace + # on sm90; at 1<<30 that's 128 GiB, won't fit on H100's 80 GB. Held exactly + # at b=2 + power-of-2 S in our sweep; for b>=3 the workspace was observed to + # grow super-linearly (b=4 took ~4x the b=2 amount, not 2x) — revisit if a + # config uses b>2. + SM90_DET_FUSED_THD_BWD_MAX_BHSS = 1 << 30 + if ( + _deterministic + and qkv_format == "thd" + and get_device_compute_capability() == (9, 0) + and config.batch_size * config.num_heads * config.max_seqlen_q * config.max_seqlen_kv + >= SM90_DET_FUSED_THD_BWD_MAX_BHSS + ): + pytest.skip( + "Deterministic FusedAttention backward with THD format OOMs on sm90" + " for large bHSS configs (known cuDNN issue)." + ) _submit( pool, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 6e097265f..6c6adc6e3 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -822,10 +822,13 @@ def forward( fp8: bool = False, fp8_meta: Optional[Dict[str, Any]] = None, quantizers=None, + pad_between_seqs: Optional[bool] = False, inference_params: Optional[InferenceParams] = None, flash_attention_backend: Optional[PkgVersion] = PkgVersion("0"), fp8_output: bool = False, num_splits: Optional[int] = 1, + cu_seqlens_q_padded: Optional[torch.Tensor] = None, + cu_seqlens_kv_padded: Optional[torch.Tensor] = None, ) -> torch.Tensor: """flash-attn fprop""" @@ -1024,8 +1027,16 @@ def forward( cu_seqlens_kv, max_seqlen_q, max_seqlen_kv, - cu_seqlens_q if qkv_format == "thd" else None, - cu_seqlens_kv if qkv_format == "thd" else None, + ( + cu_seqlens_q_padded + if pad_between_seqs + else (cu_seqlens_q if qkv_format == "thd" else None) + ), + ( + cu_seqlens_kv_padded + if pad_between_seqs + else (cu_seqlens_kv if qkv_format == "thd" else None) + ), self.attention_dropout if self.training else 0.0, cp_group, cp_global_ranks, @@ -1037,7 +1048,7 @@ def forward( deterministic=self.deterministic, window_size=window_size, quantizers=quantizers, - pad_between_seqs=False, + pad_between_seqs=pad_between_seqs, use_flash_attn_3=use_flash_attn_3, fp8_output=fp8_output, ) @@ -1082,8 +1093,12 @@ def forward( else: func = flash_attn_with_kvcache_v3 # pylint: disable=possibly-used-before-assignment if not use_flash_attn_4 and (not use_flash_attn_3 or inference_params is None): - fa_optional_forward_args_thd.append(cu_seqlens_q) - fa_optional_forward_args_thd.append(cu_seqlens_kv) + fa_optional_forward_args_thd.append( + cu_seqlens_q_padded if pad_between_seqs else cu_seqlens_q + ) + fa_optional_forward_args_thd.append( + cu_seqlens_kv_padded if pad_between_seqs else cu_seqlens_kv + ) fa_optional_forward_args_thd.append(max_seqlen_q) fa_optional_forward_args_thd.append(max_seqlen_kv) if use_flash_attn_4: @@ -1139,6 +1154,13 @@ def forward( fa_3_optional_forward_kwargs = {} fa_3_optional_forward_kwargs["window_size"] = window_size fa_3_optional_forward_kwargs["num_splits"] = num_splits + if pad_between_seqs: + fa_3_optional_forward_kwargs["seqused_q"] = ( + cu_seqlens_q[1:] - cu_seqlens_q[:-1] + ) + fa_3_optional_forward_kwargs["seqused_k"] = ( + cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] + ) if inference_params is None: fa_3_optional_forward_kwargs["deterministic"] = self.deterministic else: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py index 35684625a..36847e40e 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/context_parallel.py @@ -663,6 +663,8 @@ def get_fa_args( dq=None, dk=None, dv=None, + seqused_q=None, + seqused_k=None, ): """Get forward/backward arguments for flash-attn v2 and v3.""" if use_flash_attn_3: @@ -672,7 +674,9 @@ def get_fa_args( *[None] * 4, # k_new, v_new, qv, out cu_seqlens_q, cu_seqlens_kv, - *[None] * 3, # cu_seqlens_k_new, seqused_q, seqused_k + None, # cu_seqlens_k_new + seqused_q, + seqused_k, max_seqlen_q, max_seqlen_kv, *[None] @@ -690,8 +694,8 @@ def get_fa_args( return [ cu_seqlens_q, cu_seqlens_kv, - None, # sequed_q - None, # sequed_k + seqused_q, + seqused_k, max_seqlen_q, max_seqlen_kv, dq, @@ -701,8 +705,8 @@ def get_fa_args( return [ None, # cu_seqlens_q None, # cu_seqlens_kv - None, # sequed_q - None, # sequed_k + None, # seqused_q + None, # seqused_k max_seqlen_q, max_seqlen_kv, dq, @@ -1020,6 +1024,9 @@ def cp_p2p_fwd_flash_attn( flash_attn_fwd, max_seqlen_q, max_seqlen_kv, + pad_between_seqs, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, q_part, k_part, v_part, @@ -1046,6 +1053,20 @@ def cp_p2p_fwd_flash_attn( fa_forward_kwargs["window_size_left"] = -1 fa_forward_kwargs["window_size_right"] = -1 + seqused_q = None + seqused_k = None + if pad_between_seqs and use_flash_attn_3 and qkv_format == "thd": + # Derive actual token counts per batch element from cu_seqlens + seqused_q = cu_seqlens_q_per_step[1:] - cu_seqlens_q_per_step[:-1] + seqused_k = cu_seqlens_kv_per_step[1:] - cu_seqlens_kv_per_step[:-1] + # Override cu_seqlens to padded layout for tensor memory layout + cu_seqlens_q_ = cu_seqlens_q_padded + cu_seqlens_kv_ = cu_seqlens_kv_padded + if section == "lower-triangle": + cu_seqlens_kv_ = cu_seqlens_kv_padded // 2 + elif section == "upper-triangle": + cu_seqlens_q_ = cu_seqlens_q_padded // 2 + fa_forward_args_thd = get_fa_args( True, use_flash_attn_3, @@ -1054,6 +1075,8 @@ def cp_p2p_fwd_flash_attn( cu_seqlens_kv=cu_seqlens_kv_, max_seqlen_q=max_seqlen_q_, max_seqlen_kv=max_seqlen_kv_, + seqused_q=seqused_q, + seqused_k=seqused_k, ) fa_outputs = flash_attn_fwd( q_part, @@ -1296,6 +1319,9 @@ def cp_p2p_bwd_flash_attn( rng_states, softmax_lse, softmax_lse_, + pad_between_seqs, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, q_part, k_part, v_part, @@ -1304,7 +1330,10 @@ def cp_p2p_bwd_flash_attn( section, ): """Per-tile backward call of CP P2P with FlashAttention backend""" - dq, dk, dv = [torch.empty_like(x) for x in [q_part, k_part, v_part]] + if pad_between_seqs: + dq, dk, dv = [torch.zeros_like(x) for x in [q_part, k_part, v_part]] + else: + dq, dk, dv = [torch.empty_like(x) for x in [q_part, k_part, v_part]] if fa_utils.v2_3_plus and not fa_utils.v2_7_0_plus: fa_backward_kwargs["window_size"] = (-1, -1) elif use_flash_attn_3 or fa_utils.v2_7_0_plus: @@ -1329,17 +1358,33 @@ def cp_p2p_bwd_flash_attn( max_seqlen_q_ = max_seqlen_q // 2 softmax_lse__ = softmax_lse_ + seqused_q = None + seqused_k = None + cu_seqlens_q_bwd = cu_seqlens_q_per_step[cp_size - step - 1] + cu_seqlens_kv_bwd = cu_seqlens_kv_per_step[cp_size - step - 1] + if pad_between_seqs and use_flash_attn_3 and qkv_format == "thd": + seqused_q = cu_seqlens_q_bwd[1:] - cu_seqlens_q_bwd[:-1] + seqused_k = cu_seqlens_kv_bwd[1:] - cu_seqlens_kv_bwd[:-1] + cu_seqlens_q_bwd = cu_seqlens_q_padded + cu_seqlens_kv_bwd = cu_seqlens_kv_padded + if section == "lower-triangle": + cu_seqlens_kv_bwd = cu_seqlens_kv_padded // 2 + elif section == "upper-triangle": + cu_seqlens_q_bwd = cu_seqlens_q_padded // 2 + fa_backward_args_thd = get_fa_args( False, use_flash_attn_3, qkv_format, - cu_seqlens_q=cu_seqlens_q_per_step[cp_size - step - 1], - cu_seqlens_kv=cu_seqlens_kv_per_step[cp_size - step - 1], + cu_seqlens_q=cu_seqlens_q_bwd, + cu_seqlens_kv=cu_seqlens_kv_bwd, max_seqlen_q=max_seqlen_q_, max_seqlen_kv=max_seqlen_kv_, dq=dq, dk=dk, dv=dv, + seqused_q=seqused_q, + seqused_k=seqused_k, ) if use_flash_attn_3: fa_backward_kwargs["is_causal"] = causal_ @@ -1779,6 +1824,9 @@ def forward( flash_attn_fwd, max_seqlen_q, max_seqlen_kv, + pad_between_seqs, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, ] # cp_size = 4: @@ -1821,7 +1869,9 @@ def forward( else: out_per_step[i], softmax_lse_per_step[i], rng_states[i] = ( cp_p2p_fwd_flash_attn( - *flash_attn_inputs, *prepare_outputs, section + *flash_attn_inputs, + *prepare_outputs, + section, ) ) elif i <= rank: @@ -1848,7 +1898,9 @@ def forward( else: out_per_step[i], softmax_lse_per_step[i], rng_states[i] = ( cp_p2p_fwd_flash_attn( - *flash_attn_inputs, *prepare_outputs, section + *flash_attn_inputs, + *prepare_outputs, + section, ) ) else: @@ -1875,7 +1927,9 @@ def forward( else: out_per_step[i], softmax_lse_per_step[i], rng_states[i] = ( cp_p2p_fwd_flash_attn( - *flash_attn_inputs, *prepare_outputs, section + *flash_attn_inputs, + *prepare_outputs, + section, ) ) else: @@ -1900,7 +1954,11 @@ def forward( ) = cp_p2p_fwd_fused_attn(*fused_attn_inputs, *prepare_outputs, section) else: out_per_step[i], softmax_lse_per_step[i], rng_states[i] = ( - cp_p2p_fwd_flash_attn(*flash_attn_inputs, *prepare_outputs, section) + cp_p2p_fwd_flash_attn( + *flash_attn_inputs, + *prepare_outputs, + section, + ) ) # softmax_lse correction @@ -2150,6 +2208,7 @@ def forward( ctx.attn_bias_shape = None if attn_bias is None else attn_bias.shape ctx.deterministic = deterministic ctx.use_fused_attention = use_fused_attention + ctx.pad_between_seqs = pad_between_seqs ctx.softmax_lse_in_packed_format = softmax_lse_in_packed_format ctx.second_half_lse_seqlen = second_half_lse_seqlen ctx.fp8_meta = fp8_meta @@ -2560,6 +2619,9 @@ def backward(ctx, dout, *_args): rng_states, softmax_lse, softmax_lse_, + ctx.pad_between_seqs, + cu_seqlens_q_padded, + cu_seqlens_kv_padded, ] # Reverse the steps in forward. In the cp_size x cp_size (i.e. GPU x step) matrix, @@ -2575,7 +2637,9 @@ def backward(ctx, dout, *_args): ) else: dq_, dk_, dv_ = cp_p2p_bwd_flash_attn( - *flash_attn_inputs, *prepare_outputs, section + *flash_attn_inputs, + *prepare_outputs, + section, ) elif i >= (cp_size - rank - 1): section = "lower-triangle" @@ -2586,7 +2650,9 @@ def backward(ctx, dout, *_args): ) else: dq_, dk_, dv_ = cp_p2p_bwd_flash_attn( - *flash_attn_inputs, *prepare_outputs, section + *flash_attn_inputs, + *prepare_outputs, + section, ) else: section = "upper-triangle" @@ -2597,7 +2663,9 @@ def backward(ctx, dout, *_args): ) else: dq_, dk_, dv_ = cp_p2p_bwd_flash_attn( - *flash_attn_inputs, *prepare_outputs, section + *flash_attn_inputs, + *prepare_outputs, + section, ) else: section = "all" @@ -2608,7 +2676,9 @@ def backward(ctx, dout, *_args): ) else: dq_, dk_, dv_ = cp_p2p_bwd_flash_attn( - *flash_attn_inputs, *prepare_outputs, section + *flash_attn_inputs, + *prepare_outputs, + section, ) # dq, dk, dv are reduced across steps in higher precision @@ -3838,6 +3908,7 @@ def forward( cp_group, cp_stream, quantizers, + pad_between_seqs, use_flash_attn_3, softmax_type, softmax_offset, @@ -4073,14 +4144,25 @@ def forward( out_f16 = out_.dequantize(dtype=fwd_nominal_dtype) out_part = out_f16 else: + seqused_q = None + seqused_k = None + fa_cu_seqlens_q = cu_seqlens_q + fa_cu_seqlens_kv = cu_seqlens_kv + if pad_between_seqs and use_flash_attn_3 and qkv_format == "thd": + seqused_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1] + seqused_k = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] + fa_cu_seqlens_q = cu_seqlens_q_padded + fa_cu_seqlens_kv = cu_seqlens_kv_padded fa_forward_args_thd = get_fa_args( True, use_flash_attn_3, qkv_format, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_kv=cu_seqlens_kv, + cu_seqlens_q=fa_cu_seqlens_q, + cu_seqlens_kv=fa_cu_seqlens_kv, max_seqlen_q=max_seqlen_q, max_seqlen_kv=max_seqlen_kv, + seqused_q=seqused_q, + seqused_k=seqused_k, ) fa_outputs = flash_attn_fwd( q_part, @@ -4217,6 +4299,7 @@ def forward( ctx.fwd_nominal_dtype = fwd_nominal_dtype ctx.fp8_recipe = fp8_recipe ctx.use_flash_attn_3 = use_flash_attn_3 + ctx.pad_between_seqs = pad_between_seqs ctx.softmax_type = softmax_type ctx.dQKV_quantizer = dQKV_quantizer @@ -4405,18 +4488,32 @@ def backward(ctx, dout, *_args): dq, dk, dv = [x._data for x in [dq, dk, dv]] else: softmax_lse, rng_state = aux_ctx_tensors - dq, dk, dv = [torch.empty_like(x) for x in [q, k, v]] + if ctx.pad_between_seqs: + dq, dk, dv = [torch.zeros_like(x) for x in [q, k, v]] + else: + dq, dk, dv = [torch.empty_like(x) for x in [q, k, v]] + seqused_q = None + seqused_k = None + fa_cu_seqlens_q = cu_seqlens_q + fa_cu_seqlens_kv = cu_seqlens_kv + if ctx.pad_between_seqs and ctx.use_flash_attn_3 and ctx.dqkv_format == "thd": + seqused_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1] + seqused_k = cu_seqlens_kv[1:] - cu_seqlens_kv[:-1] + fa_cu_seqlens_q = cu_seqlens_q_padded + fa_cu_seqlens_kv = cu_seqlens_kv_padded fa_backward_args_thd = get_fa_args( False, ctx.use_flash_attn_3, ctx.dqkv_format, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_kv=cu_seqlens_kv, + cu_seqlens_q=fa_cu_seqlens_q, + cu_seqlens_kv=fa_cu_seqlens_kv, max_seqlen_q=ctx.max_seqlen_q, max_seqlen_kv=ctx.max_seqlen_kv, dq=dq, dk=dk, dv=dv, + seqused_q=seqused_q, + seqused_k=seqused_k, ) if not ctx.use_flash_attn_3: fa_backward_kwargs["rng_state"] = rng_state @@ -4524,6 +4621,7 @@ def backward(ctx, dout, *_args): None, None, None, + None, d_softmax_offset, None, ) @@ -4740,6 +4838,7 @@ def attn_forward_func_with_cp( cp_group, cp_stream, quantizers, + pad_between_seqs, use_flash_attn_3, softmax_type, softmax_offset, diff --git a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py index b38b66c3e..ca848a948 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py @@ -1658,10 +1658,13 @@ def forward( fp8=self.fp8 and self.fp8_meta["recipe"].fp8_dpa, fp8_meta=self.fp8_meta, quantizers=self.quantizers, + pad_between_seqs=pad_between_seqs, inference_params=inference_params, flash_attention_backend=flash_attention_backend, fp8_output=fp8_output, num_splits=num_splits, + cu_seqlens_q_padded=cu_seqlens_q_padded, + cu_seqlens_kv_padded=cu_seqlens_kv_padded, ) if use_fused_attention: diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 1f1637cec..6565e9f6f 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -651,7 +651,7 @@ def get_attention_backend( # backend | precision | KV cache | architecture | qkv_format | page_size # --------------------------------------------------------------------------------------- # Fused | FP16/BF16 | non-paged/paged | sm80+ | bshd,sbhd,thd | >= 1 - # Flash v2 | FP16/BF16 | non-paged/paged | sm80+ | bshd,sbhd,thd | >= 256 + # Flash v2 | FP16/BF16 | non-paged/paged | sm80+ | bshd,sbhd,thd | % 256 == 0 # Flash v3 | FP16/BF16 | non-paged/paged | sm90 | bshd,sbhd,thd | >= 1 # | FP8 | non-paged/paged | sm90 | thd | >= 1 # Flash v4 | FP16/BF16 | TODO | sm80+ | bshd,sbhd,thd | TODO @@ -691,9 +691,9 @@ def get_attention_backend( use_fused_attention = False use_unfused_attention = False if inference_params.is_paged: - if use_flash_attention_2 and inference_params.page_size < 256: + if use_flash_attention_2 and inference_params.page_size % 256 != 0: if FlashAttentionUtils.is_installed: - logger.debug("Disabling FlashAttention 2 for page size < 256") + logger.debug("Disabling FlashAttention 2 for page size not divisible by 256") use_flash_attention_2 = False if use_flash_attention_2: if not FlashAttentionUtils.is_installed: @@ -703,6 +703,16 @@ def get_attention_backend( "Disabling FlashAttention 2 as paged attention requires flash-attn 2.5+" ) use_flash_attention_2 = False + else: + # Non-paged KV cache still passes a block_table to FA2 for thd_2bshd support, + # and FA2 enforces page_size % 256 == 0 on the effective page size (max_seqlen_kv). + if use_flash_attention_2 and max_seqlen_kv % 256 != 0: + if FlashAttentionUtils.is_installed: + logger.debug( + "Disabling FlashAttention 2 for non-paged KV cache" + " with max_seqlen_kv not divisible by 256" + ) + use_flash_attention_2 = False if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: logger.debug("Disabling FlashAttention 4 as it does not support KV cache.") use_flash_attention_4 = False @@ -844,15 +854,18 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt if qkv_format == "thd": if pad_between_seqs: if ( # pylint: disable=too-many-boolean-expressions - (use_flash_attention_2 and FlashAttentionUtils.is_installed) - or (use_flash_attention_3 and FlashAttentionUtils.v3_is_installed) - or (use_flash_attention_4 and FlashAttentionUtils.v4_is_installed) - ): + use_flash_attention_2 and FlashAttentionUtils.is_installed + ) or (use_flash_attention_4 and FlashAttentionUtils.v4_is_installed): logger.debug( - "Disabling FlashAttention for qkv_format = thd when there is " + "Disabling FlashAttention 2 and 4 for qkv_format = thd when there is " "padding between sequences, i.e. [a, a, PAD, b, b, b, PAD, c, PAD]" ) - use_flash_attention = False + use_flash_attention_2 = False + use_flash_attention_4 = False + # FA3 supports pad_between_seqs via seqused_q/seqused_k + if use_unfused_attention: + logger.debug("Disabling UnfusedDotProductAttention for pad_between_seqs = True") + use_unfused_attention = False if device_compute_capability == (12, 0): if cudnn_version < (9, 18, 1): if use_fused_attention: @@ -1273,9 +1286,12 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ) use_flash_attention_2 = False if use_flash_attention_3 and deterministic and FlashAttentionUtils.v3_is_installed: - if head_dim_qk >= 256: + if is_training and max(head_dim_qk, head_dim_v) >= 256: logger.debug( - "Disabling FlashAttention 3 for deterministic execution with head_dim_qk >= 256." + "Disabling FlashAttention 3 for deterministic backward with" + " max(head_dim_qk, head_dim_v) >= 256. Found: head_dim_qk = %s, head_dim_v = %s.", + head_dim_qk, + head_dim_v, ) use_flash_attention_3 = False if use_fused_attention and deterministic: From 7e6ffccca7d19087620b7019e8a50944c5aaaa5f Mon Sep 17 00:00:00 2001 From: hx Date: Wed, 27 May 2026 00:22:21 +0800 Subject: [PATCH 081/180] [Common/PyTorch/JAX] make offset of ClampedSwiGLU configurable (#2938) * swiglu offset Signed-off-by: Hongxiao Bai * fix fusion pattern check Signed-off-by: Hongxiao Bai * use swiglu_v2 Signed-off-by: Hongxiao Bai * add default value to v1 Signed-off-by: Hongxiao Bai * fix test Signed-off-by: Hongxiao Bai * add default value to jax version Signed-off-by: Hongxiao Bai * revert the default value change Signed-off-by: Hongxiao Bai * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * update the fusion path Signed-off-by: Hongxiao Bai * update cudnn-frontend to 1.24.0 Signed-off-by: Hongxiao Bai * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Hongxiao Bai Signed-off-by: vthumbe1503 Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> Co-authored-by: vthumbe1503 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Tim Moon <4406448+timmoon10@users.noreply.github.com> --- tests/jax/test_custom_call_compute.py | 2 +- tests/pytorch/test_fusible_ops.py | 48 +++++++++++++--- .../common/activation/swiglu.cu | 24 +++++++- .../common/cast/fp8/gated_fp8.cuh | 5 +- .../common/cast/mxfp8/gated_mxfp8.cuh | 10 ++-- .../include/transformer_engine/activation.h | 55 +++++++++++++++++++ transformer_engine/common/util/math.h | 3 +- .../common/util/vectorized_pointwise.h | 8 +-- .../jax/cpp_extensions/activation.py | 15 +++-- transformer_engine/jax/csrc/extensions.h | 4 +- .../jax/csrc/extensions/activation.cpp | 10 ++-- transformer_engine/pytorch/csrc/extensions.h | 5 +- .../pytorch/csrc/extensions/activation.cpp | 11 ++-- .../pytorch/csrc/extensions/pybind.cpp | 5 +- .../pytorch/module/layernorm_mlp.py | 15 +++-- transformer_engine/pytorch/ops/_common.py | 21 ++++++- .../pytorch/ops/basic/swiglu.py | 16 +++++- .../pytorch/ops/fused/backward_grouped_mlp.py | 19 +++++++ .../pytorch/ops/fused/forward_grouped_mlp.py | 20 +++++++ transformer_engine/pytorch/transformer.py | 5 +- 20 files changed, 240 insertions(+), 61 deletions(-) diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py index 14d28d95b..0ed5645eb 100644 --- a/tests/jax/test_custom_call_compute.py +++ b/tests/jax/test_custom_call_compute.py @@ -245,7 +245,7 @@ def test_act_grad(self, shape, activation_type): value_and_grad(self.primitive_func, (0,)), static_argnums=(1, 3) ) act_args = ( - {"limit": 0.75, "alpha": 1.702} + {"limit": 0.75, "alpha": 1.702, "glu_linear_offset": 0.5} if activation_type == ("clamped_silu", "clamped_linear") else {} ) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 8e63caa98..1ced32e1a 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -1846,6 +1846,7 @@ def test_interleaved_swiglu(self): @pytest.mark.parametrize("quantization", _quantization_list) @pytest.mark.parametrize("quantize_forward", (False, True)) @pytest.mark.parametrize("quantize_backward", (False, True)) + @pytest.mark.parametrize("glu_linear_offset", (1.0, 0.0)) def test_clamped_swiglu( self, *, @@ -1856,6 +1857,7 @@ def test_clamped_swiglu( quantization: Optional[str], quantize_forward: bool, quantize_backward: bool, + glu_linear_offset: float, limit: float = 0.75, alpha: float = 1.702, ): @@ -1898,7 +1900,7 @@ def test_clamped_swiglu( x_glu = x_glu.clamp(min=None, max=limit) x_linear = x_linear.clamp(min=-limit, max=limit) out_glu = x_glu * torch.sigmoid(alpha * x_glu) - y_ref = out_glu * (x_linear + 1) + y_ref = out_glu * (x_linear + glu_linear_offset) y_ref.backward(dy_ref) # Implementation with fusible operation @@ -1909,6 +1911,7 @@ def test_clamped_swiglu( te_ops.ClampedSwiGLU( limit=limit, alpha=alpha, + glu_linear_offset=glu_linear_offset, glu_interleave_size=glu_interleave_size, ), te_ops.Quantize(forward=quantize_forward, backward=False), @@ -1938,6 +1941,7 @@ def test_interleaved_clamped_swiglu(self): quantize_forward=False, quantize_backward=False, glu_interleave_size=32, + glu_linear_offset=1.0, ) @pytest.mark.parametrize("scale", (1, 0, -2.5, 3.5)) @@ -2594,6 +2598,7 @@ def test_scaled_activation_recompute_in_mlp_config(self, op_cls) -> None: @pytest.mark.parametrize("in_shape", ((71, 192), (5, 7, 128))) @pytest.mark.parametrize("input_requires_grad", (False, True)) @pytest.mark.parametrize("scales_requires_grad", (False, True)) + @pytest.mark.parametrize("glu_linear_offset", (1.0, 0.0)) def test_scaled_clamped_qgeglu( self, *, @@ -2603,6 +2608,7 @@ def test_scaled_clamped_qgeglu( device: torch.device = "cuda", input_requires_grad: bool, scales_requires_grad: bool, + glu_linear_offset: float, limit: float = 7.0, alpha: float = 1.702, ) -> None: @@ -2647,7 +2653,7 @@ def test_scaled_clamped_qgeglu( x_glu = x_glu.clamp(min=None, max=limit) x_linear = x_linear.clamp(min=-limit, max=limit) out_glu = x_glu * torch.sigmoid(alpha * x_glu) - y = out_glu * (x_linear + 1) + y = out_glu * (x_linear + glu_linear_offset) y_ref = scales_ref.unsqueeze(-1) * y if input_requires_grad or scales_requires_grad: y_ref.backward(dy_ref) @@ -2656,6 +2662,7 @@ def test_scaled_clamped_qgeglu( glu_interleave_size=glu_interleave_size, limit=limit, alpha=alpha, + glu_linear_offset=glu_linear_offset, ) y_test = op(x_test, scales_test) if input_requires_grad or scales_requires_grad: @@ -2674,6 +2681,7 @@ def test_interleaved_scaled_clamped_qgeglu(self): glu_interleave_size=32, input_requires_grad=True, scales_requires_grad=True, + glu_linear_offset=1.0, ) @@ -3685,7 +3693,13 @@ def test_layernorm_mlp( @pytest.mark.parametrize("delay_wgrad_compute", (False, True)) @pytest.mark.parametrize("hidden_size", (128, 256)) @pytest.mark.parametrize( - "activation", ("scaled_swiglu", "scaled_clamped_qgeglu", "scaled_srelu") + "activation", + ( + "scaled_swiglu", + "scaled_clamped_qgeglu", + "scaled_clamped_qgeglu_custom", + "scaled_srelu", + ), ) def test_grouped_mlp( self, @@ -3719,7 +3733,7 @@ def test_grouped_mlp( with_quantization = quantization is not None if activation == "scaled_swiglu": scaled_act = te_ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) - elif activation == "scaled_clamped_qgeglu": + elif activation.startswith("scaled_clamped_qgeglu"): scaled_act = te_ops.ScaledClampedQGeGLU(glu_interleave_size=glu_interleave_size) elif activation == "scaled_srelu": scaled_act = te_ops.ScaledSReLU() @@ -3742,13 +3756,23 @@ def test_grouped_mlp( if ( with_quantization and quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6") - and activation == "scaled_clamped_qgeglu" + and activation.startswith("scaled_clamped_qgeglu") and bias ): # TODO: ksivaman: Need to debug numerics for this case. pytest.skip("Bias/dbias not yet supported in NVFP4 fused grouped MLP with GeGLU") fc1_out_features = 2 * hidden_size if activation_is_glu else hidden_size + # Activation parameters for clamped QGeGLU variants + if activation == "scaled_clamped_qgeglu_custom": + geglu_limit = 5.0 + geglu_alpha = 1.5 + geglu_offset = 0.5 + else: + geglu_limit = 7.0 + geglu_alpha = 1.702 + geglu_offset = 1.0 + # Random data x_ref, x_test = make_reference_and_test_tensors( in_shape, @@ -3840,13 +3864,12 @@ def test_grouped_mlp( if activation == "scaled_swiglu": x1, x2 = x.chunk(2, dim=-1) x = torch.nn.functional.silu(x1) * x2 - elif activation == "scaled_clamped_qgeglu": + elif activation.startswith("scaled_clamped_qgeglu"): x1, x2 = x.chunk(2, dim=-1) - lim = torch.tensor(7.0, device=x1.device, dtype=x1.dtype) - geglu_alpha = 1.702 + lim = torch.tensor(geglu_limit, device=x1.device, dtype=x1.dtype) x1c = torch.minimum(x1, lim) x2c = torch.clamp(x2, -lim, lim) - x = (x2c + 1) * (x1c * torch.sigmoid(geglu_alpha * x1c)) + x = (x2c + geglu_offset) * (x1c * torch.sigmoid(geglu_alpha * x1c)) elif activation == "scaled_srelu": x = torch.nn.functional.relu(x).square() else: @@ -3861,6 +3884,13 @@ def test_grouped_mlp( # Construct operations recipe = make_recipe(quantization) + if activation == "scaled_clamped_qgeglu_custom": + scaled_act = te_ops.ScaledClampedQGeGLU( + glu_interleave_size=glu_interleave_size, + limit=geglu_limit, + alpha=geglu_alpha, + glu_linear_offset=geglu_offset, + ) with te.quantized_model_init(enabled=with_quantization, recipe=recipe): fc1 = te_ops.GroupedLinear( group_size, diff --git a/transformer_engine/common/activation/swiglu.cu b/transformer_engine/common/activation/swiglu.cu index 0b5b6069b..7120a7eb6 100644 --- a/transformer_engine/common/activation/swiglu.cu +++ b/transformer_engine/common/activation/swiglu.cu @@ -39,7 +39,16 @@ void nvte_clamped_swiglu(const NVTETensor input, NVTETensor output, float limit, cudaStream_t stream) { NVTE_API_CALL(nvte_clamped_swiglu); using namespace transformer_engine; - ClampedSwiGLUParam param = {limit, alpha}; + // Preserve original behavior: linear (gate) component offset is hard-coded to 1.0f. + ClampedSwiGLUParam param = {limit, alpha, /*glu_linear_offset=*/1.0f}; + gated_act_fn>(input, output, param, stream); +} + +void nvte_clamped_swiglu_v2(const NVTETensor input, NVTETensor output, float limit, float alpha, + float glu_linear_offset, cudaStream_t stream) { + NVTE_API_CALL(nvte_clamped_swiglu_v2); + using namespace transformer_engine; + ClampedSwiGLUParam param = {limit, alpha, glu_linear_offset}; gated_act_fn>(input, output, param, stream); } @@ -47,7 +56,18 @@ void nvte_clamped_dswiglu(const NVTETensor grad, const NVTETensor input, NVTETen float limit, float alpha, cudaStream_t stream) { NVTE_API_CALL(nvte_clamped_dswiglu); using namespace transformer_engine; - ClampedSwiGLUParam param = {limit, alpha}; + // Preserve original behavior: linear (gate) component offset is hard-coded to 1.0f. + ClampedSwiGLUParam param = {limit, alpha, /*glu_linear_offset=*/1.0f}; + dgated_act_fn, clamped_dsilu>( + grad, input, output, param, stream); +} + +void nvte_clamped_dswiglu_v2(const NVTETensor grad, const NVTETensor input, NVTETensor output, + float limit, float alpha, float glu_linear_offset, + cudaStream_t stream) { + NVTE_API_CALL(nvte_clamped_dswiglu_v2); + using namespace transformer_engine; + ClampedSwiGLUParam param = {limit, alpha, glu_linear_offset}; dgated_act_fn, clamped_dsilu>( grad, input, output, param, stream); } diff --git a/transformer_engine/common/cast/fp8/gated_fp8.cuh b/transformer_engine/common/cast/fp8/gated_fp8.cuh index 6123d7130..522a9add8 100644 --- a/transformer_engine/common/cast/fp8/gated_fp8.cuh +++ b/transformer_engine/common/cast/fp8/gated_fp8.cuh @@ -169,9 +169,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) float gate_elt = static_cast(in_gate_sh_curr[shmem_idx]); bool dgate_elt = true; // gating is ideally an identity function if constexpr (std::is_same::value) { - // In case of GPT OSS, clamp the activation and gate values - dgate_elt = gate_elt <= p.limit && gate_elt >= -p.limit; // Derivative of clamp - gate_elt = min(max(-p.limit, gate_elt), p.limit) + 1; + dgate_elt = gate_elt <= p.limit && gate_elt >= -p.limit; + gate_elt = min(max(-p.limit, gate_elt), p.limit) + p.glu_linear_offset; } if constexpr (IS_BWD) { diff --git a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh index 49169a4e1..83b5a49ca 100644 --- a/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/gated_mxfp8.cuh @@ -245,9 +245,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) float after_gate_elt; bool dgate_elt = true; // gating is ideally an identity function if constexpr (std::is_same::value) { - // In case of GPT OSS, clamp the activation and gate values - dgate_elt = gate_elt <= p.limit && gate_elt >= -p.limit; // Derivative of clamp - gate_elt = min(max(-p.limit, gate_elt), p.limit) + 1.0f; + dgate_elt = gate_elt <= p.limit && gate_elt >= -p.limit; + gate_elt = min(max(-p.limit, gate_elt), p.limit) + p.glu_linear_offset; } if constexpr (IS_BWD) { float grad_elt = static_cast(in_grad_sh[shmem_offset_colwise]); @@ -510,9 +509,8 @@ __global__ void __launch_bounds__(THREADS_PER_CHUNK) float after_gate_elt; bool dgate_elt = true; if constexpr (std::is_same::value) { - // In case of GPT OSS, clamp the activation and gate values - dgate_elt = gate_elt <= p.limit && gate_elt >= -p.limit; // Derivative of clamp - gate_elt = min(max(-p.limit, gate_elt), p.limit) + 1.0f; + dgate_elt = gate_elt <= p.limit && gate_elt >= -p.limit; + gate_elt = min(max(-p.limit, gate_elt), p.limit) + p.glu_linear_offset; } if constexpr (IS_BWD) { float grad_elt = static_cast(in_grad.data.elt[e]); diff --git a/transformer_engine/common/include/transformer_engine/activation.h b/transformer_engine/common/include/transformer_engine/activation.h index 854f52c20..4ed083740 100644 --- a/transformer_engine/common/include/transformer_engine/activation.h +++ b/transformer_engine/common/include/transformer_engine/activation.h @@ -322,6 +322,11 @@ void nvte_geglu(const NVTETensor input, NVTETensor output, cudaStream_t stream); void nvte_swiglu(const NVTETensor input, NVTETensor output, cudaStream_t stream); /*! \brief Computes the gated Swish activation of the input used in GPT OSS. + * + * \deprecated This function has been deprecated in favor of nvte_clamped_swiglu_v2, + * which exposes a configurable offset for the linear (gate) component. + * This API is preserved for backward compatibility and is equivalent to + * calling nvte_clamped_swiglu_v2 with glu_linear_offset = 1.0. * * See https://github.com/openai/gpt-oss/blob/a0a84273e9e0c14a233cb9befdfd159c2bcfa6cd/gpt_oss/torch/model.py#L250 * This Gated activation has two differences compared to the original SwiGLU @@ -341,6 +346,28 @@ void nvte_swiglu(const NVTETensor input, NVTETensor output, cudaStream_t stream) void nvte_clamped_swiglu(const NVTETensor input, NVTETensor output, float limit, float alpha, cudaStream_t stream); +/*! \brief Computes the gated Swish activation of the input used in GPT OSS, with a configurable + * offset for the linear (gate) component after clamping. + * + * See https://github.com/openai/gpt-oss/blob/a0a84273e9e0c14a233cb9befdfd159c2bcfa6cd/gpt_oss/torch/model.py#L250 + * This Gated activation has two differences compared to the original SwiGLU + * 1. Both gate and pre-activations are clipped based on parameter limit. + * 2. Activation uses sigmoid(alpha * x) instead of sigmoid(x) used in Swish activation inspired + * by original GELU paper https://arxiv.org/pdf/1606.08415 + * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, + * the block quantization (MXFP8) of the specified shape of the block will be used. + * + * \param[in] input Input tensor of shape [N, H * 2]. + * \param[in,out] output Output tensor of shape [N, H]. + * It computes Act(input[N, :H]) x (input[N, H:] + glu_linear_offset) + * \param[in] limit Clipping limits for gate and pre-activation. + * \param[in] alpha Scaling factor for the sigmoid function used in the activation. + * \param[in] glu_linear_offset Offset added to the linear component after clamping (typically 1.0). + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_clamped_swiglu_v2(const NVTETensor input, NVTETensor output, float limit, float alpha, + float glu_linear_offset, cudaStream_t stream); + /*! \brief Computes the gated ReLU activation of the input. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. @@ -399,6 +426,11 @@ void nvte_dswiglu(const NVTETensor grad, const NVTETensor input, NVTETensor outp cudaStream_t stream); /*! \brief Computes the gradient of gated Swish activation of the input used in GPT OSS. + * + * \deprecated This function has been deprecated in favor of nvte_clamped_dswiglu_v2, + * which exposes a configurable offset for the linear (gate) component. + * This API is preserved for backward compatibility and is equivalent to + * calling nvte_clamped_dswiglu_v2 with glu_linear_offset = 1.0. * * https://github.com/openai/gpt-oss/blob/a0a84273e9e0c14a233cb9befdfd159c2bcfa6cd/gpt_oss/torch/model.py#L250 * This activation has two differences compared to the original SwiGLU @@ -418,6 +450,29 @@ void nvte_dswiglu(const NVTETensor grad, const NVTETensor input, NVTETensor outp void nvte_clamped_dswiglu(const NVTETensor grad, const NVTETensor input, NVTETensor output, float limit, float alpha, cudaStream_t stream); +/*! \brief Computes the gradient of gated Swish activation of the input used in GPT OSS, with a + * configurable offset for the linear (gate) component after clamping. + * + * https://github.com/openai/gpt-oss/blob/a0a84273e9e0c14a233cb9befdfd159c2bcfa6cd/gpt_oss/torch/model.py#L250 + * This activation has two differences compared to the original SwiGLU + * 1. Both gate and pre-activations are clipped based on parameter limit. + * 2. Activation uses sigmoid(alpha * x) instead of sigmoid(x) used in Swish activation inspired + * by original GELU paper https://arxiv.org/pdf/1606.08415 + * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, + * the block quantization (MXFP8) of the specified shape of the block will be used. + * + * \param[in] grad Incoming gradient of shape [N, H]. + * \param[in] input Forward input tensor of shape [N, H * 2]. + * \param[in,out] output Outgoing gradient of shape [N, H * 2]. + * \param[in] limit Clipping limits for gate and pre-activation. + * \param[in] alpha Scaling factor for the sigmoid function used in the activation. + * \param[in] glu_linear_offset Offset added to the linear component after clamping (typically 1.0). + * \param[in] stream CUDA stream used for the operation. + */ +void nvte_clamped_dswiglu_v2(const NVTETensor grad, const NVTETensor input, NVTETensor output, + float limit, float alpha, float glu_linear_offset, + cudaStream_t stream); + /*! \brief Computes the gated ReLU activation gradient. * If the scaling mode of the output tensor is set to NVTE_MXFP8_1D_SCALING, * the block quantization (MXFP8) of the specified shape of the block will be used. diff --git a/transformer_engine/common/util/math.h b/transformer_engine/common/util/math.h index 05fe2f539..64b6fa2d4 100644 --- a/transformer_engine/common/util/math.h +++ b/transformer_engine/common/util/math.h @@ -13,7 +13,8 @@ struct Empty {}; struct ClampedSwiGLUParam { float limit; - float alpha = 1.702f; // Default value for QuickGELU + float alpha = 1.702f; // Default value for QuickGELU + float glu_linear_offset = 1.0f; // Offset added to the linear (gate) component after clamping }; template diff --git a/transformer_engine/common/util/vectorized_pointwise.h b/transformer_engine/common/util/vectorized_pointwise.h index 0aa2df7d2..7707c68a0 100644 --- a/transformer_engine/common/util/vectorized_pointwise.h +++ b/transformer_engine/common/util/vectorized_pointwise.h @@ -434,9 +434,8 @@ __launch_bounds__(unary_kernel_threads) __global__ ComputeType val2 = static_cast(loader1.separate()[i]); if constexpr (std::is_same::value) { - // Clamp the gated value and add 1 at the end ComputeType limit = p.limit; - val2 = std::min(std::max(-limit, val2), limit) + 1; + val2 = std::min(std::max(-limit, val2), limit) + p.glu_linear_offset; } ComputeType temp = static_cast(Activation(val, p) * val2); if (requires_amax) { @@ -542,10 +541,9 @@ __launch_bounds__(unary_kernel_threads) __global__ bool dgate_in = true; if constexpr (std::is_same::value) { - // In case of GPT OSS, clamp the activation and gate values const ComputeType limit = p.limit; - dgate_in = gate_in <= limit && gate_in >= -limit; // Derivative of clamp - gate_in = std::min(std::max(-limit, gate_in), limit) + 1.0f; + dgate_in = gate_in <= limit && gate_in >= -limit; + gate_in = std::min(std::max(-limit, gate_in), limit) + p.glu_linear_offset; } ComputeType after_dgelu = Dactivation(gelu_in, p) * grad_val * gate_in; diff --git a/transformer_engine/jax/cpp_extensions/activation.py b/transformer_engine/jax/cpp_extensions/activation.py index 8c0edae97..5058192c3 100644 --- a/transformer_engine/jax/cpp_extensions/activation.py +++ b/transformer_engine/jax/cpp_extensions/activation.py @@ -64,6 +64,7 @@ class ClampedSwigluParams: limit: float = 7.0 alpha: float = 1.702 + glu_linear_offset: float = 1.0 def __hash__(self): """Custom hash function to ensure dataclass is hashable for jax jit to work. @@ -71,7 +72,7 @@ def __hash__(self): Returns: int: Hash value of the dataclass instance. """ - return hash((self.limit, self.alpha)) + return hash((self.limit, self.alpha, self.glu_linear_offset)) def to_ffi_lowering_dict(self): """Convert the activation parameters to a dictionary format for FFI lowering. @@ -80,7 +81,11 @@ def to_ffi_lowering_dict(self): dict: A dictionary representation of the activation parameters consumable by XLA FFI bindings for activation functions. """ - return {"limit": np.float32(self.limit), "alpha": np.float32(self.alpha)} + return { + "limit": np.float32(self.limit), + "alpha": np.float32(self.alpha), + "glu_linear_offset": np.float32(self.glu_linear_offset), + } @dataclass(frozen=True) @@ -121,11 +126,9 @@ def _convert_to_activation_function(fn_or_string, act_params: ActivationParams): if fn_or_string == "linear": return lambda x: x if fn_or_string == "clamped_linear": - # This function is used for ClampedSwiGLU - # used in GPT OSS where the gates are not only clamped - # but also shifted by +1 limit = act_params.clamped_swiglu.limit - return lambda x: jnp.clip(x, min=-limit, max=limit) + 1 + offset = act_params.clamped_swiglu.glu_linear_offset + return lambda x: jnp.clip(x, min=-limit, max=limit) + offset if fn_or_string == "quick_gelu": return lambda x: jax.nn.sigmoid(1.702 * x) * x if fn_or_string == "squared_relu": diff --git a/transformer_engine/jax/csrc/extensions.h b/transformer_engine/jax/csrc/extensions.h index 2ecfedc8a..416b18ada 100644 --- a/transformer_engine/jax/csrc/extensions.h +++ b/transformer_engine/jax/csrc/extensions.h @@ -39,6 +39,7 @@ namespace jax { struct ClampedSwigluConfig { float limit; float alpha; + float glu_linear_offset; }; struct ActivationConfig { @@ -208,7 +209,8 @@ pybind11::tuple GetTopkWorkspaceSizes(int batch_size, int seq_len, int k); XLA_FFI_REGISTER_STRUCT_ATTR_DECODING(transformer_engine::jax::ClampedSwigluConfig, ::xla::ffi::StructMember("limit"), - ::xla::ffi::StructMember("alpha")); + ::xla::ffi::StructMember("alpha"), + ::xla::ffi::StructMember("glu_linear_offset")); XLA_FFI_REGISTER_STRUCT_ATTR_DECODING( transformer_engine::jax::ActivationConfig, diff --git a/transformer_engine/jax/csrc/extensions/activation.cpp b/transformer_engine/jax/csrc/extensions/activation.cpp index ce5828d6f..6325a700d 100644 --- a/transformer_engine/jax/csrc/extensions/activation.cpp +++ b/transformer_engine/jax/csrc/extensions/activation.cpp @@ -23,6 +23,7 @@ Error_Type ActLuFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_Type scal // parameters for clamped swiglu used in GPT OSS auto swiglu_limit = act_params.clamped_swiglu.limit; auto swiglu_alpha = act_params.clamped_swiglu.alpha; + auto swiglu_glu_linear_offset = act_params.clamped_swiglu.glu_linear_offset; auto in_dtype = convert_ffi_datatype_to_te_dtype(input_buf.element_type()); auto out_dtype = convert_ffi_datatype_to_te_dtype(output_buf->element_type()); @@ -137,8 +138,8 @@ Error_Type ActLuFFI(cudaStream_t stream, Buffer_Type input_buf, Buffer_Type scal nvte_sreglu(input_tensor.data(), output_tensor.data(), stream); break; case NVTE_Activation_Type::CLAMPED_SWIGLU: - nvte_clamped_swiglu(input_tensor.data(), output_tensor.data(), swiglu_limit, swiglu_alpha, - stream); + nvte_clamped_swiglu_v2(input_tensor.data(), output_tensor.data(), swiglu_limit, swiglu_alpha, + swiglu_glu_linear_offset, stream); break; default: NVTE_ERROR("Unsupported ActivationEnum"); @@ -271,6 +272,7 @@ Error_Type DActLuDBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, // parameters for clamped swiglu used in GPT OSS auto swiglu_limit = act_params.clamped_swiglu.limit; auto swiglu_alpha = act_params.clamped_swiglu.alpha; + auto swiglu_glu_linear_offset = act_params.clamped_swiglu.glu_linear_offset; auto in_dtype = convert_ffi_datatype_to_te_dtype(input_buf.element_type()); auto out_dtype = convert_ffi_datatype_to_te_dtype(output_buf->element_type()); @@ -446,8 +448,8 @@ Error_Type DActLuDBiasQuantizeFFI(cudaStream_t stream, Buffer_Type input_buf, nvte_dsreglu(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), stream); break; case NVTE_Activation_Type::CLAMPED_SWIGLU: - nvte_clamped_dswiglu(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), - swiglu_limit, swiglu_alpha, stream); + nvte_clamped_dswiglu_v2(input_tensor.data(), act_input_tensor.data(), output_tensor.data(), + swiglu_limit, swiglu_alpha, swiglu_glu_linear_offset, stream); break; default: NVTE_ERROR("Unsupported ActivationEnum"); diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index 8082ff07e..f8ac778aa 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -274,10 +274,11 @@ py::object swiglu(const at::Tensor &input, py::handle quantizer); py::object dswiglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer); -py::object clamped_swiglu(const at::Tensor &input, py::handle quantizer, float limit, float alpha); +py::object clamped_swiglu(const at::Tensor &input, py::handle quantizer, float limit, float alpha, + float glu_linear_offset); py::object clamped_dswiglu(const at::Tensor &grad, const at::Tensor &input, py::handle quantizer, - float limit, float alpha); + float limit, float alpha, float glu_linear_offset); /*************************************************************************************************** * LayerNorm **************************************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/activation.cpp b/transformer_engine/pytorch/csrc/extensions/activation.cpp index cab9fab30..58a8f84f8 100644 --- a/transformer_engine/pytorch/csrc/extensions/activation.cpp +++ b/transformer_engine/pytorch/csrc/extensions/activation.cpp @@ -330,13 +330,16 @@ py::object dswiglu(const at::Tensor& grad, const at::Tensor& input, py::handle q } /* clamped functions */ -py::object clamped_swiglu(const at::Tensor& input, py::handle quantizer, float limit, float alpha) { - return activation_helper(input, quantizer, 2, limit, alpha); +py::object clamped_swiglu(const at::Tensor& input, py::handle quantizer, float limit, float alpha, + float glu_linear_offset) { + return activation_helper(input, quantizer, 2, limit, alpha, + glu_linear_offset); } py::object clamped_dswiglu(const at::Tensor& grad, const at::Tensor& input, py::handle quantizer, - float limit, float alpha) { - return dactivation_helper(grad, input, quantizer, limit, alpha); + float limit, float alpha, float glu_linear_offset) { + return dactivation_helper(grad, input, quantizer, limit, alpha, + glu_linear_offset); } } // namespace pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index a4571c64e..b5b563882 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -187,7 +187,7 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("quantizer")); m.def("clamped_swiglu", transformer_engine::pytorch::clamped_swiglu, "SwiGLU activation used in GPT OSS", py::arg("input"), py::arg("quantizer"), - py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f); + py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f, py::arg("glu_linear_offset") = 1.0f); /* Backward of GLU */ m.def("dglu", transformer_engine::pytorch::dglu, "Backward of GLU", py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer")); @@ -216,7 +216,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("fwd_input"), py::arg("quantizer")); m.def("clamped_dswiglu", transformer_engine::pytorch::clamped_dswiglu, "Backward of SwiGLU used in GPT OSS", py::arg("grad"), py::arg("fwd_input"), - py::arg("quantizer"), py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f); + py::arg("quantizer"), py::arg("limit") = 7.0f, py::arg("alpha") = 1.702f, + py::arg("glu_linear_offset") = 1.0f); /* DBias + DAct fusions*/ m.def("dbias_dgelu", transformer_engine::pytorch::dbias_dgelu, "DGeLU + DBias + Quantize", py::arg("grad"), py::arg("fwd_input"), py::arg("quantizer")); diff --git a/transformer_engine/pytorch/module/layernorm_mlp.py b/transformer_engine/pytorch/module/layernorm_mlp.py index 46918ff0f..c837af9d3 100644 --- a/transformer_engine/pytorch/module/layernorm_mlp.py +++ b/transformer_engine/pytorch/module/layernorm_mlp.py @@ -1798,7 +1798,7 @@ class LayerNormMLP(TransformerEngineBaseModule): activation_params : dict, default = None Additional parameters for the activation function. At the moment, only used for ``'clamped_swiglu'`` activation which - supports ``'limit'`` and ``'alpha'`` parameters. + supports ``'limit'``, ``'alpha'``, and ``'glu_linear_offset'`` parameters. init_method : Callable, default = None used for initializing FC1 weights in the following way: ``init_method(weight)``. When set to ``None``, defaults to ``torch.nn.init.normal_(mean=0.0, std=0.023)``. @@ -2501,17 +2501,16 @@ def onnx_forward( fc1_out = fc1_out.to(torch.float32) # activation is computed in fp32 act_params = self.activation_params or {} - # Default params for clamped_swiglu in Transformer Engine - clamped_swiglu_limit, clamped_swiglu_alpha = act_params.get("limit", 7.0), act_params.get( - "alpha", 1.702 - ) + clamped_swiglu_limit = act_params.get("limit", 7.0) + clamped_swiglu_alpha = act_params.get("alpha", 1.702) + clamped_swiglu_offset = act_params.get("glu_linear_offset", 1.0) - def _clamped_swiglu(x, limit, alpha): + def _clamped_swiglu(x, limit, alpha, offset): x_glu, x_linear = x.chunk(2, dim=-1) x_glu = x_glu.clamp(min=None, max=limit) x_linear = x_linear.clamp(min=-limit, max=limit) out_glu = x_glu * torch.sigmoid(alpha * x_glu) - y = out_glu * (x_linear + 1) + y = out_glu * (x_linear + offset) return y activation_map = { @@ -2529,7 +2528,7 @@ def _clamped_swiglu(x, limit, alpha): "silu": torch.nn.functional.silu, "swiglu": lambda x: torch.nn.functional.silu(x.chunk(2, -1)[0]) * x.chunk(2, -1)[1], "clamped_swiglu": lambda x: _clamped_swiglu( - x, clamped_swiglu_limit, clamped_swiglu_alpha + x, clamped_swiglu_limit, clamped_swiglu_alpha, clamped_swiglu_offset ), } if self.activation not in activation_map: diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index c0474220e..717d87201 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -33,11 +33,20 @@ def _cudnn_frontend_version_at_least(min_version: str) -> bool: def _cudnn_frontend_version_supported() -> bool: """Check cuDNN frontend is at least 1.23.0. - All grouped MLP fused-kernel features require cuDNN frontend 1.23.0. + All grouped MLP fused-kernel features require cuDNN frontend >= 1.23.0. """ return _cudnn_frontend_version_at_least("1.23.0") +def _cudnn_frontend_geglu_runtime_params() -> bool: + """Check cuDNN frontend is at least 1.24.0. + + Runtime-configurable GeGLU parameters (linear_offset, geglu_alpha, + glu_clamp_max, glu_clamp_min) require cuDNN frontend >= 1.24.0. + """ + return _cudnn_frontend_version_at_least("1.24.0") + + def _cudnn_frontend_supports_grouped_gemm_srelu() -> bool: """Check cuDNN frontend min version for grouped GEMM SReLU kernels.""" return _cudnn_frontend_version_at_least("1.24.0") @@ -292,8 +301,14 @@ def fuse_grouped_mlp_ops( and isinstance(window[2], GroupedLinear) ): matches_pattern = False - elif isinstance(window[1], ScaledClampedQGeGLU) and ( - abs(window[1]._clamped.alpha - 1.702) > 0.001 + elif ( + isinstance(window[1], ScaledClampedQGeGLU) + and not _cudnn_frontend_geglu_runtime_params() + and ( + abs(window[1]._clamped.alpha - 1.702) > 0.001 + or abs(window[1]._clamped.glu_linear_offset - 1.0) > 0.001 + or abs(window[1]._clamped.limit - 7.0) > 0.001 + ) ): matches_pattern = False else: diff --git a/transformer_engine/pytorch/ops/basic/swiglu.py b/transformer_engine/pytorch/ops/basic/swiglu.py index 9267d9bbb..39571fcd8 100644 --- a/transformer_engine/pytorch/ops/basic/swiglu.py +++ b/transformer_engine/pytorch/ops/basic/swiglu.py @@ -208,6 +208,9 @@ class ClampedSwiGLU(BasicOperation): The clamp limit. alpha : float The scaling factor for the sigmoid function used in the activation. + glu_linear_offset : float + Offset added to the linear (gate) component after clamping. + Set to ``0.0`` to disable the offset. cache_quantized_input : bool, default = ``False`` Quantize input tensor when caching for use in the backward pass. glu_interleave_size : int, optional @@ -222,12 +225,14 @@ def __init__( *, limit: float = 7.0, alpha: float = 1.702, + glu_linear_offset: float = 1.0, cache_quantized_input: bool = False, glu_interleave_size: Optional[int] = None, ): super().__init__() self.limit: float = limit self.alpha: float = alpha + self.glu_linear_offset: float = glu_linear_offset self.cache_quantized_input: bool = cache_quantized_input self.glu_interleave_size: Optional[int] = glu_interleave_size @@ -236,12 +241,13 @@ def _tex_clamped_swiglu_forward( swiglu_in: torch.Tensor, next_op_input_quantizer: Optional[Quantizer], ) -> torch.Tensor: - """Call :func:`tex.clamped_swiglu` with this op's ``limit`` / ``alpha``.""" + """Call :func:`tex.clamped_swiglu` with this op's ``limit`` / ``alpha`` / ``glu_linear_offset``.""" return tex.clamped_swiglu( swiglu_in, next_op_input_quantizer, self.limit, self.alpha, + self.glu_linear_offset, ) def _tex_clamped_dswiglu( @@ -250,13 +256,14 @@ def _tex_clamped_dswiglu( swiglu_in: torch.Tensor, quantizer: Optional[Quantizer], ) -> torch.Tensor: - """Call :func:`tex.clamped_dswiglu` with this op's ``limit`` / ``alpha``.""" + """Call :func:`tex.clamped_dswiglu` with this op's ``limit`` / ``alpha`` / ``glu_linear_offset``.""" return tex.clamped_dswiglu( dy, swiglu_in, quantizer, self.limit, self.alpha, + self.glu_linear_offset, ) def op_forward( @@ -581,6 +588,9 @@ class ScaledClampedQGeGLU(_ScaledGLU): Clamp limit (see :class:`ClampedSwiGLU`). alpha : float, default ``1.702`` Sigmoid scale (see :class:`ClampedSwiGLU`). + glu_linear_offset : float, default ``1.0`` + Offset added to the linear component after clamping + (see :class:`ClampedSwiGLU`). """ @@ -591,6 +601,7 @@ def __init__( activation_recompute_in_mlp: bool = False, limit: float = 7.0, alpha: float = 1.702, + glu_linear_offset: float = 1.0, ) -> None: super().__init__( glu_interleave_size, @@ -599,6 +610,7 @@ def __init__( self._clamped: ClampedSwiGLU = ClampedSwiGLU( limit=limit, alpha=alpha, + glu_linear_offset=glu_linear_offset, ) def _glu_forward(self, swiglu_in: torch.Tensor) -> torch.Tensor: diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 3b6330b22..24aaafc1e 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -22,6 +22,7 @@ from ..fuser import register_backward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext from .._common import ( + _cudnn_frontend_geglu_runtime_params, _cudnn_frontend_version_supported, _cudnn_frontend_supports_grouped_gemm_srelu, fuse_grouped_mlp_ops, @@ -326,6 +327,17 @@ def __init__( "dgeglu" if isinstance(activation, ScaledClampedQGeGLU) else "dswiglu" ) + # cuDNN-frontend >= 1.24.0 exposes runtime-configurable GeGLU + # parameters; pass them through when available. + self._pass_geglu_runtime_params: bool = ( + isinstance(activation, ScaledClampedQGeGLU) and _cudnn_frontend_geglu_runtime_params() + ) + if self._pass_geglu_runtime_params: + self._cudnn_linear_offset: float = activation._clamped.glu_linear_offset + self._cudnn_geglu_alpha: float = activation._clamped.alpha + self._cudnn_glu_clamp_max: float = activation._clamped.limit + self._cudnn_glu_clamp_min: float = -activation._clamped.limit + def fuser_backward( self, basic_op_ctxs: list[OperationContext], @@ -484,6 +496,13 @@ def fuser_backward( fc2_dactivation_kwargs["act_func"] = self._cudnn_dact_func else: fc2_dactivation_kwargs["use_dsrelu_reuse"] = recompute_fc2_x_from_dsrelu + if self._pass_geglu_runtime_params: + fc2_dactivation_kwargs.update( + linear_offset=self._cudnn_linear_offset, + geglu_alpha=self._cudnn_geglu_alpha, + glu_clamp_max=self._cudnn_glu_clamp_max, + glu_clamp_min=self._cudnn_glu_clamp_min, + ) if fc2_op.single_grouped_weight: # Clone and swizzle scales for GEMM diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index f7ac45502..034d40443 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -23,6 +23,7 @@ from ..fuser import register_forward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext from .._common import ( + _cudnn_frontend_geglu_runtime_params, _cudnn_frontend_version_supported, _cudnn_frontend_supports_grouped_gemm_srelu, _nvidia_cudnn_frontend_supports_wgrad, @@ -128,6 +129,18 @@ def __init__( "geglu" if isinstance(activation, ScaledClampedQGeGLU) else "swiglu" ) + # cuDNN-frontend >= 1.24.0 exposes runtime-configurable GeGLU + # parameters; pass them through when the activation carries + # non-default values (or always, if available). + self._pass_geglu_runtime_params: bool = ( + isinstance(activation, ScaledClampedQGeGLU) and _cudnn_frontend_geglu_runtime_params() + ) + if self._pass_geglu_runtime_params: + self._cudnn_linear_offset: float = activation._clamped.glu_linear_offset + self._cudnn_geglu_alpha: float = activation._clamped.alpha + self._cudnn_glu_clamp_max: float = activation._clamped.limit + self._cudnn_glu_clamp_min: float = -activation._clamped.limit + def fuser_forward( self, basic_op_ctxs: list[OperationContext], @@ -333,6 +346,13 @@ def fuser_forward( } if self._cudnn_act_func is not None: fc1_activation_kwargs["act_func"] = self._cudnn_act_func + if self._pass_geglu_runtime_params: + fc1_activation_kwargs.update( + linear_offset=self._cudnn_linear_offset, + geglu_alpha=self._cudnn_geglu_alpha, + glu_clamp_max=self._cudnn_glu_clamp_max, + glu_clamp_min=self._cudnn_glu_clamp_min, + ) if fc1_op.single_grouped_weight: # Clone and swizzle scales for GEMM. diff --git a/transformer_engine/pytorch/transformer.py b/transformer_engine/pytorch/transformer.py index 4b96ccf73..d377e5f3b 100644 --- a/transformer_engine/pytorch/transformer.py +++ b/transformer_engine/pytorch/transformer.py @@ -189,8 +189,9 @@ class TransformerLayer(torch.nn.Module): activation_params : Optional[dict], default = None Additional parameters for the activation function. At the moment, only used for ``'clamped_swiglu'`` activation which - supports ``'limit'`` and ``'alpha'`` parameters. You can set these as - ``activation_params={'limit': 7.0, 'alpha': 1.702}``. + supports ``'limit'``, ``'alpha'``, and ``'glu_linear_offset'`` parameters. + You can set these as + ``activation_params={'limit': 7.0, 'alpha': 1.702, 'glu_linear_offset': 1.0}``. device : Union[torch.device, str], default = "cuda" The device on which the parameters of the model will be allocated. It is the user's responsibility to ensure all parameters are moved to the GPU before running the From 937c4de09a045f6363677eb066c72b7623e0dd0c Mon Sep 17 00:00:00 2001 From: Faradawn Yang <73060648+faradawn@users.noreply.github.com> Date: Tue, 26 May 2026 12:08:05 -0500 Subject: [PATCH 082/180] Add examples for MoE models - Mixtral in TE (#2642) * rebase and add mixtral moe example Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * edit tutorial wording and remove nvtx profiling markers Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address feedback for fused MLP and make table consistent Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * add moe description and code snippet Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * fix image header and grammar Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * add top2 expert in diagram, add perf verdict, and code highlight Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * fix grammar and code highlighting Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * add here is the expected output Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * add code highlight to mxfp8 and callout note Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * change nv mixtral to te Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * rename nv to te in all code Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * add elif Co-authored-by: Sudhakar Singh Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * fix Dtensor and address tiers Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * fix d tensor Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * fix ordering of weights Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * fix backtick and hyperlink Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * fix backtick and hyperlink 2nd Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * remove dtensor Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> * fix none mask Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> --------- Signed-off-by: Faradawn Yang <73060648+faradawn@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Sudhakar Singh Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- docs/examples/te_mixtral/collator.py | 148 ++ docs/examples/te_mixtral/hf_to_te_weights.py | 426 ++++++ .../te_mixtral/media/dense_to_sparse.drawio | 108 ++ .../te_mixtral/media/dense_to_sparse.svg | 1 + .../te_mixtral/media/fused_mlp_path.drawio | 82 ++ .../te_mixtral/media/fused_mlp_path.svg | 1 + .../media/mixtral_decoder_swap.drawio | 100 ++ .../te_mixtral/media/mixtral_decoder_swap.svg | 1 + .../media/moe_loop_vs_grouped.drawio | 85 ++ .../te_mixtral/media/moe_loop_vs_grouped.svg | 1 + docs/examples/te_mixtral/requirements.txt | 10 + docs/examples/te_mixtral/run_finetune_ep.py | 127 ++ docs/examples/te_mixtral/te_mixtral.py | 1274 +++++++++++++++++ docs/examples/te_mixtral/te_mixtral_mxfp8.py | 568 ++++++++ docs/examples/te_mixtral/te_moe_dispatch.py | 298 ++++ docs/examples/te_mixtral/test_accuracy.py | 188 +++ ...torial_accelerate_hf_mixtral_with_te.ipynb | 462 ++++++ docs/examples/te_mixtral/utils.py | 485 +++++++ docs/index.rst | 1 + 19 files changed, 4366 insertions(+) create mode 100644 docs/examples/te_mixtral/collator.py create mode 100644 docs/examples/te_mixtral/hf_to_te_weights.py create mode 100644 docs/examples/te_mixtral/media/dense_to_sparse.drawio create mode 100644 docs/examples/te_mixtral/media/dense_to_sparse.svg create mode 100644 docs/examples/te_mixtral/media/fused_mlp_path.drawio create mode 100644 docs/examples/te_mixtral/media/fused_mlp_path.svg create mode 100644 docs/examples/te_mixtral/media/mixtral_decoder_swap.drawio create mode 100644 docs/examples/te_mixtral/media/mixtral_decoder_swap.svg create mode 100644 docs/examples/te_mixtral/media/moe_loop_vs_grouped.drawio create mode 100644 docs/examples/te_mixtral/media/moe_loop_vs_grouped.svg create mode 100644 docs/examples/te_mixtral/requirements.txt create mode 100644 docs/examples/te_mixtral/run_finetune_ep.py create mode 100644 docs/examples/te_mixtral/te_mixtral.py create mode 100644 docs/examples/te_mixtral/te_mixtral_mxfp8.py create mode 100644 docs/examples/te_mixtral/te_moe_dispatch.py create mode 100644 docs/examples/te_mixtral/test_accuracy.py create mode 100644 docs/examples/te_mixtral/tutorial_accelerate_hf_mixtral_with_te.ipynb create mode 100644 docs/examples/te_mixtral/utils.py diff --git a/docs/examples/te_mixtral/collator.py b/docs/examples/te_mixtral/collator.py new file mode 100644 index 000000000..b9a53cf54 --- /dev/null +++ b/docs/examples/te_mixtral/collator.py @@ -0,0 +1,148 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Data collator for THD sequence packing (variable-length flash attention).""" + +import logging +from dataclasses import dataclass +from typing import Any + +import torch +from transformers import DataCollatorForLanguageModeling + + +logger = logging.getLogger(__name__) + + +def _pt_flatten_collate(features: list[dict[str, list[int]]], return_position_ids: bool = False): + """Flatten a list of tokenized samples into a single packed batch with cumulative sequence lengths.""" + is_labels_provided = "labels" in features[0] + sample_lengths = [len(sample["input_ids"]) for sample in features] + + batch = {} + batch["max_length_q"] = batch["max_length_k"] = max(sample_lengths) + batch["input_ids"] = torch.tensor( + [[token for sample in features for token in sample["input_ids"]]], dtype=torch.int64 + ) + if is_labels_provided: + batch["labels"] = torch.tensor( + [[label for sample in features for label in sample["labels"]]], dtype=torch.int64 + ) + cu_seq_lens = torch.zeros(len(features) + 1, dtype=torch.int32) + cu_seq_lens[1:] = torch.cumsum(torch.tensor(sample_lengths), dim=0, dtype=torch.int32) + batch["cu_seq_lens_q"] = batch["cu_seq_lens_k"] = cu_seq_lens + if "attention_mask" in features[0]: + batch["attention_mask"] = torch.tensor( + [[v for sample in features for v in sample["attention_mask"]]], dtype=torch.int64 + ) + if return_position_ids: + batch["position_ids"] = torch.hstack( + [torch.arange(sample_len, dtype=torch.int64) for sample_len in sample_lengths] + ).unsqueeze(0) + + return batch + + +def _pt_pad_to_multiple_of( + batch: dict[str, Any], pad_to_multiple_of: int, token_pad: int, label_pad: int +): + """Pad a batch to a multiple of ``pad_to_multiple_of`` by appending a mock sequence.""" + remainder = -batch["input_ids"].numel() % pad_to_multiple_of + if remainder == 0: + return batch + + batch["input_ids"] = torch.cat( + [batch["input_ids"], torch.full((1, remainder), token_pad, dtype=batch["input_ids"].dtype)], + dim=1, + ) + if "labels" in batch: + batch["labels"] = torch.cat( + [batch["labels"], torch.full((1, remainder), label_pad, dtype=batch["labels"].dtype)], + dim=1, + ) + if "cu_seq_lens_q" in batch: + batch["cu_seq_lens_q"] = torch.cat( + [ + batch["cu_seq_lens_q"], + torch.tensor( + [batch["cu_seq_lens_q"][-1] + remainder], dtype=batch["cu_seq_lens_q"].dtype + ), + ], + dim=0, + ) + batch["cu_seq_lens_k"] = batch["cu_seq_lens_q"] + if "max_length_q" in batch: + batch["max_length_q"] = max(batch["max_length_q"], remainder) + batch["max_length_k"] = batch["max_length_q"] + if "attention_mask" in batch: + batch["attention_mask"] = torch.cat( + [ + batch["attention_mask"], + torch.zeros((1, remainder), dtype=batch["attention_mask"].dtype), + ], + dim=1, + ) + if "position_ids" in batch: + batch["position_ids"] = torch.cat( + [ + batch["position_ids"], + torch.arange(remainder, dtype=batch["position_ids"].dtype).unsqueeze(0), + ], + dim=1, + ) + + return batch + + +@dataclass +class DataCollatorWithFlattening: + """Data collator that flattens variable-length sequences into a single packed tensor for flash attention. + + Wraps a ``DataCollatorForLanguageModeling`` and produces THD-format batches with + ``cu_seq_lens_q`` / ``cu_seq_lens_k`` metadata for TE's fused attention kernels. + + Args: + collator: The base collator for MLM/CLM masking. + pad_to_multiple_of: If set, pads the total token count to be divisible by this number. + separator_id: Label value inserted at sequence boundaries (typically -100 for causal LM). + """ + + collator: DataCollatorForLanguageModeling + pad_to_multiple_of: int | None = None + separator_id: int | None = None + + def __call__(self, features, return_tensors=None): + """Pack features into a single THD batch with flash-attention metadata.""" + if return_tensors is not None and return_tensors != "pt": + raise NotImplementedError( + f"Only return_tensors='pt' is supported, got '{return_tensors}'" + ) + + bshd_batch = self.collator(features, return_tensors=return_tensors) + packed_batch = _pt_flatten_collate(features) + + masked_input_ids = bshd_batch["input_ids"][bshd_batch["attention_mask"].bool()].unsqueeze(0) + masked_labels = bshd_batch["labels"][bshd_batch["attention_mask"].bool()].unsqueeze(0) + + if self.separator_id is not None: + masked_labels[:, packed_batch["cu_seq_lens_q"][1:-1]] = self.separator_id + + packed_batch["input_ids"] = masked_input_ids + packed_batch["labels"] = masked_labels + + if self.pad_to_multiple_of is not None: + pad_token_id = self.collator.tokenizer.pad_token_id + if not isinstance(pad_token_id, int): + logger.warning( + f"tokenizer.pad_token_id is not an integer, using 1 instead: {pad_token_id}" + ) + pad_token_id = 1 + packed_batch = _pt_pad_to_multiple_of( + packed_batch, + self.pad_to_multiple_of, + token_pad=pad_token_id, + label_pad=-100, + ) + + return packed_batch diff --git a/docs/examples/te_mixtral/hf_to_te_weights.py b/docs/examples/te_mixtral/hf_to_te_weights.py new file mode 100644 index 000000000..86f89d7ab --- /dev/null +++ b/docs/examples/te_mixtral/hf_to_te_weights.py @@ -0,0 +1,426 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""HuggingFace Mixtral -> Transformer Engine state-dict mapping. + +Two top-level entry points share the same top-level / attention / layernorm +plumbing and only differ in how they place the expert MoE weights: + + * :func:`replace_params_bf16` — used by ``te_mixtral.py`` (Improvements 1/2, + BF16). Expert weights land in stacked ``mlp.experts_{gate_up,down}_weight`` + parameters (loop path) and/or per-expert ``mlp.experts_{gate_up,down}.weight{i}`` + parameters (Sequential-Op ``GroupedLinear``). + + * :func:`replace_params_mxfp8` — used by ``te_mixtral_mxfp8.py`` (Improvement 3, + MXFP8). Expert gate (``w1``) and up (``w3``) rows are row-interleaved in + blocks of 32 to match the GLU layout the fused MXFP8 grouped-MLP kernel + expects. +""" + +from __future__ import annotations + +import re + +import torch +import torch.distributed as dist +from transformers import MixtralConfig + + +# Block size for the gate/up interleaved layout. Must match the +# ``glu_interleave_size`` configured on ``ScaledSwiGLU`` in the MXFP8 MoE +# block (see ``te_mixtral_mxfp8.py``). +GLU_INTERLEAVE_SIZE = 32 + + +# --------------------------------------------------------------------------- +# Low-level copy helpers +# --------------------------------------------------------------------------- + + +def _copy_param(target: torch.Tensor, source: torch.Tensor) -> None: + """Copy ``source`` into ``target`` preserving the target's dtype/device.""" + target.copy_(source.to(device=target.device, dtype=target.dtype)) + + +def _copy_qkv_proj_to_fused( + fused_qkv: torch.Tensor, + proj_weight: torch.Tensor, + proj_kind: str, + config: MixtralConfig, +) -> None: + """Copy one HF Q/K/V projection into the TE fused QKV layout. + + TE interleaves the heads as ``[Q_g_0, ..., Q_g_{h-1}, K_g, V_g]`` per + KV group ``g``; HF stores Q/K/V as separate projections. + """ + head_num = config.num_attention_heads + num_query_groups = config.num_key_value_heads + heads_per_group = head_num // num_query_groups + hidden_size = config.hidden_size + head_size = hidden_size // head_num + qkv_total_dim = head_num + 2 * num_query_groups + + fused_view = fused_qkv.view(qkv_total_dim, head_size, hidden_size) + proj_weight = proj_weight.to(device=fused_view.device, dtype=fused_view.dtype) + + if proj_kind == "q": + q_view = proj_weight.view(head_num, head_size, hidden_size) + for i in range(num_query_groups): + start = (heads_per_group + 2) * i + end = start + heads_per_group + fused_view[start:end].copy_(q_view[i * heads_per_group : (i + 1) * heads_per_group]) + elif proj_kind == "k": + k_view = proj_weight.view(num_query_groups, head_size, hidden_size) + for i in range(num_query_groups): + fused_view[(heads_per_group + 2) * i + heads_per_group].copy_(k_view[i]) + elif proj_kind == "v": + v_view = proj_weight.view(num_query_groups, head_size, hidden_size) + for i in range(num_query_groups): + fused_view[(heads_per_group + 2) * i + heads_per_group + 1].copy_(v_view[i]) + else: + raise ValueError(f"Unsupported proj_kind: {proj_kind}") + + +def _interleave_gate_up( + gate: torch.Tensor, + up: torch.Tensor, + interleave: int = GLU_INTERLEAVE_SIZE, +) -> torch.Tensor: + """Interleave HF gate (``w1``) and up (``w3``) rows in blocks of ``interleave``. + + HF stacks gate-then-up along the output dim (``[I gate rows; I up rows]``); + the fused MXFP8 kernel reads gate_up's output in the GLU-interleaved layout + ``[B gate; B up; B gate; B up; ...]`` with ``B = interleave``. + """ + intermediate_size, hidden = gate.shape + if up.shape != (intermediate_size, hidden): + raise ValueError(f"gate and up shape mismatch: {gate.shape} vs {up.shape}") + if intermediate_size % interleave != 0: + raise ValueError(f"intermediate_size {intermediate_size} must be divisible by {interleave}") + g = gate.reshape(intermediate_size // interleave, interleave, hidden) + u = up.reshape(intermediate_size // interleave, interleave, hidden) + stacked = torch.stack([g, u], dim=1) # [I/B, 2, B, H] + return stacked.reshape(2 * intermediate_size, hidden).contiguous() + + +# --------------------------------------------------------------------------- +# Shared per-layer plumbing (top-level, attention, router gate) +# --------------------------------------------------------------------------- + + +def _ep_rank_from_config(config: MixtralConfig) -> tuple[int, int]: + """Return (ep_size, ep_rank). EP rank is global_rank % ep_size.""" + ep_size = int(getattr(config, "expert_parallel_size", 1)) + world_rank = dist.get_rank() if dist.is_available() and dist.is_initialized() else 0 + ep_rank = world_rank % ep_size if ep_size > 1 else 0 + return ep_size, ep_rank + + +def _collect_layer_prefixes(hf_state_dict: dict) -> set[str]: + prefixes = set() + for key in hf_state_dict.keys(): + m = re.match(r"model\.layers\.\d+\.", key) + if m is not None: + prefixes.add(m.group()) + return prefixes + + +def _copy_top_level(hf_state_dict: dict, te_state_dict: dict) -> None: + direct = { + "model.embed_tokens.weight": "model.embed_tokens.weight", + "model.norm.weight": "model.norm.weight", + "lm_head.weight": "lm_head.weight", + "model.rotary_emb.inv_freq": "model.rotary_emb.inv_freq", + } + for hf_key, te_key in direct.items(): + if hf_key in hf_state_dict and te_key in te_state_dict: + _copy_param(te_state_dict[te_key], hf_state_dict[hf_key]) + + +def _copy_attention_and_layernorms( + hf_state_dict: dict, te_state_dict: dict, layer_prefix: str, config: MixtralConfig +) -> None: + direct = { + layer_prefix + + "input_layernorm.weight": layer_prefix + + "self_attention.layernorm_qkv.layer_norm_weight", + layer_prefix + "self_attn.o_proj.weight": layer_prefix + "self_attention.proj.weight", + layer_prefix + + "post_attention_layernorm.weight": layer_prefix + + "post_attention_layernorm.weight", + } + for hf_key, te_key in direct.items(): + if hf_key in hf_state_dict and te_key in te_state_dict: + _copy_param(te_state_dict[te_key], hf_state_dict[hf_key]) + + fused_qkv_key = layer_prefix + "self_attention.layernorm_qkv.weight" + if fused_qkv_key in te_state_dict: + qkv_sources = { + "q": layer_prefix + "self_attn.q_proj.weight", + "k": layer_prefix + "self_attn.k_proj.weight", + "v": layer_prefix + "self_attn.v_proj.weight", + } + for proj_kind, hf_key in qkv_sources.items(): + if hf_key in hf_state_dict: + _copy_qkv_proj_to_fused( + te_state_dict[fused_qkv_key], hf_state_dict[hf_key], proj_kind, config + ) + + +def _copy_router_gate(hf_state_dict: dict, te_state_dict: dict, layer_prefix: str) -> None: + candidates = ( + layer_prefix + "mlp.gate.weight", + layer_prefix + "block_sparse_moe.gate.weight", + ) + te_gate_key = layer_prefix + "mlp.gate.weight" + for hf_key in candidates: + if hf_key in hf_state_dict and te_gate_key in te_state_dict: + _copy_param(te_state_dict[te_gate_key], hf_state_dict[hf_key]) + return + + +def _packed_expert_candidates(layer_prefix: str) -> tuple[tuple[str, ...], tuple[str, ...]]: + gate_up = ( + layer_prefix + "mlp.experts.gate_up_proj", + layer_prefix + "block_sparse_moe.experts.gate_up_proj", + ) + down = ( + layer_prefix + "mlp.experts.down_proj", + layer_prefix + "block_sparse_moe.experts.down_proj", + ) + return gate_up, down + + +def _sequential_op_keys(te_state_dict: dict, layer_prefix: str) -> tuple[list[str], list[str]]: + """Return sorted lists of TE Sequential-Op ``weight{i}`` keys, if present.""" + gate_up_prefix = layer_prefix + "mlp.experts_gate_up." + down_prefix = layer_prefix + "mlp.experts_down." + + def _weight_index(key: str) -> int: + match = re.search(r"weight(\d+)$", key) + assert match is not None + return int(match.group(1)) + + gate_up_keys = sorted( + (k for k in te_state_dict if k.startswith(gate_up_prefix) and re.search(r"weight\d+$", k)), + key=_weight_index, + ) + down_keys = sorted( + (k for k in te_state_dict if k.startswith(down_prefix) and re.search(r"weight\d+$", k)), + key=_weight_index, + ) + return gate_up_keys, down_keys + + +# --------------------------------------------------------------------------- +# Public entry points +# --------------------------------------------------------------------------- + + +def replace_params_bf16( + hf_state_dict: dict, te_state_dict: dict, config: MixtralConfig +) -> set[str]: + """Map HF Mixtral weights into the BF16 TE state dict. + + Expert weights are placed in stacked ``mlp.experts_{gate_up,down}_weight`` + parameters (loop path) and/or per-expert Sequential-Op ``weight{i}`` + parameters (grouped_op path). Both formats are written so a single + checkpoint loader supports either ``expert_ffn_mode``. + + Supports both packed HF MoE tensors (``mlp.experts.gate_up_proj``) and + older per-expert tensors (``experts.{i}.w{1,2,3}.weight``). + """ + ep_size, ep_rank = _ep_rank_from_config(config) + layer_prefixes = _collect_layer_prefixes(hf_state_dict) + _copy_top_level(hf_state_dict, te_state_dict) + + for layer_prefix in layer_prefixes: + _copy_attention_and_layernorms(hf_state_dict, te_state_dict, layer_prefix, config) + _copy_router_gate(hf_state_dict, te_state_dict, layer_prefix) + + packed_gate_up_candidates, packed_down_candidates = _packed_expert_candidates(layer_prefix) + te_gate_up_key = layer_prefix + "mlp.experts_gate_up_weight" + te_down_key = layer_prefix + "mlp.experts_down_weight" + + # Path A: stacked param (loop path) <- packed HF tensor. + for hf_key in packed_gate_up_candidates: + if hf_key in hf_state_dict and te_gate_up_key in te_state_dict: + te_gate_up = te_state_dict[te_gate_up_key] + local_experts = te_gate_up.shape[0] + expert_start = ep_rank * local_experts if ep_size > 1 else 0 + expert_end = expert_start + local_experts + _copy_param( + te_state_dict[te_gate_up_key], + hf_state_dict[hf_key][expert_start:expert_end], + ) + break + for hf_key in packed_down_candidates: + if hf_key in hf_state_dict and te_down_key in te_state_dict: + te_down = te_state_dict[te_down_key] + local_experts = te_down.shape[0] + expert_start = ep_rank * local_experts if ep_size > 1 else 0 + expert_end = expert_start + local_experts + _copy_param( + te_state_dict[te_down_key], + hf_state_dict[hf_key][expert_start:expert_end], + ) + break + + # Path B: Sequential-Op per-expert params <- packed HF tensor. + te_gate_up_op_keys, te_down_op_keys = _sequential_op_keys(te_state_dict, layer_prefix) + if te_gate_up_op_keys and te_down_op_keys: + num_local_experts = len(te_gate_up_op_keys) + expert_start = ep_rank * num_local_experts if ep_size > 1 else 0 + expert_end = expert_start + num_local_experts + for hf_key in packed_gate_up_candidates: + if hf_key in hf_state_dict: + hf_gate_up = hf_state_dict[hf_key][expert_start:expert_end] + for expert_idx, te_key in enumerate(te_gate_up_op_keys): + _copy_param(te_state_dict[te_key], hf_gate_up[expert_idx]) + break + for hf_key in packed_down_candidates: + if hf_key in hf_state_dict: + hf_down = hf_state_dict[hf_key][expert_start:expert_end] + for expert_idx, te_key in enumerate(te_down_op_keys): + _copy_param(te_state_dict[te_key], hf_down[expert_idx]) + break + + # Path C: older HF format with per-expert w1/w2/w3 -> stacked params. + if te_gate_up_key in te_state_dict and te_down_key in te_state_dict: + te_gate_up = te_state_dict[te_gate_up_key] + te_down = te_state_dict[te_down_key] + num_local_experts = te_gate_up.shape[0] + expert_start = ep_rank * num_local_experts if ep_size > 1 else 0 + for expert_idx in range(num_local_experts): + global_expert_idx = expert_start + expert_idx + for expert_prefix in ( + layer_prefix + f"mlp.experts.{global_expert_idx}.", + layer_prefix + f"block_sparse_moe.experts.{global_expert_idx}.", + ): + w1_key = expert_prefix + "w1.weight" + w3_key = expert_prefix + "w3.weight" + w2_key = expert_prefix + "w2.weight" + if w1_key in hf_state_dict: + te_gate_up[expert_idx, : config.intermediate_size].copy_( + hf_state_dict[w1_key].to( + device=te_gate_up.device, dtype=te_gate_up.dtype + ) + ) + if w3_key in hf_state_dict: + te_gate_up[expert_idx, config.intermediate_size :].copy_( + hf_state_dict[w3_key].to( + device=te_gate_up.device, dtype=te_gate_up.dtype + ) + ) + if w2_key in hf_state_dict: + te_down[expert_idx].copy_( + hf_state_dict[w2_key].to(device=te_down.device, dtype=te_down.dtype) + ) + + # Path D: older HF format -> Sequential-Op per-expert params. + if te_gate_up_op_keys and te_down_op_keys: + num_local_experts = len(te_gate_up_op_keys) + expert_start = ep_rank * num_local_experts if ep_size > 1 else 0 + for expert_idx in range(num_local_experts): + global_expert_idx = expert_start + expert_idx + gate_up = te_state_dict[te_gate_up_op_keys[expert_idx]] + down = te_state_dict[te_down_op_keys[expert_idx]] + for expert_prefix in ( + layer_prefix + f"mlp.experts.{global_expert_idx}.", + layer_prefix + f"block_sparse_moe.experts.{global_expert_idx}.", + ): + w1_key = expert_prefix + "w1.weight" + w3_key = expert_prefix + "w3.weight" + w2_key = expert_prefix + "w2.weight" + if w1_key in hf_state_dict: + gate_up[: config.intermediate_size].copy_( + hf_state_dict[w1_key].to(device=gate_up.device, dtype=gate_up.dtype) + ) + if w3_key in hf_state_dict: + gate_up[config.intermediate_size :].copy_( + hf_state_dict[w3_key].to(device=gate_up.device, dtype=gate_up.dtype) + ) + if w2_key in hf_state_dict: + down.copy_(hf_state_dict[w2_key].to(device=down.device, dtype=down.dtype)) + + return layer_prefixes + + +def replace_params_mxfp8( + hf_state_dict: dict, te_state_dict: dict, config: MixtralConfig +) -> set[str]: + """Map HF Mixtral weights into the MXFP8 TE state dict. + + Per-expert gate/up rows are interleaved in blocks of 32 to match the GLU + layout that the fused MXFP8 grouped-MLP kernel reads. + + Supports both packed HF tensors (``mlp.experts.gate_up_proj`` of shape + ``[E, 2I, H]``) and older per-expert tensors + (``mlp.experts.{i}.w{1,2,3}.weight``). + """ + ep_size, ep_rank = _ep_rank_from_config(config) + layer_prefixes = _collect_layer_prefixes(hf_state_dict) + _copy_top_level(hf_state_dict, te_state_dict) + + for layer_prefix in layer_prefixes: + _copy_attention_and_layernorms(hf_state_dict, te_state_dict, layer_prefix, config) + _copy_router_gate(hf_state_dict, te_state_dict, layer_prefix) + + te_gate_up_op_keys, te_down_op_keys = _sequential_op_keys(te_state_dict, layer_prefix) + if not (te_gate_up_op_keys and te_down_op_keys): + continue + + num_local_experts = len(te_gate_up_op_keys) + expert_start = ep_rank * num_local_experts if ep_size > 1 else 0 + expert_end = expert_start + num_local_experts + intermediate_size = config.intermediate_size + + packed_gate_up_candidates, packed_down_candidates = _packed_expert_candidates(layer_prefix) + + # Path A: newer HF format with packed gate_up tensor [E, 2I, H]. + packed_gate_up_done = False + for hf_key in packed_gate_up_candidates: + if hf_key not in hf_state_dict: + continue + hf_gate_up = hf_state_dict[hf_key][expert_start:expert_end] + for expert_idx, te_key in enumerate(te_gate_up_op_keys): + gate = hf_gate_up[expert_idx, :intermediate_size] + up = hf_gate_up[expert_idx, intermediate_size:] + interleaved = _interleave_gate_up(gate, up, GLU_INTERLEAVE_SIZE) + _copy_param(te_state_dict[te_key], interleaved) + packed_gate_up_done = True + break + + packed_down_done = False + for hf_key in packed_down_candidates: + if hf_key not in hf_state_dict: + continue + hf_down = hf_state_dict[hf_key][expert_start:expert_end] + for expert_idx, te_key in enumerate(te_down_op_keys): + _copy_param(te_state_dict[te_key], hf_down[expert_idx]) + packed_down_done = True + break + + if packed_gate_up_done and packed_down_done: + continue + + # Path B: older HF format with per-expert w1/w2/w3 weights. + for expert_idx in range(num_local_experts): + global_expert_idx = expert_start + expert_idx + for expert_prefix in ( + layer_prefix + f"mlp.experts.{global_expert_idx}.", + layer_prefix + f"block_sparse_moe.experts.{global_expert_idx}.", + ): + w1_key = expert_prefix + "w1.weight" + w3_key = expert_prefix + "w3.weight" + w2_key = expert_prefix + "w2.weight" + if w1_key in hf_state_dict and w3_key in hf_state_dict and w2_key in hf_state_dict: + gate = hf_state_dict[w1_key] + up = hf_state_dict[w3_key] + interleaved = _interleave_gate_up(gate, up, GLU_INTERLEAVE_SIZE) + _copy_param(te_state_dict[te_gate_up_op_keys[expert_idx]], interleaved) + _copy_param(te_state_dict[te_down_op_keys[expert_idx]], hf_state_dict[w2_key]) + break + + return layer_prefixes diff --git a/docs/examples/te_mixtral/media/dense_to_sparse.drawio b/docs/examples/te_mixtral/media/dense_to_sparse.drawio new file mode 100644 index 000000000..f222bc07e --- /dev/null +++ b/docs/examples/te_mixtral/media/dense_to_sparse.drawio @@ -0,0 +1,108 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/docs/examples/te_mixtral/media/dense_to_sparse.svg b/docs/examples/te_mixtral/media/dense_to_sparse.svg new file mode 100644 index 000000000..971a94a9f --- /dev/null +++ b/docs/examples/te_mixtral/media/dense_to_sparse.svg @@ -0,0 +1 @@ +
Dense Transformer Block
Dense Transformer Block
Sparse Transformer Block
Sparse Transformer Block
hidden_states
hidden_states
Self-Attention
Self-Attention
Residual + RMSNorm
Residual + RMSNorm
Dense MLP
Dense MLP
gate_proj, up_proj
gate_proj, up_proj
down_proj
down_proj
Residual + RMSNorm
Residual + RMSNorm
output
output
hidden_states
hidden_states
Self-Attention
Self-Attention
Residual + RMSNorm
Residual + RMSNorm
Sparse MoE
Sparse MoE
E0
E0
E1
E1
E2
E2
E3
E3
E4
E4
E5
E5
E6
E6
E7
E7
Residual + RMSNorm
Residual + RMSNorm
output
output
Router
RouterText is not SVG - cannot display \ No newline at end of file diff --git a/docs/examples/te_mixtral/media/fused_mlp_path.drawio b/docs/examples/te_mixtral/media/fused_mlp_path.drawio new file mode 100644 index 000000000..a70123ab7 --- /dev/null +++ b/docs/examples/te_mixtral/media/fused_mlp_path.drawio @@ -0,0 +1,82 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/docs/examples/te_mixtral/media/fused_mlp_path.svg b/docs/examples/te_mixtral/media/fused_mlp_path.svg new file mode 100644 index 000000000..cbd777962 --- /dev/null +++ b/docs/examples/te_mixtral/media/fused_mlp_path.svg @@ -0,0 +1 @@ +
Unfused MLP
Unfused MLP
Fused MLP
Fused MLP
Quantize
Quantize
Gate Up
Gate Up
SwiGLU
SwiGLU
De-quantize
De-quanti...
Gate Down
Gate Down
Quantize
Quantize
Fused Group MLP
Fused Group MLP
Gate Down
Gate DownText is not SVG - cannot display \ No newline at end of file diff --git a/docs/examples/te_mixtral/media/mixtral_decoder_swap.drawio b/docs/examples/te_mixtral/media/mixtral_decoder_swap.drawio new file mode 100644 index 000000000..093d2dfa9 --- /dev/null +++ b/docs/examples/te_mixtral/media/mixtral_decoder_swap.drawio @@ -0,0 +1,100 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/docs/examples/te_mixtral/media/mixtral_decoder_swap.svg b/docs/examples/te_mixtral/media/mixtral_decoder_swap.svg new file mode 100644 index 000000000..c35a47c30 --- /dev/null +++ b/docs/examples/te_mixtral/media/mixtral_decoder_swap.svg @@ -0,0 +1 @@ +
HF Transformer Block
HF Transformer Block
TE Transformer Block
TE Transformer Block
input_layernorm
input_layernorm
q_proj
q_proj
k_proj
k_proj
v_proj
v_proj
o_proj
o_proj
post_attention_layernorm
post_attention_layernorm
mlp.gate
mlp.gate
mlp.experts.gate_up_proj
mlp.experts.gate_up_proj
mlp.experts.down_proj
mlp.experts.down_proj
layernorm_qkv.layer_norm
query, key, value
self_attention.proj
layernorm_qkv.layer_norm...
post_attention_layernorm
post_attention_layernorm
mlp.gate
mlp.gate
mlp.experts_gate_up
mlp.experts_gate_up
mlp.experts_down
mlp.experts_downText is not SVG - cannot display \ No newline at end of file diff --git a/docs/examples/te_mixtral/media/moe_loop_vs_grouped.drawio b/docs/examples/te_mixtral/media/moe_loop_vs_grouped.drawio new file mode 100644 index 000000000..72fb7ba9b --- /dev/null +++ b/docs/examples/te_mixtral/media/moe_loop_vs_grouped.drawio @@ -0,0 +1,85 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + \ No newline at end of file diff --git a/docs/examples/te_mixtral/media/moe_loop_vs_grouped.svg b/docs/examples/te_mixtral/media/moe_loop_vs_grouped.svg new file mode 100644 index 000000000..7907a3238 --- /dev/null +++ b/docs/examples/te_mixtral/media/moe_loop_vs_grouped.svg @@ -0,0 +1 @@ +
HF MoE — Python loop 
HF MoE — Python loop 
TE MoE — Grouped GEMM
TE MoE — Grouped GEMM
time
time
E0
E0
E1
E1
E2
E2
E3
E3
E4
E4
E5
E5
E6
E6
E7
E7
GroupedLinear
GroupedLinear
E0
E0
E1
E1
E2
E2
E3
E3
E4
E4
E5
E5
E6
E6
E7
E7Text is not SVG - cannot display \ No newline at end of file diff --git a/docs/examples/te_mixtral/requirements.txt b/docs/examples/te_mixtral/requirements.txt new file mode 100644 index 000000000..5ad5e71db --- /dev/null +++ b/docs/examples/te_mixtral/requirements.txt @@ -0,0 +1,10 @@ +torchao!=0.14.0 +transformer_engine[pytorch] +transformers==5.8.0 +accelerate==1.13.0 +datasets==4.8.4 +safetensors==0.7.0 +huggingface_hub==1.10.1 +tokenizers==0.22.2 +flash-attn +nvidia-cudnn-frontend>=1.23.0 diff --git a/docs/examples/te_mixtral/run_finetune_ep.py b/docs/examples/te_mixtral/run_finetune_ep.py new file mode 100644 index 000000000..e389b7986 --- /dev/null +++ b/docs/examples/te_mixtral/run_finetune_ep.py @@ -0,0 +1,127 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""EP fine-tune launcher for TE Mixtral. + + python3 run_finetune_ep.py --improvement 0 + torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 1 + torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 2 + torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 3 + +Improvements (``--ep-size 2`` => 4 experts/rank on 8 GPUs, DP=4): + + 0 = HF baseline BF16 (single process, ``device_map="auto"``). + 1 = TE EP BF16, Python loop over experts. + 2 = TE EP BF16, GroupedLinear. + 3 = TE EP MXFP8 + fused MXFP8 grouped-MLP kernel. +""" + +import argparse +import os +import sys + +# Improvement 3 needs ``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1`` set before TE is imported, +# because the fused-grouped-MLP fusion is registered at module-import time +# inside ``if ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8.is_supported(): ...``. +for _i, _arg in enumerate(sys.argv[1:]): + if _arg == "--improvement" and _i + 2 < len(sys.argv) and sys.argv[_i + 2] == "3": + os.environ["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] = "1" + break + if _arg == "--improvement=3": + os.environ["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] = "1" + break + +from utils import HyperParameters, run_hf_baseline_finetune, run_te_mixtral_finetune + + +IMPROVEMENT_LABELS = { + 0: "HF baseline BF16", + 1: "TE EP BF16, Python expert loop", + 2: "TE EP BF16, GroupedLinear", + 3: "TE EP MXFP8 fused MLP", +} + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description="Run TE Mixtral fine-tuning with Expert Parallelism." + ) + parser.add_argument("--hf-token", type=str, default=os.environ.get("HF_TOKEN", "")) + parser.add_argument( + "--ep-size", + type=int, + default=2, + help="Expert-parallel group size. Default 2 -> 4 experts/rank with 8 GPUs (DP=4).", + ) + parser.add_argument( + "--improvement", + type=int, + choices=(0, 1, 2, 3), + default=1, + help=( + "Improvement: " + "0=HF baseline BF16, " + "1=TE EP BF16 Python loop, " + "2=TE EP BF16 GroupedLinear, " + "3=TE EP MXFP8 fused MLP." + ), + ) + parser.add_argument("--batch-size", type=int, default=1) + parser.add_argument("--max-seq-length", type=int, default=256) + parser.add_argument("--warmup-steps", type=int, default=1) + parser.add_argument("--train-steps", type=int, default=2) + return parser.parse_args() + + +def main() -> None: + args = parse_args() + + hp = HyperParameters() + hp.model_name = "mistralai/Mixtral-8x7B-v0.1" + hp.hf_access_token = args.hf_token + hp.batch_size = args.batch_size + hp.max_seq_length = args.max_seq_length + hp.num_warmup_steps = args.warmup_steps + hp.num_training_steps = args.train_steps + + if args.improvement == 0: + hp.expert_parallel_size = 1 + hp.mixed_precision = "bf16" + hp.expert_ffn_mode = "loop" # unused: HF baseline doesn't use TE MoE + elif args.improvement == 1: + hp.expert_parallel_size = args.ep_size + hp.mixed_precision = "bf16" + hp.expert_ffn_mode = "loop" + elif args.improvement == 2: + hp.expert_parallel_size = args.ep_size + hp.mixed_precision = "bf16" + hp.expert_ffn_mode = "grouped_op" + elif args.improvement == 3: + hp.expert_parallel_size = args.ep_size + hp.mixed_precision = "mxfp8" + hp.expert_ffn_mode = "grouped_op" + hp.model_impl = "te_mixtral_mxfp8" + + print( + f"[Improvement {args.improvement}] {IMPROVEMENT_LABELS[args.improvement]}\n" + f" mixed_precision={hp.mixed_precision}, ep_size={hp.expert_parallel_size}, " + f"expert_ffn_mode={hp.expert_ffn_mode}\n" + f" batch_size={hp.batch_size}, max_seq_length={hp.max_seq_length}" + ) + + if args.improvement == 0: + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if world_size != 1: + raise ValueError( + "HF baseline must run as a single process (device_map='auto'). " + "Use plain python or torchrun --nproc_per_node=1." + ) + run_hf_baseline_finetune(hp) + return + + run_te_mixtral_finetune(hp) + + +if __name__ == "__main__": + main() diff --git a/docs/examples/te_mixtral/te_mixtral.py b/docs/examples/te_mixtral/te_mixtral.py new file mode 100644 index 000000000..20c6cd6d6 --- /dev/null +++ b/docs/examples/te_mixtral/te_mixtral.py @@ -0,0 +1,1274 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""TransformerEngine-optimized Mixtral model with Mixture of Experts.""" + +import logging +import warnings +from collections import OrderedDict +from contextlib import nullcontext +from dataclasses import dataclass + +from typing import Any, ClassVar, ContextManager, Protocol +from typing_extensions import Unpack + +import torch +import torch.distributed as dist +import torch.nn as nn +import transformer_engine.common.recipe +import transformer_engine.pytorch +import transformers + +from transformer_engine.pytorch.ops import GroupedLinear as TEOpsGroupedLinear +from transformer_engine.pytorch.ops import Sequential as TEOpsSequential +from transformer_engine.pytorch.ops import SwiGLU as TEOpsSwiGLU +from transformer_engine.pytorch.attention import InferenceParams +from transformer_engine.pytorch.attention.inference import PagedKVCacheManager +from transformer_engine.pytorch.attention.rope import RotaryPositionEmbedding +from transformer_engine.pytorch.quantization import FP8GlobalStateManager +from transformer_engine.pytorch.router import fused_moe_aux_loss +from transformers import MixtralConfig, PreTrainedModel +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding +from transformers.utils.generic import TransformersKwargs + +logger = logging.getLogger(__name__) + + +AUTO_MAP = { + "AutoConfig": "modeling_mixtral_te.TEMixtralConfig", + "AutoModel": "modeling_mixtral_te.TEMixtralModel", + "AutoModelForCausalLM": "modeling_mixtral_te.TEMixtralForCausalLM", +} + + +# HF->TE checkpoint mapping lives in ``hf_to_te_weights.py`` so both +# ``te_mixtral.py`` (BF16) and ``te_mixtral_mxfp8.py`` (MXFP8) can share it. +from hf_to_te_weights import ( # noqa: E402 + replace_params_bf16 as replace_params, +) + + +class TEMixtralConfig(MixtralConfig): + """TEMixtral configuration.""" + + # Attention input format: + # "bshd" = Batch, Sequence, Head, Dimension (standard padded format) + # "thd" = Total tokens (packed/unpadded), Head, Dimension (sequence packing format) + attn_input_format: str = "thd" + self_attn_mask_type: str = "padding_causal" + layer_precision: list[str | None] | None = None + use_quantized_model_init: bool = False + expert_parallel_size: int = 1 + moe_aux_loss_coeff: float = 0.0 + # Expert FFN execution mode: + # "grouped_op" — fuse all per-rank experts via the TE Sequential-Op + # ``transformer_engine.pytorch.ops.GroupedLinear``. On + # Blackwell (SM100+) this automatically dispatches to the + # graph-safe ``general_grouped_gemm_for_grouped_tensor`` + # path added in https://github.com/NVIDIA/TransformerEngine/pull/2923 . + # "loop" — naive Python loop, one F.linear per expert. Pedagogical + # baseline so the tutorial can isolate the GroupedLinear win. + expert_ffn_mode: str = "grouped_op" + + def __init__(self, **kwargs): + """Initialize the TEMixtralConfig with additional TE-related config options.""" + super().__init__(**kwargs) + + if self.layer_precision is not None: + if len(self.layer_precision) != self.num_hidden_layers: + raise ValueError( + f"layer_precision must be a list of length {self.num_hidden_layers}" + ) + for precision in self.layer_precision: + if precision not in {"fp8", "fp4", None}: + raise ValueError( + f'layer_precision element must be "fp8", "fp4", or None, got {precision!r}' + ) + + if self.expert_ffn_mode not in ("grouped_op", "loop"): + raise ValueError( + f'expert_ffn_mode must be "grouped_op" or "loop", got {self.expert_ffn_mode!r}' + ) + + if self.num_local_experts % self.expert_parallel_size != 0: + raise ValueError( + f"num_local_experts ({self.num_local_experts}) must be divisible by " + f"expert_parallel_size ({self.expert_parallel_size})" + ) + + +@dataclass +class DispatchOutput: + """Output of TokenDispatcher.dispatch(). + + Attributes: + expert_input: Tokens sorted by local expert, shape ``[total_recv_tokens, H]``. + tokens_per_expert: Token count per local expert. + handle: Opaque state needed by ``combine()`` to reverse the dispatch. + """ + + expert_input: torch.Tensor + tokens_per_expert: list[int] + handle: Any + + +class TokenDispatcher(Protocol): + """Protocol for MoE token dispatch/combine strategies. + + Encapsulates the full dispatch cycle (permute -> communicate -> sort) and + combine cycle (unsort -> communicate -> unpermute) so that the MoE block + is agnostic to the communication backend (NCCL all-to-all, HybridEP, etc.). + """ + + def dispatch( + self, + hidden_states: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + ) -> DispatchOutput: + """Dispatch tokens to their assigned experts. + + Args: + hidden_states: Flattened input tensor of shape ``[N, H]``. + selected_experts: Expert assignments, shape ``[N, top_k]``, int. + routing_weights: Normalized routing probabilities, shape ``[N, top_k]``, float32. + + Returns: + DispatchOutput with expert-sorted tokens, per-expert counts, and an opaque handle. + """ + ... + + def combine( + self, + expert_output: torch.Tensor, + handle: Any, + ) -> torch.Tensor: + """Combine expert outputs back to the original token order. + + Args: + expert_output: Expert output tensor of shape ``[total_recv_tokens, H]``. + handle: Opaque state from ``dispatch()``. + + Returns: + Combined output tensor of shape ``[N, H]`` with routing weights applied. + """ + ... + + def set_ep_group(self, ep_group: dist.ProcessGroup) -> None: + """Set the expert-parallel process group for communication.""" + ... + + +class TEMixtralPreTrainedModel(PreTrainedModel): + """Base class for TEMixtral models.""" + + config_class = TEMixtralConfig + base_model_prefix = "model" + _no_split_modules = ("TEMixtralDecoderLayer",) + _skip_keys_device_placement = ("past_key_values",) + _do_not_quantize = ( + "lm_head", + "model.layers.*.mlp.gate", + ) # Flag for testing that these layers are not quantized. + + def init_empty_weights(self): + """Handles moving the model from the meta device to the cuda device and initializing the weights.""" + for module in self.modules(): + if hasattr(module, "reset_parameters"): + module.reset_parameters() + + # After reset_parameters materializes GroupedLinear views on CUDA, + # re-stack them into the authoritative stacked parameters. + for module in self.modules(): + if isinstance(module, TEMixtralSparseMoeBlock): + module._restack_from_views() + + self.model.embed_tokens.to_empty(device="cuda") + self.model.embed_tokens.apply(self._init_weights) + + self.model.rotary_emb.inv_freq = LlamaRotaryEmbedding(config=self.model.config).inv_freq.to( + "cuda" + ) + + self.tie_weights() + + def _init_weights(self, module): + """Initialize module weights. + + We only use this method for standard pytorch modules, TE modules handle their own weight initialization through + `init_method` parameters and the `reset_parameters` method. + """ + if module.__module__.startswith("transformer_engine.pytorch"): + return + + super()._init_weights(module) + + def state_dict(self, *args, **kwargs): + """Override state_dict to filter out TransformerEngine's _extra_state keys.""" + state_dict = super().state_dict(*args, **kwargs) + return {k: v for k, v in state_dict.items() if not k.endswith("_extra_state")} + + +class TEMixtralSparseMoeBlock(nn.Module): + """Mixture of Experts block using TransformerEngine GroupedLinear.""" + + def __init__(self, config: MixtralConfig, dispatcher: TokenDispatcher | None = None): + """Initialize the sparse MoE block.""" + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.num_experts = config.num_local_experts + self.top_k = config.num_experts_per_tok + self.jitter_noise = config.router_jitter_noise + + self.ep_size = getattr(config, "expert_parallel_size", 1) + self.num_local_experts = self.num_experts // self.ep_size + self.expert_ffn_mode = getattr(config, "expert_ffn_mode", "grouped_op") + self._uses_stacked_expert_weights = self.expert_ffn_mode != "grouped_op" + self.moe_aux_loss_coeff = getattr(config, "moe_aux_loss_coeff", 0.0) + self._aux_loss: torch.Tensor = torch.tensor(0.0) + self.initializer_range = config.initializer_range + + self.dispatcher: TokenDispatcher = dispatcher or AllToAllTokenDispatcher( + self.num_experts, + self.num_local_experts, + self.hidden_size, + self.ep_size, + ) + + device = "meta" if torch.get_default_device() == torch.device("meta") else "cuda" + + def _init_method(x): + torch.nn.init.normal_(x, mean=0.0, std=config.initializer_range) + + # Router always outputs num_experts logits (replicated across EP ranks) + with transformer_engine.pytorch.quantized_model_init(enabled=False): + self.gate = transformer_engine.pytorch.Linear( + self.hidden_size, + self.num_experts, + bias=False, + device=device, + params_dtype=config.dtype, + init_method=_init_method, + ) + + # Expert FFNs — only num_local_experts per rank when EP > 1. + # Both ``grouped_op`` (improvement 2) and ``loop`` (improvement 1) allocate the same + # pair of GroupedLinear ops; ``loop`` just routes its tokens through + # them one-expert-at-a-time in ``_expert_ffn``. + self.experts_gate_up = TEOpsGroupedLinear( + num_groups=self.num_local_experts, + in_features=self.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + dtype=config.dtype, + device=device, + ) + self.experts_down = TEOpsGroupedLinear( + num_groups=self.num_local_experts, + in_features=self.intermediate_size, + out_features=self.hidden_size, + bias=False, + dtype=config.dtype, + device=device, + ) + # ``grouped_op`` runs the two GroupedLinears + SwiGLU through TE's + # fusible Sequential wrapper so the OperationFuser can collapse them. + if self.expert_ffn_mode == "grouped_op": + object.__setattr__( + self, + "_experts_ffn_op", + TEOpsSequential(self.experts_gate_up, TEOpsSwiGLU(), self.experts_down), + ) + + if self._uses_stacked_expert_weights: + # Stack per-expert weights into single parameters (authoritative weight store). + # GroupedLinear's _parameters dict is emptied; weight attributes are set as views + # so that reset_parameters() / _get_weight_tensors() can still find them. + self.experts_gate_up_weight = nn.Parameter( + torch.stack( + [ + self.experts_gate_up._parameters.pop(f"weight{i}").data + for i in range(self.num_local_experts) + ] + ) + ) # [num_local_experts, 2*intermediate_size, hidden_size] + + self.experts_down_weight = nn.Parameter( + torch.stack( + [ + self.experts_down._parameters.pop(f"weight{i}").data + for i in range(self.num_local_experts) + ] + ) + ) # [num_local_experts, hidden_size, intermediate_size] + + # Set views back on GroupedLinear so getattr(self, "weight{i}") still works + # (needed by GroupedLinear.reset_parameters and _get_weight_tensors). + self._sync_expert_views() + + def _restack_from_views(self) -> None: + """Re-create stacked parameters on CUDA after meta init. + + Called by ``init_empty_weights()`` after ``reset_parameters()`` has been called + on all TE modules. Since GroupedLinear has no registered parameters (we popped them), + its ``reset_parameters()`` cannot move them from meta to CUDA. This method explicitly + creates the stacked parameters on CUDA and reinitializes them. + """ + if not self._uses_stacked_expert_weights: + return + + device = torch.cuda.current_device() + for attr_name in ("experts_gate_up_weight", "experts_down_weight"): + old_param = getattr(self, attr_name) + new_data = torch.empty_like(old_param, device=device) + torch.nn.init.normal_(new_data, mean=0.0, std=self.initializer_range) + setattr(self, attr_name, nn.Parameter(new_data)) + + # Re-sync views to point to the new stacked parameter + self._sync_expert_views() + + def _sync_expert_views(self) -> None: + """Set GroupedLinear weight attributes as views of the stacked parameters. + + GroupedLinear internally uses ``getattr(self, f"weight{i}")`` in methods like + ``reset_parameters()`` and ``_get_weight_tensors()``. After popping the original + parameters, we set views of the stacked tensor so these methods keep working. + Uses ``object.__setattr__`` to bypass ``nn.Module.__setattr__`` and avoid + re-registering them as parameters. + """ + if not self._uses_stacked_expert_weights: + return + gate_up_w = self.experts_gate_up_weight + for i in range(self.num_local_experts): + object.__setattr__(self.experts_gate_up, f"weight{i}", gate_up_w[i]) + + down_w = self.experts_down_weight + for i in range(self.num_local_experts): + object.__setattr__(self.experts_down, f"weight{i}", down_w[i]) + + def set_ep_group(self, ep_group: dist.ProcessGroup) -> None: + """Set the expert-parallel process group for token dispatch. + + Must be called before the first forward pass when ``ep_size > 1``. + """ + self.dispatcher.set_ep_group(ep_group) + + def _expert_ffn(self, tokens: torch.Tensor, m_splits: list[int]) -> torch.Tensor: + """Run the expert SwiGLU FFN (gate_up -> silu -> down) per local expert.""" + if self.expert_ffn_mode == "grouped_op": + # Run gate_up -> SwiGLU -> down as one TE fusible op group. + # Each GroupedLinear consumes the same per-expert split sizes. + split_sizes = torch.tensor(m_splits, dtype=torch.int32, device=tokens.device) + return self._experts_ffn_op(tokens, split_sizes, split_sizes) + elif self.expert_ffn_mode == "loop": + # Naive HF-style loop: one F.linear per expert against a slice of the + # stacked weight. Same checkpoint as the grouped path; only kernel + # dispatch differs. + # IMPORTANT: do NOT go through ``.data`` here — that detaches + # the tensor from autograd, so backward never reaches the + # expert ``nn.Parameter`` and the optimizer silently skips + # ~95% of the model. + gate_up_w = self.experts_gate_up_weight + down_w = self.experts_down_weight + outputs = [] + for i, chunk in enumerate(torch.split(tokens, m_splits, dim=0)): + if chunk.shape[0] == 0: + outputs.append(chunk) + continue + gate, up = torch.nn.functional.linear(chunk, gate_up_w[i]).chunk(2, dim=-1) + outputs.append( + torch.nn.functional.linear(torch.nn.functional.silu(gate) * up, down_w[i]) + ) + return torch.cat(outputs, dim=0) + + raise RuntimeError(f"Unknown expert_ffn_mode: {self.expert_ffn_mode!r}") + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Forward pass for the MoE block. + + Args: + hidden_states: Input tensor of shape [B, S, H] (bshd) or [T, H] (thd). + + Returns: + Output tensor of the same shape as the input. + """ + original_shape = hidden_states.shape + + # Apply multiplicative jitter noise to hidden states during training to encourage load balancing + if self.training and self.jitter_noise > 0: + hidden_states = hidden_states * torch.empty_like(hidden_states).uniform_( + 1.0 - self.jitter_noise, 1.0 + self.jitter_noise + ) + + # Flatten to [N, H] for routing + if hidden_states.dim() == 3: + hidden_states = hidden_states.reshape(-1, self.hidden_size) + + # Router: compute expert assignments + with transformer_engine.pytorch.autocast(enabled=False): + # Keep the router logits in bf16 during FP8 training + router_logits = self.gate(hidden_states) # [N, num_experts] + + # Compute the full (N, E) softmax probs once and reuse them for both + # top-k and the fused aux loss kernel. + softmax_probs = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float32) + routing_weights, selected_experts = torch.topk( + softmax_probs, self.top_k, dim=-1 + ) # [N, top_k] + # Normalize routing weights + routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True) + + # Auxiliary load-balancing loss (switch transformer style). Use TE's + # fused router kernel — fold bincount + softmax-mean + sum into one + # CUDA launch. + if self.moe_aux_loss_coeff > 0: + num_tokens = hidden_states.shape[0] + tokens_per_expert = torch.bincount( + selected_experts.reshape(-1), minlength=self.num_experts + ).to(torch.int32) + self._aux_loss = fused_moe_aux_loss( + probs=softmax_probs, + tokens_per_expert=tokens_per_expert, + total_num_tokens=num_tokens, + num_experts=self.num_experts, + topk=self.top_k, + coeff=self.moe_aux_loss_coeff, + ) + else: + self._aux_loss = torch.tensor(0.0, device=hidden_states.device) + + # Populate GroupedLinear weight attributes from stacked parameters. + self._sync_expert_views() + + if isinstance(self.dispatcher, AllToAllTokenDispatcher): + pad_to_multiple = None + if ( + self.expert_ffn_mode == "grouped_op" + and FP8GlobalStateManager.is_fp8_enabled() + and FP8GlobalStateManager.get_fp8_recipe().mxfp8() + ): + pad_to_multiple = 128 + self.dispatcher.pad_to_multiple = pad_to_multiple + + dispatch_output = self.dispatcher.dispatch(hidden_states, selected_experts, routing_weights) + + expert_input = dispatch_output.expert_input + tokens_per_expert = dispatch_output.tokens_per_expert + + expert_output = self._expert_ffn(expert_input, tokens_per_expert) + + output = self.dispatcher.combine(expert_output, dispatch_output.handle) + + return output.reshape(original_shape) + + +class TEMixtralDecoderLayer(nn.Module): + """Mixtral decoder layer using TE attention and MoE MLP.""" + + def __init__( + self, config: MixtralConfig, layer_idx: int, dispatcher: TokenDispatcher | None = None + ): + """Initialize the decoder layer.""" + super().__init__() + self.hidden_size = config.hidden_size + + device = "meta" if torch.get_default_device() == torch.device("meta") else "cuda" + + def _init_method(x): + torch.nn.init.normal_(x, mean=0.0, std=config.initializer_range) + + self.self_attention = transformer_engine.pytorch.MultiheadAttention( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_gqa_groups=config.num_key_value_heads, + bias=False, + layernorm_epsilon=config.rms_norm_eps, + attention_dropout=0, + fuse_qkv_params=True, + qkv_weight_interleaved=True, + normalization="RMSNorm", + input_layernorm=True, + qkv_format=config.attn_input_format, + attn_mask_type=config.self_attn_mask_type, + layer_number=layer_idx + 1, + params_dtype=config.dtype, + device=device, + init_method=_init_method, + output_layer_init_method=_init_method, + ) + + self.post_attention_layernorm = transformer_engine.pytorch.RMSNorm( + config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.dtype, + device=device, + ) + + self.mlp = TEMixtralSparseMoeBlock(config, dispatcher) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + rotary_pos_emb: torch.Tensor | None = None, + inference_params: InferenceParams | None = None, + **kwargs, + ) -> torch.Tensor: + """Forward pass for the decoder layer.""" + # Self attention with fused input layernorm + attn_output = self.self_attention( + hidden_states, + attention_mask=attention_mask, + rotary_pos_emb=rotary_pos_emb, + inference_params=inference_params, + cu_seqlens_q=kwargs.get("cu_seqlens_q", None), + cu_seqlens_kv=kwargs.get("cu_seqlens_kv", None), + cu_seqlens_q_padded=kwargs.get("cu_seqlens_q_padded", None), + cu_seqlens_kv_padded=kwargs.get("cu_seqlens_kv_padded", None), + max_seqlen_q=kwargs.get("max_seqlen_q", None), + max_seqlen_kv=kwargs.get("max_seqlen_kv", None), + pad_between_seqs=kwargs.get("pad_between_seqs", None), + ) + + # Residual connection + hidden_states = hidden_states + attn_output + + # Post-attention layernorm + MoE MLP + residual + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + + return hidden_states + + +class TEMixtralModel(TEMixtralPreTrainedModel): + """Mixtral model implemented in Transformer Engine.""" + + def __init__( + self, + config: MixtralConfig, + fp8_recipe: transformer_engine.common.recipe.Recipe | None = None, + fp4_recipe: transformer_engine.common.recipe.Recipe | None = None, + dispatcher: TokenDispatcher | None = None, + ): + """Initialize the TEMixtral model. + + Args: + config: The configuration of the model. + fp8_recipe: The FP8 recipe for the model. + fp4_recipe: The FP4 recipe for the model. + dispatcher: The token dispatcher for the model. If None, the default AllToAllTokenDispatcher will be used. + """ + super().__init__(config) + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self._fp8_recipe: transformer_engine.common.recipe.Recipe | None = fp8_recipe + self._fp4_recipe: transformer_engine.common.recipe.Recipe | None = fp4_recipe + + if fp8_recipe is not None and self.config.layer_precision is None: + if fp4_recipe is not None: + raise RuntimeError( + "Both FP8 and FP4 recipes provided, but no layer precision provided." + ) + + warnings.warn( + "No layer precision provided, using FP8 recipe for all layers.", UserWarning + ) + self.config.layer_precision = ["fp8"] * self.config.num_hidden_layers + + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx, dtype=config.dtype + ) + + layers: list[TEMixtralDecoderLayer] = [] + for layer_idx in range(config.num_hidden_layers): + with self.get_autocast_context(layer_idx, init=True): + layers += [TEMixtralDecoderLayer(config, layer_idx, dispatcher)] + + self.layers = nn.ModuleList(layers) + + self.norm = transformer_engine.pytorch.RMSNorm( + config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.dtype, + device="meta" if torch.get_default_device() == torch.device("meta") else "cuda", + ) + + self.rotary_emb = RotaryPositionEmbedding(config.hidden_size // config.num_attention_heads) + self.rotary_emb.inv_freq = LlamaRotaryEmbedding(config=config).inv_freq + + self.gradient_checkpointing = False + + self.post_init() + + def set_ep_groups(self, ep_group: dist.ProcessGroup) -> None: + """Propagate an expert-parallel process group to every MoE block. + + Args: + ep_group: The EP process group to set on each ``TEMixtralSparseMoeBlock``. + """ + for layer in self.layers: + layer.mlp.set_ep_group(ep_group) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + past_key_values: InferenceParams | None = None, + inputs_embeds: torch.Tensor | None = None, + use_cache: bool | None = None, + **kwargs: Unpack[TransformersKwargs], + ) -> BaseModelOutputWithPast: + """Forward pass for the TEMixtral model.""" + all_hidden_states = [] + output_hidden_states = kwargs.get("output_hidden_states", False) + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if inputs_embeds is None: + inputs_embeds: torch.Tensor = self.embed_tokens(input_ids) + + hidden_states = inputs_embeds + + # TE-specific input handling + has_thd_input = [ + x in kwargs for x in ["cu_seq_lens_q", "cu_seq_lens_k", "max_length_q", "max_length_k"] + ] + decode_without_mask = ( + isinstance(past_key_values, InferenceParams) + and hidden_states.dim() == 3 + and hidden_states.size(1) == 1 + ) + should_pack_inputs = ( + not any(has_thd_input) + and self.config.attn_input_format == "thd" + and not decode_without_mask + ) + + if should_pack_inputs: + assert ( + attention_mask is not None + ), "Attention mask is required when packing BSHD inputs." + batch_size = hidden_states.size(0) + padded_seq_len = hidden_states.size(1) + hidden_states, indices, cu_seqlens, max_seqlen, _ = _unpad_input( + hidden_states, attention_mask + ) + + # MXFP8 block scaling requires the token dim divisible by 32. + # After THD unpadding the total token count is data-dependent. + thd_orig_tokens = hidden_states.shape[0] + thd_remainder = thd_orig_tokens % 32 + if thd_remainder != 0: + thd_pad = 32 - thd_remainder + hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, thd_pad)) + # Extend cu_seqlens: add padding tokens to the last sequence + cu_seqlens = cu_seqlens.clone() + cu_seqlens[-1] = cu_seqlens[-1] + thd_pad + max_seqlen = max_seqlen + thd_pad + + kwargs["cu_seq_lens_q"] = kwargs["cu_seq_lens_k"] = cu_seqlens + kwargs["max_length_q"] = kwargs["max_length_k"] = max_seqlen + + if ( + self.config.attn_input_format == "thd" + and hidden_states.dim() == 3 + and hidden_states.size(0) == 1 + ): + hidden_states = hidden_states.squeeze(0) + + if ( + self.config.attn_input_format == "bshd" + and attention_mask is not None + and attention_mask.dim() == 2 + ): + # Convert HF mask (1=attend, 0=pad) to TE boolean mask (True=masked, False=attend) + attention_mask = ~attention_mask[:, None, None, :].bool() + + if isinstance(past_key_values, InferenceParams): + _ref = input_ids if input_ids is not None else inputs_embeds + lengths = ( + attention_mask.sum(dim=1).tolist() + if attention_mask is not None and attention_mask.shape[:2] == _ref.shape[:2] + else [1] * _ref.shape[0] + ) + past_key_values.pre_step(OrderedDict(zip(list(range(len(lengths))), lengths))) + + with torch.autocast(device_type="cuda", enabled=False): + te_rope_emb = self.rotary_emb(max_seq_len=self.config.max_position_embeddings) + + with self.get_autocast_context(None, outer=True): + for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): + if output_hidden_states: + all_hidden_states = (*all_hidden_states, hidden_states) + + with self.get_autocast_context(layer_idx): + hidden_states = decoder_layer( + hidden_states, + attention_mask=( + None if self.config.attn_input_format == "thd" else attention_mask + ), + rotary_pos_emb=te_rope_emb, + inference_params=past_key_values, + cu_seqlens_q=kwargs.get("cu_seq_lens_q", None), + cu_seqlens_kv=kwargs.get("cu_seq_lens_k", None), + cu_seqlens_q_padded=kwargs.get("cu_seq_lens_q_padded", None), + cu_seqlens_kv_padded=kwargs.get("cu_seq_lens_k_padded", None), + max_seqlen_q=kwargs.get("max_length_q", None), + max_seqlen_kv=kwargs.get("max_length_k", None), + pad_between_seqs=kwargs.get("pad_between_seqs", None), + ) + + hidden_states = self.norm(hidden_states) + + if output_hidden_states: + all_hidden_states = (*all_hidden_states, hidden_states) + + if should_pack_inputs: + if thd_remainder != 0: + hidden_states = hidden_states[:thd_orig_tokens] + hidden_states = _pad_input(hidden_states, indices, batch_size, padded_seq_len) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=all_hidden_states if output_hidden_states else None, + ) + + def get_autocast_context( + self, layer_number: int | None, init: bool = False, outer: bool = False + ) -> ContextManager: + """Return the appropriate TE autocast context manager for a given layer. + + This function handles both the quantized_model_init during layer creation and the te.autocast() during layer + forward pass. + + Args: + layer_number: The 0-indexed layer number. + init: Whether to return a `quantized_model_init` context for layer initialization. + outer: Whether to return a global te.autocast() context to wrap the entire model stack. + """ + if self.config.layer_precision is None: + return nullcontext() + + if outer: + if "fp8" not in self.config.layer_precision: + return nullcontext() + if self._fp8_recipe is None: + warnings.warn("No FP8 recipe provided, using default recipe.", UserWarning) + return transformer_engine.pytorch.autocast(enabled=True, recipe=self._fp8_recipe) + + precision = self.config.layer_precision[layer_number] + recipe = {"fp8": self._fp8_recipe, "fp4": self._fp4_recipe}.get(precision) + + if init and self.config.use_quantized_model_init: + if precision in ("fp8", "fp4"): + return transformer_engine.pytorch.quantized_model_init(recipe=recipe) + return nullcontext() + + if precision == "fp8": + if recipe is None: + warnings.warn("No FP8 recipe provided, using default recipe.", UserWarning) + return transformer_engine.pytorch.autocast(enabled=True, recipe=recipe) + if precision == "fp4": + if recipe is None: + raise RuntimeError("No FP4 recipe provided, but layer precision is set to FP4.") + return transformer_engine.pytorch.autocast(enabled=True, recipe=recipe) + return transformer_engine.pytorch.autocast(enabled=False) + + +class TEMixtralForCausalLM(TEMixtralPreTrainedModel, transformers.GenerationMixin): + """Mixtral model with causal language head.""" + + _tied_weights_keys: ClassVar[list[str]] = [] + + def __init__( + self, + config, + fp8_recipe: transformer_engine.common.recipe.Recipe | None = None, + fp4_recipe: transformer_engine.common.recipe.Recipe | None = None, + dispatcher: TokenDispatcher | None = None, + ): + """Initialize the TEMixtralForCausalLM model. + + Args: + config: The configuration of the model. + fp8_recipe: The FP8 recipe for the model. + fp4_recipe: The FP4 recipe for the model. + dispatcher: The token dispatcher for expert parallelism. If None, the default + AllToAllTokenDispatcher will be used. + """ + super().__init__(config) + self.model = TEMixtralModel( + config, fp8_recipe=fp8_recipe, fp4_recipe=fp4_recipe, dispatcher=dispatcher + ) + self.vocab_size = config.vocab_size + + with transformer_engine.pytorch.quantized_model_init(enabled=False): + self.lm_head = transformer_engine.pytorch.Linear( + config.hidden_size, + config.vocab_size, + bias=False, + params_dtype=config.dtype, + device="meta" if torch.get_default_device() == torch.device("meta") else "cuda", + init_method=lambda x: torch.nn.init.normal_( + x, mean=0.0, std=config.initializer_range + ), + ) + + self.post_init() + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + past_key_values: tuple[tuple[torch.Tensor, ...], ...] | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + shift_labels: torch.Tensor | None = None, + use_cache: bool | None = None, + cache_position: torch.Tensor | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Unpack[TransformersKwargs], + ) -> CausalLMOutputWithPast: + """Forward pass for the TEMixtralForCausalLM model.""" + outputs: BaseModelOutputWithPast = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + cache_position=cache_position, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + slice_indices = ( + slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + ) + + with transformer_engine.pytorch.autocast(enabled=False): + if hidden_states.ndim == 3: + logits = self.lm_head(hidden_states[:, slice_indices, :]) + else: + logits = self.lm_head(hidden_states[slice_indices, :]) + + loss = None + if labels is not None or shift_labels is not None: + loss = self.loss_function( + logits=logits, + labels=labels, + shift_labels=shift_labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) + + # Collect auxiliary load-balancing loss from all MoE layers + if self.config.moe_aux_loss_coeff > 0 and loss is not None: + aux_loss = sum(layer.mlp._aux_loss for layer in self.model.layers) + loss = loss + aux_loss + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +# Required for torch.compile'd functions below (_pad_input, _unpad_input, _build_expert_sort_indices) +# that use data-dependent scalar values (e.g., max_seqlen_in_batch.item()) or produce tensors +# whose shape depends on input data (e.g., repeat_interleave with tensor counts). +# These must be set at module level because torch.compile traces lazily on first call, +# so a scoped setting would not be active at trace time. +torch._dynamo.config.capture_scalar_outputs = True +torch._dynamo.config.capture_dynamic_output_shape_ops = True + + +@torch.compile +def _pad_input(hidden_states, indices, batch, seqlen): + """Convert a THD tensor to a BSHD equivalent tensor. + + Adapted from huggingface/transformers/modeling_flash_attention_utils.py + """ + dim = hidden_states.shape[1:] + output = torch.zeros( + (batch * seqlen), *dim, device=hidden_states.device, dtype=hidden_states.dtype + ) + output[indices] = hidden_states + return output.view(batch, seqlen, *dim) + + +@torch.compile +def _unpad_input(hidden_states, attention_mask, unused_mask=None): + """Convert a BSHD tensor to a THD equivalent tensor. + + Adapted from huggingface/transformers/modeling_flash_attention_utils.py + """ + batch_size = hidden_states.size(0) + seq_length = hidden_states.size(1) + + if attention_mask.shape[1] != seq_length: + return ( + hidden_states.squeeze(1), + torch.arange(batch_size, dtype=torch.int64, device=hidden_states.device), + torch.arange(batch_size + 1, dtype=torch.int32, device=hidden_states.device), + 1, + 1, + ) + + all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask + seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32) + used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = torch.nn.functional.pad( + torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0) + ) + + return ( + hidden_states.reshape(-1, *hidden_states.shape[2:])[indices], + indices, + cu_seqlens, + max_seqlen_in_batch, + used_seqlens_in_batch, + ) + + +class HFInferenceParams(InferenceParams): + """Extension of the InferenceParams class to support HF generate() and beam search.""" + + def get_seq_length(self, layer_idx: int = 0) -> int: + """Return the current cached sequence length. + + Required by HuggingFace transformers generate() to determine how many + tokens have already been cached. + """ + if not self.sequences: + return 0 + return max(self.sequences.values()) + + def reorder_cache(self, beam_idx: torch.LongTensor): + """Reorder the cache based on the beam indices.""" + if isinstance(self.cache_manager, PagedKVCacheManager): + raise NotImplementedError("Beam search is not supported for paged cache manager.") + for layer_number, (key_cache, value_cache) in self.cache_manager.cache.items(): + updated_key_cache = key_cache.index_select(0, beam_idx) + updated_value_cache = value_cache.index_select(0, beam_idx) + self.cache_manager.cache[layer_number] = (updated_key_cache, updated_value_cache) + + +@torch.compile(fullgraph=True) +def _build_expert_sort_indices(recv_counts: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Build sort and unsort index tensors for reordering received tokens by local expert. + + After all-to-all, tokens arrive grouped by source rank: + ``[src0_exp0..src0_expL, src1_exp0..src1_expL, ...]``. ``GroupedLinear`` expects them + grouped by expert: ``[all_exp0, all_exp1, ...]``. + + Uses only vectorized tensor operations (no ``.item()`` calls or Python-level loops) + so that it is compatible with ``torch.compile(fullgraph=True)``. + + Args: + recv_counts: Integer tensor of shape ``[ep_size, num_local_experts]`` giving the + number of tokens received from each source rank for each local expert. + + Returns: + A ``(sort_indices, unsort_indices)`` pair of 1-D ``int64`` tensors that can be + used to reorder and restore the token dimension. + """ + ep_size, num_local_experts = recv_counts.shape + device = recv_counts.device + num_blocks = ep_size * num_local_experts + + # Source-grouped (row-major) block offsets: [s0e0, s0e1, ..., s1e0, s1e1, ...] + counts_src = recv_counts.reshape(-1).long() + offsets_src = torch.zeros(num_blocks, dtype=torch.long, device=device) + offsets_src[1:] = counts_src[:-1].cumsum(0) + + # Expert-grouped (column-major) block offsets: [e0s0, e0s1, ..., e1s0, e1s1, ...] + counts_exp = recv_counts.t().contiguous().reshape(-1).long() + offsets_exp = torch.zeros(num_blocks, dtype=torch.long, device=device) + offsets_exp[1:] = counts_exp[:-1].cumsum(0) + + total = counts_src.sum() + + # Mapping from source block index (s * L + e) to expert block index (e * S + s) + s_idx = torch.arange(ep_size, device=device).unsqueeze(1).expand(ep_size, num_local_experts) + e_idx = ( + torch.arange(num_local_experts, device=device) + .unsqueeze(0) + .expand(ep_size, num_local_experts) + ) + src_to_exp = (e_idx * ep_size + s_idx).reshape(-1) + + # Per-block positional shift from source layout to expert layout + shifts = offsets_exp[src_to_exp] - offsets_src + + # Expand per-block shifts to per-token + token_shifts = shifts.repeat_interleave(counts_src) + + # Map each source-grouped position to its expert-grouped destination + src_positions = torch.arange(total, device=device) + dst_positions = src_positions + token_shifts + + # sort_indices[exp_pos] = src_pos (gathers source tokens into expert order) + sort_indices = torch.empty(total, dtype=torch.long, device=device) + sort_indices[dst_positions] = src_positions + + # unsort_indices: inverse permutation (restores expert-ordered output to source order) + unsort_indices = torch.empty_like(sort_indices) + unsort_indices[sort_indices] = torch.arange(total, device=device) + + return sort_indices, unsort_indices + + +@dataclass +class _AllToAllHandle: + """Opaque handle for AllToAllTokenDispatcher, storing state between dispatch and combine.""" + + row_id_map: torch.Tensor + routing_weights: torch.Tensor + restore_shape: torch.Size + map_type: str = "index" + pad_offsets: torch.Tensor | None = None + unsort_indices: torch.Tensor | None = None + input_split_sizes: list[int] | None = None + output_split_sizes: list[int] | None = None + + +class _DifferentiableAllToAll(torch.autograd.Function): + """Differentiable wrapper around dist.all_to_all_single. + + The forward pass performs the standard all-to-all communication. + The backward pass reverses the communication direction (swapping + input/output split sizes) so that gradients flow correctly. + """ + + @staticmethod + def forward( + ctx, + input: torch.Tensor, + output_split_sizes: list[int], + input_split_sizes: list[int], + group: dist.ProcessGroup, + ) -> torch.Tensor: + """Perform all-to-all forward and save sizes for backward.""" + ctx.input_split_sizes = input_split_sizes + ctx.output_split_sizes = output_split_sizes + ctx.group = group + output = torch.empty( + sum(output_split_sizes), + input.shape[1], + device=input.device, + dtype=input.dtype, + ) + dist.all_to_all_single( + output, input.contiguous(), output_split_sizes, input_split_sizes, group=group + ) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor) -> tuple[torch.Tensor, None, None, None]: + """Reverse all-to-all: swap input and output split sizes.""" + grad_input = torch.empty( + sum(ctx.input_split_sizes), + grad_output.shape[1], + device=grad_output.device, + dtype=grad_output.dtype, + ) + dist.all_to_all_single( + grad_input, + grad_output.contiguous(), + ctx.input_split_sizes, + ctx.output_split_sizes, + group=ctx.group, + ) + return grad_input, None, None, None + + +class AllToAllTokenDispatcher: + """TokenDispatcher using NCCL all-to-all for expert-parallel communication. + + Handles both EP=1 (no communication, just permute/unpermute) and EP>1 + (all-to-all token exchange between ranks) cases transparently. + + Args: + num_experts: Total number of experts (global). + num_local_experts: Number of experts on this rank. + hidden_size: Hidden dimension size. + ep_size: Expert parallel world size. + """ + + def __init__(self, num_experts: int, num_local_experts: int, hidden_size: int, ep_size: int): + """Initialize the AllToAllTokenDispatcher.""" + self.num_experts = num_experts + self.num_local_experts = num_local_experts + self.hidden_size = hidden_size + self.ep_size = ep_size + self._ep_group: dist.ProcessGroup | None = None + self.pad_to_multiple: int | None = None + + def set_ep_group(self, ep_group: dist.ProcessGroup) -> None: + """Set the expert-parallel process group for all-to-all communication.""" + self._ep_group = ep_group + + def dispatch( + self, + hidden_states: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + ) -> DispatchOutput: + """Dispatch tokens to their assigned experts via permute and optional all-to-all. + + Args: + hidden_states: Flattened input tensor of shape ``[N, H]``. + selected_experts: Expert assignments, shape ``[N, top_k]``, int. + routing_weights: Normalized routing probabilities, shape ``[N, top_k]``, float32. + + Returns: + DispatchOutput with expert-sorted tokens, per-expert counts, and an opaque handle. + """ + # Compute m_splits: number of tokens per expert + m_splits_tensor = torch.bincount( + selected_experts.reshape(-1), minlength=self.num_experts + ).int() + + pad_offsets = None + if self.pad_to_multiple is not None: + routing_map = torch.zeros( + hidden_states.shape[0], + self.num_experts, + dtype=torch.bool, + device=hidden_states.device, + ) + routing_map.scatter_(1, selected_experts, True) + routing_probs = torch.zeros( + hidden_states.shape[0], + self.num_experts, + dtype=routing_weights.dtype, + device=hidden_states.device, + ) + routing_probs.scatter_(1, selected_experts, routing_weights) + ( + permuted_hidden, + _, + row_id_map, + pad_offsets, + m_splits_tensor, + ) = transformer_engine.pytorch.moe_permute_and_pad_with_probs( + hidden_states, + routing_probs, + routing_map, + m_splits_tensor, + self.pad_to_multiple, + ) + m_splits_tensor = m_splits_tensor.int() + routing_weights_for_unpermute = routing_probs + map_type = "mask" + else: + # Permute tokens by expert using TE moe_permute. + permuted_hidden, row_id_map = transformer_engine.pytorch.moe_permute( + hidden_states, + selected_experts.to(torch.int32), + num_out_tokens=selected_experts.numel(), + map_type="index", + ) + routing_weights_for_unpermute = routing_weights + map_type = "index" + + if self._ep_group is not None: + ep_group = self._ep_group + + # Token counts per expert, reshaped to [ep_size, num_local_experts] + send_counts = m_splits_tensor.reshape(self.ep_size, self.num_local_experts) + + # Exchange per-expert token counts between EP ranks + recv_counts = torch.empty_like(send_counts) + dist.all_to_all_single(recv_counts.flatten(), send_counts.flatten(), group=ep_group) + + # Derive split sizes for the token all-to-all + input_split_sizes = send_counts.sum(dim=1).tolist() + output_split_sizes = recv_counts.sum(dim=1).tolist() + local_m_splits = recv_counts.sum(dim=0).int().tolist() + + # Dispatch tokens to expert-owning ranks (differentiable) + recv_tokens = _DifferentiableAllToAll.apply( + permuted_hidden, output_split_sizes, input_split_sizes, ep_group + ) + + # Sort received tokens by local expert index. + # After all_to_all layout is [src0_exp0..src0_expL, src1_exp0..src1_expL, ...]. + # GroupedLinear needs [all_exp0, all_exp1, ...]. + sort_indices, unsort_indices = _build_expert_sort_indices(recv_counts) + + handle = _AllToAllHandle( + row_id_map=row_id_map, + routing_weights=routing_weights_for_unpermute, + restore_shape=hidden_states.shape, + map_type=map_type, + pad_offsets=pad_offsets, + unsort_indices=unsort_indices, + input_split_sizes=input_split_sizes, + output_split_sizes=output_split_sizes, + ) + return DispatchOutput( + expert_input=recv_tokens[sort_indices], + tokens_per_expert=local_m_splits, + handle=handle, + ) + + handle = _AllToAllHandle( + row_id_map=row_id_map, + routing_weights=routing_weights_for_unpermute, + restore_shape=hidden_states.shape, + map_type=map_type, + pad_offsets=pad_offsets, + ) + return DispatchOutput( + expert_input=permuted_hidden, + tokens_per_expert=m_splits_tensor.tolist(), + handle=handle, + ) + + def combine(self, expert_output: torch.Tensor, handle: _AllToAllHandle) -> torch.Tensor: + """Combine expert outputs back to the original token order. + + Args: + expert_output: Expert output tensor of shape ``[total_recv_tokens, H]``. + handle: Handle from ``dispatch()`` containing state for the reverse operation. + + Returns: + Combined output tensor of shape ``[N, H]`` with routing weights applied. + """ + if self._ep_group is not None: + assert handle.unsort_indices is not None + # Unsort back to source-rank-grouped order and reverse all_to_all (differentiable) + combined = _DifferentiableAllToAll.apply( + expert_output[handle.unsort_indices], + handle.input_split_sizes, + handle.output_split_sizes, + self._ep_group, + ) + else: + combined = expert_output + + # Unpermute and combine with routing weights (keep probs in float32 for numerical stability) + return transformer_engine.pytorch.moe_unpermute( + combined, + handle.row_id_map, + merging_probs=handle.routing_weights, + restore_shape=handle.restore_shape, + map_type=handle.map_type, + pad_offsets=handle.pad_offsets, + ) diff --git a/docs/examples/te_mixtral/te_mixtral_mxfp8.py b/docs/examples/te_mixtral/te_mixtral_mxfp8.py new file mode 100644 index 000000000..cca7d2308 --- /dev/null +++ b/docs/examples/te_mixtral/te_mixtral_mxfp8.py @@ -0,0 +1,568 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""TE-native MXFP8 Mixtral model (improvement 3). + +MoE FFN is a TE ``Sequential`` of three fusible ops — ``GroupedLinear`` +(gate_up), ``ScaledSwiGLU(glu_interleave_size=32)``, ``GroupedLinear`` +(down) — that the OperationFuser collapses into the fused +``ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8`` and backward kernels under +MXFP8. HF gate (``w1``) and up (``w3``) weights are row-interleaved in +blocks of 32 to match the GLU interleaved layout that fused kernel reads. + +The fused kernel is enabled by ``utils._enable_fused_mxfp8_grouped_mlp()`` +(sets ``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1`` and patches the SM-version / +cudnn-frontend signature checks). Requires +``nvidia-cudnn-frontend >= 1.23.0`` and SM>=10 (Blackwell B100/B200/B300+). +""" + +from __future__ import annotations + +import logging +from collections import OrderedDict +from contextlib import nullcontext +from typing import Any, ClassVar, ContextManager + +import torch +import torch.distributed as dist +import torch.nn as nn + +import transformer_engine.common.recipe as te_recipe +import transformer_engine.pytorch as te +from transformer_engine.pytorch.attention.inference import InferenceParams +from transformer_engine.pytorch.attention.rope import RotaryPositionEmbedding +from transformer_engine.pytorch.ops import ( + GroupedLinear as TEOpsGroupedLinear, + ScaledSwiGLU, + Sequential as TEOpsSequential, +) +from transformer_engine.pytorch.router import fused_moe_aux_loss +from transformers import MixtralConfig, PreTrainedModel +from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast +from transformers.models.llama.modeling_llama import LlamaRotaryEmbedding + +from te_moe_dispatch import AllToAllTokenDispatcher +from te_mixtral import ( + _pad_input, + _unpad_input, +) + +logger = logging.getLogger(__name__) + + +# HF->TE checkpoint mapping is shared with the BF16 path in te_mixtral.py. +# ``GLU_INTERLEAVE_SIZE`` is the gate/up interleave block (32); the fused +# MXFP8 forward op only fires when ``ScaledSwiGLU`` is configured with it. +from hf_to_te_weights import ( + GLU_INTERLEAVE_SIZE, + replace_params_mxfp8 as replace_params, +) + + +class TEMixtralMXFP8Config(MixtralConfig): + """Improvement-3 config. Same surface as :class:`te_mixtral.TEMixtralConfig` but + with the FFN mode locked to ``grouped_op`` + MXFP8.""" + + attn_input_format: str = "thd" + self_attn_mask_type: str = "padding_causal" + expert_parallel_size: int = 1 + moe_aux_loss_coeff: float = 0.0 + + def __init__(self, **kwargs): + super().__init__(**kwargs) + if self.num_local_experts % self.expert_parallel_size != 0: + raise ValueError( + f"num_local_experts ({self.num_local_experts}) must be divisible by " + f"expert_parallel_size ({self.expert_parallel_size})" + ) + + +class TEMixtralMXFP8PreTrainedModel(PreTrainedModel): + """HF integration boilerplate for the improvement-3 model.""" + + config_class = TEMixtralMXFP8Config + base_model_prefix = "model" + _no_split_modules = ("TEMixtralMXFP8DecoderLayer",) + _skip_keys_device_placement = ("past_key_values",) + _do_not_quantize = ("lm_head", "model.layers.*.mlp.gate") + + def _init_weights(self, module): + if module.__module__.startswith("transformer_engine.pytorch"): + return + super()._init_weights(module) + + def state_dict(self, *args, **kwargs): + sd = super().state_dict(*args, **kwargs) + return {k: v for k, v in sd.items() if not k.endswith("_extra_state")} + + +class TEMixtralMXFP8SparseMoeBlock(nn.Module): + """MoE block: router + EP dispatcher + fused MXFP8 grouped MLP.""" + + def __init__( + self, + config: TEMixtralMXFP8Config, + dispatcher: AllToAllTokenDispatcher | None = None, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + self.intermediate_size = config.intermediate_size + self.num_experts = config.num_local_experts + self.top_k = config.num_experts_per_tok + self.jitter_noise = config.router_jitter_noise + + self.ep_size = getattr(config, "expert_parallel_size", 1) + self.num_local_experts = self.num_experts // self.ep_size + self.moe_aux_loss_coeff = getattr(config, "moe_aux_loss_coeff", 0.0) + self._aux_loss: torch.Tensor = torch.tensor(0.0) + + if self.intermediate_size % GLU_INTERLEAVE_SIZE != 0: + raise ValueError( + f"intermediate_size ({self.intermediate_size}) must be divisible by " + f"GLU_INTERLEAVE_SIZE ({GLU_INTERLEAVE_SIZE})" + ) + + self.dispatcher = dispatcher or AllToAllTokenDispatcher( + num_experts=self.num_experts, + num_local_experts=self.num_local_experts, + hidden_size=self.hidden_size, + ep_size=self.ep_size, + ) + + device = "meta" if torch.get_default_device() == torch.device("meta") else "cuda" + + def _init_method(x: torch.Tensor) -> None: + torch.nn.init.normal_(x, mean=0.0, std=config.initializer_range) + + with te.quantized_model_init(enabled=False): + self.gate = te.Linear( + self.hidden_size, + self.num_experts, + bias=False, + device=device, + params_dtype=config.dtype, + init_method=_init_method, + ) + + self.experts_gate_up = TEOpsGroupedLinear( + num_groups=self.num_local_experts, + in_features=self.hidden_size, + out_features=2 * self.intermediate_size, + bias=False, + dtype=config.dtype, + device=device, + ) + self.experts_swiglu = ScaledSwiGLU(glu_interleave_size=GLU_INTERLEAVE_SIZE) + self.experts_down = TEOpsGroupedLinear( + num_groups=self.num_local_experts, + in_features=self.intermediate_size, + out_features=self.hidden_size, + bias=False, + dtype=config.dtype, + device=device, + ) + # Wrap as TE Sequential to enable forward/backward op fusion + # (ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 / dswiglu). + object.__setattr__( + self, + "_experts_ffn_op", + TEOpsSequential(self.experts_gate_up, self.experts_swiglu, self.experts_down), + ) + + def set_ep_group(self, ep_group: dist.ProcessGroup) -> None: + """Set the EP communication group on the dispatcher. + + Each EP rank owns its local slice of expert weights as ordinary + Parameters (``weight0..weight{N-1}``) because per-expert parameters + are never replicated across the EP group. + """ + self.dispatcher.set_ep_group(ep_group) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + original_shape = hidden_states.shape + + if self.training and self.jitter_noise > 0: + hidden_states = hidden_states * torch.empty_like(hidden_states).uniform_( + 1.0 - self.jitter_noise, 1.0 + self.jitter_noise + ) + + if hidden_states.dim() == 3: + hidden_states = hidden_states.reshape(-1, self.hidden_size) + + with te.autocast(enabled=False): + router_logits = self.gate(hidden_states) # [N, E] + + # Top-k routing weights, two algebraically equivalent forms. + # Old:: + # + # probs = softmax(logits) # (N, E) + # weights, idx = topk(probs, k) + # weights = weights / weights.sum(-1, keepdim=True) + # + # New (used here):: + # + # topk_logits, idx = topk(logits, k) + # weights = softmax(topk_logits) # softmax over (N, k) + topk_logits, selected_experts = torch.topk(router_logits, self.top_k, dim=-1) + routing_weights = torch.nn.functional.softmax(topk_logits, dim=-1, dtype=torch.float32) + + # Bincount once, in the MoE block. ``AllToAllTokenDispatcher`` + # takes this as a required argument, so the dispatcher never + # bincounts again. + tokens_per_expert = torch.bincount( + selected_experts.reshape(-1), minlength=self.num_experts + ).to(torch.int32) + + if self.moe_aux_loss_coeff > 0: + num_tokens = hidden_states.shape[0] + softmax_probs = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float32) + self._aux_loss = fused_moe_aux_loss( + probs=softmax_probs, + tokens_per_expert=tokens_per_expert, + total_num_tokens=num_tokens, + num_experts=self.num_experts, + topk=self.top_k, + coeff=self.moe_aux_loss_coeff, + ) + else: + self._aux_loss = torch.tensor(0.0, device=hidden_states.device) + + dispatch_out = self.dispatcher.dispatch( + hidden_states, + selected_experts, + routing_weights, + tokens_per_expert, + ) + expert_input = dispatch_out.expert_input + expert_probs = dispatch_out.expert_probs + split_sizes = torch.tensor( + dispatch_out.tokens_per_expert, dtype=torch.int32, device=expert_input.device + ) + + # Fused gate_up -> ScaledSwiGLU(probs) -> down. + expert_output = self._experts_ffn_op(expert_input, split_sizes, expert_probs, split_sizes) + + output = self.dispatcher.combine(expert_output, dispatch_out.handle) + return output.reshape(original_shape) + + +class TEMixtralMXFP8DecoderLayer(nn.Module): + """Self-attention + improvement-3 MoE block.""" + + def __init__( + self, + config: TEMixtralMXFP8Config, + layer_idx: int, + dispatcher: AllToAllTokenDispatcher | None = None, + ) -> None: + super().__init__() + self.hidden_size = config.hidden_size + device = "meta" if torch.get_default_device() == torch.device("meta") else "cuda" + + def _init_method(x: torch.Tensor) -> None: + torch.nn.init.normal_(x, mean=0.0, std=config.initializer_range) + + self.self_attention = te.MultiheadAttention( + hidden_size=config.hidden_size, + num_attention_heads=config.num_attention_heads, + num_gqa_groups=config.num_key_value_heads, + bias=False, + layernorm_epsilon=config.rms_norm_eps, + attention_dropout=0, + fuse_qkv_params=True, + qkv_weight_interleaved=True, + normalization="RMSNorm", + input_layernorm=True, + qkv_format=config.attn_input_format, + attn_mask_type=config.self_attn_mask_type, + layer_number=layer_idx + 1, + params_dtype=config.dtype, + device=device, + init_method=_init_method, + output_layer_init_method=_init_method, + ) + self.post_attention_layernorm = te.RMSNorm( + config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.dtype, + device=device, + ) + self.mlp = TEMixtralMXFP8SparseMoeBlock(config, dispatcher) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + rotary_pos_emb: torch.Tensor | None = None, + inference_params: InferenceParams | None = None, + **kwargs: Any, + ) -> torch.Tensor: + attn_output = self.self_attention( + hidden_states, + attention_mask=attention_mask, + rotary_pos_emb=rotary_pos_emb, + inference_params=inference_params, + cu_seqlens_q=kwargs.get("cu_seqlens_q", None), + cu_seqlens_kv=kwargs.get("cu_seqlens_kv", None), + cu_seqlens_q_padded=kwargs.get("cu_seqlens_q_padded", None), + cu_seqlens_kv_padded=kwargs.get("cu_seqlens_kv_padded", None), + max_seqlen_q=kwargs.get("max_seqlen_q", None), + max_seqlen_kv=kwargs.get("max_seqlen_kv", None), + pad_between_seqs=kwargs.get("pad_between_seqs", None), + ) + hidden_states = hidden_states + attn_output + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + return residual + hidden_states + + +class TEMixtralMXFP8Model(TEMixtralMXFP8PreTrainedModel): + """Embedding + N decoder layers + RMSNorm. THD-packed under MXFP8.""" + + def __init__( + self, + config: TEMixtralMXFP8Config, + fp8_recipe: te_recipe.Recipe | None = None, + dispatcher: AllToAllTokenDispatcher | None = None, + ) -> None: + super().__init__(config) + self.config = config + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + self._fp8_recipe = fp8_recipe + + self.embed_tokens = nn.Embedding( + config.vocab_size, config.hidden_size, self.padding_idx, dtype=config.dtype + ) + + layers: list[TEMixtralMXFP8DecoderLayer] = [ + TEMixtralMXFP8DecoderLayer(config, i, dispatcher) + for i in range(config.num_hidden_layers) + ] + self.layers = nn.ModuleList(layers) + + self.norm = te.RMSNorm( + config.hidden_size, + eps=config.rms_norm_eps, + dtype=config.dtype, + device="meta" if torch.get_default_device() == torch.device("meta") else "cuda", + ) + + self.rotary_emb = RotaryPositionEmbedding(config.hidden_size // config.num_attention_heads) + self.rotary_emb.inv_freq = LlamaRotaryEmbedding(config=config).inv_freq + + self.gradient_checkpointing = False + self.post_init() + + def set_ep_groups(self, ep_group: dist.ProcessGroup) -> None: + for layer in self.layers: + layer.mlp.set_ep_group(ep_group) + + def _outer_autocast(self) -> ContextManager: + if self._fp8_recipe is None: + return nullcontext() + return te.autocast(enabled=True, recipe=self._fp8_recipe) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + past_key_values: InferenceParams | None = None, + inputs_embeds: torch.Tensor | None = None, + use_cache: bool | None = None, + **kwargs: Any, + ) -> BaseModelOutputWithPast: + del position_ids, use_cache # not used in this minimal forward + + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("Specify exactly one of input_ids or inputs_embeds") + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + hidden_states = inputs_embeds + + has_thd_input = [ + x in kwargs for x in ("cu_seq_lens_q", "cu_seq_lens_k", "max_length_q", "max_length_k") + ] + decode_without_mask = ( + isinstance(past_key_values, InferenceParams) + and hidden_states.dim() == 3 + and hidden_states.size(1) == 1 + ) + should_pack_inputs = ( + not any(has_thd_input) + and self.config.attn_input_format == "thd" + and not decode_without_mask + ) + + thd_remainder = 0 + thd_orig_tokens = 0 + indices = None + batch_size = 0 + padded_seq_len = 0 + if should_pack_inputs: + assert attention_mask is not None, "attention_mask required when packing BSHD." + batch_size = hidden_states.size(0) + padded_seq_len = hidden_states.size(1) + hidden_states, indices, cu_seqlens, max_seqlen, _ = _unpad_input( + hidden_states, attention_mask + ) + + # MXFP8 requires total tokens divisible by 32; pad the last seq. + thd_orig_tokens = hidden_states.shape[0] + thd_remainder = thd_orig_tokens % 32 + if thd_remainder != 0: + thd_pad = 32 - thd_remainder + hidden_states = torch.nn.functional.pad(hidden_states, (0, 0, 0, thd_pad)) + cu_seqlens = cu_seqlens.clone() + cu_seqlens[-1] = cu_seqlens[-1] + thd_pad + max_seqlen = max_seqlen + thd_pad + + kwargs["cu_seq_lens_q"] = kwargs["cu_seq_lens_k"] = cu_seqlens + kwargs["max_length_q"] = kwargs["max_length_k"] = max_seqlen + + if ( + self.config.attn_input_format == "thd" + and hidden_states.dim() == 3 + and hidden_states.size(0) == 1 + ): + hidden_states = hidden_states.squeeze(0) + + if ( + self.config.attn_input_format == "bshd" + and attention_mask is not None + and attention_mask.dim() == 2 + ): + attention_mask = ~attention_mask[:, None, None, :].bool() + + if isinstance(past_key_values, InferenceParams): + _ref = input_ids if input_ids is not None else inputs_embeds + lengths = ( + attention_mask.sum(dim=1).tolist() + if attention_mask is not None and attention_mask.shape[:2] == _ref.shape[:2] + else [1] * _ref.shape[0] + ) + past_key_values.pre_step(OrderedDict(zip(list(range(len(lengths))), lengths))) + + with torch.autocast(device_type="cuda", enabled=False): + te_rope_emb = self.rotary_emb(max_seq_len=self.config.max_position_embeddings) + + with self._outer_autocast(): + for layer_idx, decoder_layer in enumerate(self.layers): + hidden_states = decoder_layer( + hidden_states, + attention_mask=( + None if self.config.attn_input_format == "thd" else attention_mask + ), + rotary_pos_emb=te_rope_emb, + inference_params=past_key_values, + cu_seqlens_q=kwargs.get("cu_seq_lens_q", None), + cu_seqlens_kv=kwargs.get("cu_seq_lens_k", None), + cu_seqlens_q_padded=kwargs.get("cu_seq_lens_q_padded", None), + cu_seqlens_kv_padded=kwargs.get("cu_seq_lens_k_padded", None), + max_seqlen_q=kwargs.get("max_length_q", None), + max_seqlen_kv=kwargs.get("max_length_k", None), + pad_between_seqs=kwargs.get("pad_between_seqs", None), + ) + + hidden_states = self.norm(hidden_states) + + if should_pack_inputs: + if thd_remainder != 0: + hidden_states = hidden_states[:thd_orig_tokens] + hidden_states = _pad_input(hidden_states, indices, batch_size, padded_seq_len) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + hidden_states=None, + ) + + +class TEMixtralMXFP8ForCausalLM(TEMixtralMXFP8PreTrainedModel): + """Causal LM wrapper with MXFP8 autocast.""" + + _tied_weights_keys: ClassVar[list[str]] = [] + + def __init__( + self, + config: TEMixtralMXFP8Config, + fp8_recipe: te_recipe.Recipe | None = None, + dispatcher: AllToAllTokenDispatcher | None = None, + ) -> None: + super().__init__(config) + self.model = TEMixtralMXFP8Model(config, fp8_recipe=fp8_recipe, dispatcher=dispatcher) + self.vocab_size = config.vocab_size + with te.quantized_model_init(enabled=False): + self.lm_head = te.Linear( + config.hidden_size, + config.vocab_size, + bias=False, + params_dtype=config.dtype, + device="meta" if torch.get_default_device() == torch.device("meta") else "cuda", + init_method=lambda x: torch.nn.init.normal_( + x, mean=0.0, std=config.initializer_range + ), + ) + self.post_init() + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.Tensor | None = None, + past_key_values: Any = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + shift_labels: torch.Tensor | None = None, + use_cache: bool | None = None, + cache_position: torch.Tensor | None = None, + logits_to_keep: int | torch.Tensor = 0, + **kwargs: Any, + ) -> CausalLMOutputWithPast: + outputs = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=inputs_embeds, + use_cache=use_cache, + cache_position=cache_position, + **kwargs, + ) + + hidden_states = outputs.last_hidden_state + slice_indices = ( + slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep + ) + with te.autocast(enabled=False): + if hidden_states.ndim == 3: + logits = self.lm_head(hidden_states[:, slice_indices, :]) + else: + logits = self.lm_head(hidden_states[slice_indices, :]) + + loss = None + if labels is not None or shift_labels is not None: + loss = self.loss_function( + logits=logits, + labels=labels, + shift_labels=shift_labels, + vocab_size=self.config.vocab_size, + **kwargs, + ) + + if self.config.moe_aux_loss_coeff > 0 and loss is not None: + aux_loss = sum(layer.mlp._aux_loss for layer in self.model.layers) + loss = loss + aux_loss + + return CausalLMOutputWithPast( + loss=loss, + logits=logits, + past_key_values=outputs.past_key_values, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) diff --git a/docs/examples/te_mixtral/te_moe_dispatch.py b/docs/examples/te_mixtral/te_moe_dispatch.py new file mode 100644 index 000000000..0fc06a639 --- /dev/null +++ b/docs/examples/te_mixtral/te_moe_dispatch.py @@ -0,0 +1,298 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Token dispatch / combine for the MXFP8 TE-native MoE block. + +Wraps the permute/pad/all-to-all/sort-by-expert plumbing that moves tokens +from data-parallel ranks to their owning expert ranks, and reverses the +operation on the way back. The transport is NCCL ``all_to_all_single``; +the all-to-all is just the mechanism — the public API is ``dispatch()`` +/ ``combine()``. + +Per-expert MoE permute pads to 128 (grouped MXFP8 GEMM M-tile). Both +hidden states *and* per-token routing probabilities are transmitted so +the destination-side ``ScaledSwiGLU(glu_interleave_size=32)`` has its +scales locally — that's what trips the fused +``ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8`` kernel. ``combine()`` does not +re-apply routing weights (already applied inside ScaledSwiGLU). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any + +import torch +import torch.distributed as dist + +import transformer_engine.pytorch as te + +# Required so the ``@torch.compile`` helpers below can capture data-dependent +# tensor shapes (e.g. ``repeat_interleave`` with tensor counts) without +# bailing out to Python. Must be set at module level — torch.compile traces +# lazily on the first call, so a scoped setting wouldn't be active. +torch._dynamo.config.capture_scalar_outputs = True +torch._dynamo.config.capture_dynamic_output_shape_ops = True + + +# Per-expert token-count alignment required by the fused MXFP8 grouped-MLP +# CuTe-DSL kernel (ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8). The kernel +# rejects an input whose per-group token count is not a multiple of 256 +# ("Invalid a.shape[0] ... expected to be divisible by 256"). 128 was the +# old value that worked when the fused kernel wasn't firing on B300; with +# the SM10.x gate now passing on B300 we must pad to 256. +_MXFP8_GROUP_ALIGN = 256 + + +@dataclass +class DispatchOutput: + """Tokens, per-token probs, and split sizes routed to the local experts.""" + + expert_input: torch.Tensor + expert_probs: torch.Tensor + tokens_per_expert: list[int] + handle: Any + + +@dataclass +class _Handle: + row_id_map: torch.Tensor + restore_shape: torch.Size + pad_offsets: torch.Tensor | None + unsort_indices: torch.Tensor | None = None + input_split_sizes: list[int] | None = None + output_split_sizes: list[int] | None = None + + +class _DifferentiableAllToAll(torch.autograd.Function): + """``dist.all_to_all_single`` wrapped in autograd (works for 1-D and 2-D).""" + + @staticmethod + def forward( + ctx, + input: torch.Tensor, + output_split_sizes: list[int], + input_split_sizes: list[int], + group: dist.ProcessGroup, + ) -> torch.Tensor: + ctx.input_split_sizes = input_split_sizes + ctx.output_split_sizes = output_split_sizes + ctx.group = group + total_out = sum(output_split_sizes) + if input.dim() == 1: + output = torch.empty(total_out, device=input.device, dtype=input.dtype) + else: + output = torch.empty( + total_out, *input.shape[1:], device=input.device, dtype=input.dtype + ) + dist.all_to_all_single( + output, input.contiguous(), output_split_sizes, input_split_sizes, group=group + ) + return output + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + total_in = sum(ctx.input_split_sizes) + if grad_output.dim() == 1: + grad_input = torch.empty(total_in, device=grad_output.device, dtype=grad_output.dtype) + else: + grad_input = torch.empty( + total_in, *grad_output.shape[1:], device=grad_output.device, dtype=grad_output.dtype + ) + dist.all_to_all_single( + grad_input, + grad_output.contiguous(), + ctx.input_split_sizes, + ctx.output_split_sizes, + group=ctx.group, + ) + return grad_input, None, None, None + + +@torch.compile(fullgraph=True) +def _build_expert_sort_indices(recv_counts: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Build sort/unsort indices that regroup ``[src*expert]`` tokens by expert. + + After ``all_to_all`` tokens arrive grouped by source rank + (``[src0_exp0..src0_expL, src1_exp0..src1_expL, ...]``). ``GroupedLinear`` + needs them grouped by expert (``[all_exp0, all_exp1, ...]``). + + Uses only vectorized tensor ops (no ``.item()`` calls or Python-level + loops) so this is ``torch.compile(fullgraph=True)``-safe. + """ + ep_size, num_local_experts = recv_counts.shape + device = recv_counts.device + num_blocks = ep_size * num_local_experts + + counts_src = recv_counts.reshape(-1).long() + offsets_src = torch.zeros(num_blocks, dtype=torch.long, device=device) + offsets_src[1:] = counts_src[:-1].cumsum(0) + + counts_exp = recv_counts.t().contiguous().reshape(-1).long() + offsets_exp = torch.zeros(num_blocks, dtype=torch.long, device=device) + offsets_exp[1:] = counts_exp[:-1].cumsum(0) + + total = counts_src.sum() + + s_idx = torch.arange(ep_size, device=device).unsqueeze(1).expand(ep_size, num_local_experts) + e_idx = ( + torch.arange(num_local_experts, device=device) + .unsqueeze(0) + .expand(ep_size, num_local_experts) + ) + src_to_exp = (e_idx * ep_size + s_idx).reshape(-1) + shifts = offsets_exp[src_to_exp] - offsets_src + token_shifts = shifts.repeat_interleave(counts_src) + src_positions = torch.arange(total, device=device) + dst_positions = src_positions + token_shifts + sort_indices = torch.empty(total, dtype=torch.long, device=device) + sort_indices[dst_positions] = src_positions + unsort_indices = torch.empty_like(sort_indices) + unsort_indices[sort_indices] = torch.arange(total, device=device) + return sort_indices, unsort_indices + + +class AllToAllTokenDispatcher: + """NCCL all-to-all dispatcher for the TE-native MXFP8 MoE block. + + Args: + num_experts: Total global experts. + num_local_experts: Experts owned by this rank. + hidden_size: Hidden feature dim. + ep_size: Expert-parallel world size (1 = single-process). + pad_align: Per-expert split alignment. Must be a multiple of 128 for + the grouped MXFP8 GEMM. + """ + + def __init__( + self, + num_experts: int, + num_local_experts: int, + hidden_size: int, + ep_size: int, + pad_align: int = _MXFP8_GROUP_ALIGN, + ) -> None: + self.num_experts = num_experts + self.num_local_experts = num_local_experts + self.hidden_size = hidden_size + self.ep_size = ep_size + self.pad_align = pad_align + self._ep_group: dist.ProcessGroup | None = None + + def set_ep_group(self, ep_group: dist.ProcessGroup) -> None: + self._ep_group = ep_group + + def dispatch( + self, + hidden_states: torch.Tensor, + selected_experts: torch.Tensor, + routing_weights: torch.Tensor, + tokens_per_expert: torch.Tensor, + ) -> DispatchOutput: + """Permute -> pad -> (all-to-all) -> sort-by-expert. + + ``tokens_per_expert`` is required: the MoE block already computes the + per-expert token count (for the fused aux loss + the routing tables), + so the dispatcher takes it as input rather than launching another + ``torch.bincount`` kernel. + """ + num_tokens = hidden_states.shape[0] + + # Dense per-expert routing tables required by ``moe_permute_and_pad_with_probs``. + routing_map = torch.zeros( + num_tokens, self.num_experts, dtype=torch.bool, device=hidden_states.device + ) + routing_map.scatter_(1, selected_experts, True) + routing_probs = torch.zeros( + num_tokens, self.num_experts, dtype=routing_weights.dtype, device=hidden_states.device + ) + routing_probs.scatter_(1, selected_experts, routing_weights) + + ( + permuted_hidden, + permuted_probs, + row_id_map, + pad_offsets, + padded_tokens_per_expert, + ) = te.moe_permute_and_pad_with_probs( + hidden_states, + routing_probs, + routing_map, + tokens_per_expert, + self.pad_align, + ) + padded_tokens_per_expert = padded_tokens_per_expert.int() + + if self._ep_group is None or self.ep_size == 1: + handle = _Handle( + row_id_map=row_id_map, + restore_shape=hidden_states.shape, + pad_offsets=pad_offsets, + ) + return DispatchOutput( + expert_input=permuted_hidden, + expert_probs=permuted_probs, + tokens_per_expert=padded_tokens_per_expert.tolist(), + handle=handle, + ) + + # EP > 1: ship both tokens and probs across ranks. A single packed + # all_to_all was tried; the extra ``.contiguous()`` slicing on the + # receive side cost more than the saved NCCL collective at Mixtral + # batch=8 / seq=8192 (the probs comm is ~1/2048 of the token comm, + # so the all_to_all is bandwidth-bound, not latency-bound). + ep_group = self._ep_group + send_counts = padded_tokens_per_expert.reshape(self.ep_size, self.num_local_experts) + recv_counts = torch.empty_like(send_counts) + dist.all_to_all_single(recv_counts.flatten(), send_counts.flatten(), group=ep_group) + + input_split_sizes = send_counts.sum(dim=1).tolist() + output_split_sizes = recv_counts.sum(dim=1).tolist() + local_m_splits = recv_counts.sum(dim=0).int().tolist() + + recv_tokens = _DifferentiableAllToAll.apply( + permuted_hidden, output_split_sizes, input_split_sizes, ep_group + ) + recv_probs = _DifferentiableAllToAll.apply( + permuted_probs, output_split_sizes, input_split_sizes, ep_group + ) + + sort_indices, unsort_indices = _build_expert_sort_indices(recv_counts) + + handle = _Handle( + row_id_map=row_id_map, + restore_shape=hidden_states.shape, + pad_offsets=pad_offsets, + unsort_indices=unsort_indices, + input_split_sizes=input_split_sizes, + output_split_sizes=output_split_sizes, + ) + return DispatchOutput( + expert_input=recv_tokens[sort_indices], + expert_probs=recv_probs[sort_indices], + tokens_per_expert=local_m_splits, + handle=handle, + ) + + def combine(self, expert_output: torch.Tensor, handle: _Handle) -> torch.Tensor: + """Reverse the dispatch. ``ScaledSwiGLU`` already applied per-token probs, + so ``moe_unpermute`` is called without ``merging_probs``.""" + if handle.unsort_indices is not None: + combined = _DifferentiableAllToAll.apply( + expert_output[handle.unsort_indices], + handle.input_split_sizes, + handle.output_split_sizes, + self._ep_group, + ) + else: + combined = expert_output + + return te.moe_unpermute( + combined, + handle.row_id_map, + merging_probs=None, + restore_shape=handle.restore_shape, + map_type="mask", + pad_offsets=handle.pad_offsets, + ) diff --git a/docs/examples/te_mixtral/test_accuracy.py b/docs/examples/te_mixtral/test_accuracy.py new file mode 100644 index 000000000..46ed2a5bf --- /dev/null +++ b/docs/examples/te_mixtral/test_accuracy.py @@ -0,0 +1,188 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Forward + backward parity check for te_mixtral (BF16 and MXFP8) vs HF. + +Compares logits, loss, and weight gradients (``model.embed_tokens.weight`` +and ``lm_head.weight``) between the HuggingFace reference and the +TransformerEngine port in both BF16 and MXFP8 modes. +""" + +import torch +from transformers import MixtralConfig, MixtralForCausalLM + +import transformer_engine.pytorch as te +from transformer_engine.common import recipe as te_recipe + +from te_mixtral import TEMixtralForCausalLM, replace_params as replace_params_bf16 +from te_mixtral_mxfp8 import ( + TEMixtralMXFP8ForCausalLM, + replace_params as replace_params_mxfp8, +) + + +# BF16 should match HF very closely. +BF16_TOL = 0.01 + +# MXFP8 quantizes activations to FP8 with per-tile bf16 scales, so we expect +# larger logit / gradient drift than the BF16 case. +MXFP8_LOGITS_ATOL = 1.5 +MXFP8_LOGITS_RTOL = 0.05 +MXFP8_LOSS_ATOL = 0.05 +MXFP8_LOSS_RTOL = 0.05 +MXFP8_GRAD_ATOL = 1.0 +MXFP8_GRAD_RTOL = 0.1 + + +def _build_config(): + return MixtralConfig( + hidden_size=256, + intermediate_size=512, + num_local_experts=4, + num_experts_per_tok=2, + num_hidden_layers=2, + num_attention_heads=8, + num_key_value_heads=8, + vocab_size=1024, + max_position_embeddings=128, + router_jitter_noise=0.0, + rms_norm_eps=1e-5, + ) + + +def _load_te_weights(model_te, model_hf, replace_params_fn): + te_state_dict = model_te.state_dict() + replace_params_fn(model_hf.state_dict(), te_state_dict, model_te.config) + missing, unexpected = model_te.load_state_dict(te_state_dict, strict=False) + if unexpected: + raise RuntimeError(f"Unexpected TE keys during load: {unexpected}") + allowed_missing = [k for k in missing if k.endswith("_extra_state")] + if len(allowed_missing) != len(missing): + raise RuntimeError(f"Unexpected missing TE keys during load: {missing}") + + +def _zero_grads(model): + for p in model.parameters(): + if p.grad is not None: + p.grad = None + + +def _forward_backward(model, input_ids, attention_mask, labels, *, fp8_recipe=None): + """Return (logits, loss, embed_grad, lm_head_grad), all detached as float32.""" + _zero_grads(model) + if fp8_recipe is not None: + with te.autocast(enabled=True, recipe=fp8_recipe): + out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + out.loss.backward() + else: + out = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels) + out.loss.backward() + return ( + out.logits.detach().float(), + out.loss.detach().float(), + model.model.embed_tokens.weight.grad.detach().float().clone(), + model.lm_head.weight.grad.detach().float().clone(), + ) + + +def _compare(label, hf, te_, *, atol, rtol): + diff = (hf - te_).abs() + max_diff = diff.max().item() + mean_diff = diff.mean().item() + print(f" {label:<14s} max={max_diff:.6f} mean={mean_diff:.6f}") + torch.testing.assert_close(te_, hf, atol=atol, rtol=rtol) + + +def _make_inputs(cfg, device): + torch.manual_seed(1) + # seq divisible by 32 so the MXFP8 path is happy. + input_ids = torch.randint(0, cfg.vocab_size, (2, 64), device=device) + attention_mask = torch.ones_like(input_ids, device=device) + labels = input_ids.clone() + return input_ids, attention_mask, labels + + +def _build_hf(cfg, device, dtype): + torch.manual_seed(0) + model = MixtralForCausalLM(cfg).to(device=device, dtype=dtype) + model.eval() + return model + + +def _run_bf16(cfg, model_hf, inputs, device, dtype): + print("=" * 64) + print("BF16 parity check (forward + backward)") + print("=" * 64) + + te_cfg = TEMixtralForCausalLM.config_class(**cfg.to_dict()) + model_te = TEMixtralForCausalLM(te_cfg).to(device=device, dtype=dtype) + _load_te_weights(model_te, model_hf, replace_params_bf16) + model_te.eval() + + input_ids, attention_mask, labels = inputs + hf_logits, hf_loss, hf_embed_g, hf_lm_g = _forward_backward( + model_hf, input_ids, attention_mask, labels + ) + te_logits, te_loss, te_embed_g, te_lm_g = _forward_backward( + model_te, input_ids, attention_mask, labels + ) + + print(f" logits shape {tuple(hf_logits.shape)}") + print(f" HF loss {hf_loss.item():.6f}") + print(f" TE loss {te_loss.item():.6f}") + _compare("logits", hf_logits, te_logits, atol=BF16_TOL, rtol=0.0) + _compare("loss", hf_loss, te_loss, atol=BF16_TOL, rtol=0.0) + _compare("embed.grad", hf_embed_g, te_embed_g, atol=BF16_TOL, rtol=0.0) + _compare("lm_head.grad", hf_lm_g, te_lm_g, atol=BF16_TOL, rtol=0.0) + print("BF16 parity OK.\n") + + +def _run_mxfp8(cfg, model_hf, inputs, device, dtype): + print("=" * 64) + print("MXFP8 parity check (forward + backward)") + print("=" * 64) + + te_cfg = TEMixtralMXFP8ForCausalLM.config_class(**cfg.to_dict()) + te_cfg.attn_input_format = "bshd" + te_cfg.self_attn_mask_type = "causal" + te_cfg.expert_parallel_size = 1 + te_cfg.dtype = dtype + recipe = te_recipe.MXFP8BlockScaling(fp8_format=te_recipe.Format.E4M3) + model_te = TEMixtralMXFP8ForCausalLM(te_cfg, fp8_recipe=recipe).to(device=device, dtype=dtype) + _load_te_weights(model_te, model_hf, replace_params_mxfp8) + model_te.eval() + + input_ids, attention_mask, labels = inputs + hf_logits, hf_loss, hf_embed_g, hf_lm_g = _forward_backward( + model_hf, input_ids, attention_mask, labels + ) + te_logits, te_loss, te_embed_g, te_lm_g = _forward_backward( + model_te, input_ids, attention_mask, labels, fp8_recipe=recipe + ) + + print(f" logits shape {tuple(hf_logits.shape)}") + print(f" HF loss {hf_loss.item():.6f}") + print(f" TE loss {te_loss.item():.6f}") + _compare("logits", hf_logits, te_logits, atol=MXFP8_LOGITS_ATOL, rtol=MXFP8_LOGITS_RTOL) + _compare("loss", hf_loss, te_loss, atol=MXFP8_LOSS_ATOL, rtol=MXFP8_LOSS_RTOL) + _compare("embed.grad", hf_embed_g, te_embed_g, atol=MXFP8_GRAD_ATOL, rtol=MXFP8_GRAD_RTOL) + _compare("lm_head.grad", hf_lm_g, te_lm_g, atol=MXFP8_GRAD_ATOL, rtol=MXFP8_GRAD_RTOL) + print("MXFP8 parity OK.\n") + + +def main() -> None: + assert torch.cuda.is_available(), "CUDA required." + + cfg = _build_config() + device = "cuda" + dtype = torch.bfloat16 + + model_hf = _build_hf(cfg, device, dtype) + inputs = _make_inputs(cfg, device) + + _run_bf16(cfg, model_hf, inputs, device, dtype) + _run_mxfp8(cfg, model_hf, inputs, device, dtype) + + +if __name__ == "__main__": + main() diff --git a/docs/examples/te_mixtral/tutorial_accelerate_hf_mixtral_with_te.ipynb b/docs/examples/te_mixtral/tutorial_accelerate_hf_mixtral_with_te.ipynb new file mode 100644 index 000000000..decdcff31 --- /dev/null +++ b/docs/examples/te_mixtral/tutorial_accelerate_hf_mixtral_with_te.ipynb @@ -0,0 +1,462 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Accelerating Hugging Face Mixtral MoE Fine-Tuning with Transformer Engine\n", + "\n", + "
\n", + "\n", + "Goal\n", + "\n", + "This tutorial showcases how to accelerate fine-tuning a mixture-of-experts model, [Mixtral-8x7B](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1), with Transformer Engine (TE) in `BF16` and `MXFP8` precision.\n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "**Setup**\n", + "\n", + "Mixtral-8x7B has 8 experts and roughly 47B total parameters. In `BF16` the model weights alone consume ~93 GB, and full `AdamW` fine-tuning needs ~370 GB. This tutorial is tested on 8x B300 GPUs with `Expert Parallelism (EP) = 2` and `Data Parallelism (DP) = 4`, so the experts are divided across 2 GPUs and there are 4 replicas. The container used is [pytorch-26.04-py3](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch?version=26.04-py3). A sequence length of 8192 and a global batch size of 48 are used across the experiments.\n", + "\n", + "Install the required Python packages using the following command in a terminal:" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "```bash\n", + "pip install -r requirements.txt \n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Table of Contents\n", + "\n", + "1. [Baseline] Running HF Mixtral -- Without Expert Parallelism (Precision: `BF16`)\n", + "2. [Improvement 1] Transformer Engine with Expert Parallelism (Precision: `BF16`)\n", + "3. [Improvement 2] Batched Expert Execution with `GroupedLinear` (Precision: `BF16`)\n", + "4. [Improvement 3] Precision Optimization and Fused MLP (Precision: `MXFP8`)\n", + "5. Conclusion\n", + "6. Appendix: Dependencies" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## [Baseline] Running HF Mixtral -- Without Expert Parallelism (Precision: `BF16`)\n", + "\n", + "Before applying any Transformer Engine optimizations, we establish a Hugging Face (HF) baseline. Mixtral replaces the standard Transformer feed-forward network (FFN) with a sparse **Mixture of Experts** (MoE): a learned router selects the top-2 experts out of 8 per token, as shown in Fig 1. \n", + "\n", + "
\n", + "\n", + "
Fig 1: Dense Transformer block (left) vs Sparse MoE Transformer block (right).
\n", + "
\n", + "\n", + "\n", + "The current HF implementation has two limitations.\n", + "\n", + "1. **Pipeline parallelism**. Because the full model does not fit on one GPU, the baseline uses pipeline parallelism to split the model across GPUs. This is the simplest way to partition a model, but GPU utilization is limited by pipeline bubbles and sequential layer dependencies.\n", + "\n", + "\n", + "2. **Excessive kernel launches.** [HF's MixtralSparseMoeBlock](https://github.com/huggingface/transformers/blob/3ef278124e47832f34406ca3ca85bc50ad8b79bb/src/transformers/models/mixtral/modeling_mixtral.py) iterates over all 8 experts in a Python loop. Each expert triggers individual kernel launches. \n", + "\n", + "```python\n", + "for expert_idx, expert_layer in enumerate(self.experts):\n", + " idx, top_x = torch.where(expert_mask[expert_idx])\n", + " current_state = hidden_states[None, top_x].reshape(-1, hidden_dim)\n", + " current_hidden = expert_layer(current_state) * routing_weights[top_x, idx, None]\n", + " final_hidden_states.index_add_(0, top_x, current_hidden)\n", + "```\n", + "\n", + "For each layer, HF loops through the experts sequentially. Each expert is much smaller than the dense FFN, so each expert GEMM is small and cannot saturate the GPU's tensor cores. Looping over many experts therefore launches many small GEMMs, leaving the FFN dominated by orchestration overhead and memory movement.\n", + "\n", + "\n", + "The script [run_finetune_ep.py](run_finetune_ep.py) initializes Hugging Face and then runs fine-tuning. For the full implementation, refer to [utils.py](utils.py). Now, let's execute the following command in the terminal." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "```bash\n", + "python3 run_finetune_ep.py --improvement 0 --batch-size 48 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Here is the expected output:\n", + "\n", + "```\n", + "30 fine-tuning steps complete!\n", + "Median time per step: 2472 ms\n", + "```\n", + "\n", + "Let's add this information in a table and keep comparing it with a few possible improvements in future sections:\n", + "\n", + "| Models | Precision | Step Time | Speedup (over baseline) |\n", + "|---|---|---:|---:|\n", + "| HF baseline | BF16 | 2472 ms | 1 |" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## [Improvement 1] Transformer Engine with Expert Parallelism (Precision: `BF16`)\n", + "\n", + "Now that we have a baseline, let's bring in Transformer Engine. This section replaces the HF Transformer block with TE modules and introduces expert parallelism (EP). \n", + "\n", + "
\n", + "\n", + "
Fig 2: HF MixtralDecoderLayer (left) wrapped by TE modules (right).
\n", + "
\n", + "\n", + "**Fused Building Blocks**\n", + "\n", + "- **Attention block.** Instead of using one module for layer norm and another module for attention, TE combines them (`RMSNorm` and attention) with `te.MultiheadAttention`, where the input `RMSNorm` is bundled with the fused QKV projection. The layer norm weights and QKV weights are stored in the same building block: `self_attention.layernorm_qkv.weight`. Here is how you can use TE's attention block:\n", + "\n", + " ```python\n", + " self.self_attention = transformer_engine.pytorch.MultiheadAttention(\n", + " hidden_size=config.hidden_size,\n", + " fuse_qkv_params=True,\n", + " qkv_weight_interleaved=True,\n", + " normalization=\"RMSNorm\",\n", + " input_layernorm=True,\n", + " ...\n", + " )\n", + " ```\n", + "\n", + "- **MoE block.** TE provides the building blocks for the MoE layer. First, the gate computes router probabilities, then `softmax` and `top-k` select the top two experts out of eight for each token. The selected experts are then passed to the dispatcher. The dispatcher determines which EP ranks host the selected MoE experts, and NCCL handles the all-to-all communication that moves tokens to those ranks. The dispatcher also uses `transformer_engine.pytorch.moe_permute_and_pad_with_probs` to handle padding requirements, such as padding to multiples of 32 required by `MXFP8`.\n", + "\n", + " Below is the overview pseudocode for the MoE block:\n", + "\n", + " ```python\n", + " router_logits = self.gate(hidden_states) \n", + "\n", + " softmax_probs = torch.nn.functional.softmax(router_logits, dim=-1)\n", + "\n", + " routing_weights, selected_experts = torch.topk(softmax_probs, self.top_k, dim=-1)\n", + "\n", + " dispatch_output = self.dispatcher.dispatch(hidden_states, selected_experts, routing_weights)\n", + " ```\n", + "\n", + "**Parallelism layout**\n", + "\n", + "In this tutorial, EP=2 is used to split the model between 2 GPUs, each hosting 4 experts. In this 8-GPU setup, the model is replicated 4 times, creating 4 data-parallel groups.\n", + "\n", + "Here is how to set up EP:\n", + "\n", + "```python\n", + "config.expert_parallel_size = 2\n", + "ep_size = config.expert_parallel_size\n", + "dp_size = world_size // ep_size\n", + "ep_group = None\n", + "for dp_rank in range(dp_size):\n", + " ranks = list(range(dp_rank * ep_size, (dp_rank + 1) * ep_size))\n", + " group = dist.new_group(ranks=ranks)\n", + " if dist.get_rank() in ranks:\n", + " ep_group = group\n", + "model.model.set_ep_groups(ep_group=ep_group)\n", + "```\n", + "\n", + "**Mapping the HF checkpoint to TE**\n", + "\n", + "Some weights/parameters need to be reshaped and also remapped to corresponding weight names in TE modules. The `replace_params` helper in [te_mixtral.py](te_mixtral.py) performs the mapping (also illustrated in Fig 2 above). The two non-trivial groups are:\n", + "\n", + "- **Attention.** HF stores Q, K, V as separate projections; TE fuses them into a single QKV weight that lives under the `layernorm_qkv` submodule:\n", + "\n", + "| HF key | TE key |\n", + "|---|---|\n", + "| `self_attn.q_proj.weight` | `self_attention.layernorm_qkv.weight` (Q slice) |\n", + "| `self_attn.k_proj.weight` | `self_attention.layernorm_qkv.weight` (K slice) |\n", + "| `self_attn.v_proj.weight` | `self_attention.layernorm_qkv.weight` (V slice) |\n", + "| `input_layernorm.weight` | `self_attention.layernorm_qkv.layer_norm_weight` |\n", + "\n", + "- **MoE experts.** HF packs all experts' `SwiGLU` projections into two tensors per layer; TE keeps the same packing under different attribute names so `replace_params` is essentially a copy:\n", + "\n", + "| HF key | TE key |\n", + "|---|---|\n", + "| `mlp.experts.gate_up_proj` `[num_experts, 2*ffn, h]` | `mlp.experts_gate_up_weight` |\n", + "| `mlp.experts.down_proj` `[num_experts, h, ffn]` | `mlp.experts_down_weight` |\n", + "| `mlp.gate.weight` | `mlp.gate.weight` |\n", + "\n", + "All other weights (embeddings, norms, LM head) are direct copies. See `replace_params` in `te_mixtral.py` for the full mapping.\n", + "\n", + "Let's launch the same fine-tuning loop -- this time across 8 GPUs via `torchrun`. See `run_finetune_ep.py` and `utils.py` for the full implementation.\n", + "\n", + "Now, let's execute the following command." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "```bash\n", + "torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 1 --ep-size 2 --batch-size 12 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Here is the expected output:\n", + "\n", + "```\n", + "30 fine-tuning steps complete!\n", + "Median time per step: 747 ms\n", + "```\n", + "\n", + "Compared to the baseline implementation, we see the following result:\n", + "\n", + "| Models | Precision | Step Time | Speedup (over baseline) |\n", + "|---|---|---:|---:|\n", + "| HF baseline | BF16 | 2472 ms | 1 |\n", + "| TE decoder, TE building blocks, and MoE layer | BF16 | 747 ms | 3.31 |\n", + "\n", + "Improvement 1 is 3.31x faster than the baseline, a **231%** speedup." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## [Improvement 2] Batched Expert Execution with `GroupedLinear` (Precision: `BF16`)\n", + "\n", + "Improvement 1 kept the per-expert FFN as a Python loop. If each rank owns 4 local experts, the loop launches the per-expert GEMMs one by one. Another limitation is that the expert GEMMs are usually small: each one sees only the tokens routed to it, which is too small to feed the tensor cores efficiently. The following section shows how to execute the GEMMs in a batch. \n", + "\n", + "
\n", + "\n", + "
Fig 3: Left: looping through experts one-by-one. Right: one grouped-GEMM over all experts.
\n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "\n", + "`GroupedLinear` applies multiple linear transformations in one call. It gathers the experts' weights and input tokens. Although each expert receives a different number of tokens, `GroupedLinear` supports this by accepting per-expert token counts (`split_sizes`). `GroupedLinear` submits the local experts through TE’s grouped GEMM path instead of launching one PyTorch Linear operation per expert. This reduces per-expert launch and scheduling overhead. \n", + "\n", + "Here are the steps to use `GroupedLinear`. Each expert keeps its own weight tensor (`weight0`, `weight1`, ...), and the call takes the per-expert token counts as an extra positional argument:\n", + "\n", + "```python\n", + "from transformer_engine.pytorch.ops import GroupedLinear\n", + "\n", + "experts_gate_up = GroupedLinear(\n", + " num_groups=num_local_experts,\n", + " in_features=hidden_size,\n", + " out_features=2 * intermediate_size,\n", + " bias=False,\n", + " dtype=torch.bfloat16,\n", + " device=\"cuda\",\n", + ")\n", + "\n", + "gate_up_output = experts_gate_up(tokens, split_sizes)\n", + "```\n", + "\n", + "Compared with the Python loop in Improvement 1, this becomes one gate-up projection per layer instead of 4 separate calls (4 is the number of experts on a GPU). The expert weights can be imported from HF. In `te_mixtral.py`, the grouped-op path keeps each local expert as a normal per-expert `weight{i}` parameter and loads the owning expert slice directly.\n", + "\n", + "To see the effect of `GroupedLinear`, we keep everything else unchanged. Execute the following command in the terminal. " + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "```bash\n", + "torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 2 --ep-size 2 --batch-size 12 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Here is the expected output:\n", + "\n", + "```\n", + "30 fine-tuning steps complete!\n", + "Median time per step: 635 ms\n", + "```\n", + "\n", + "Adding the `GroupedLinear` result gives us:\n", + "\n", + "| Models | Precision | Step Time | Speedup (over baseline) |\n", + "|---|---|---:|---:|\n", + "| HF baseline | BF16 | 2472 ms | 1 |\n", + "| TE decoder, TE building blocks, and MoE layer | BF16 | 747 ms | 3.31 |\n", + "| TE with `GroupedLinear` | BF16 | 635 ms | 3.89 |\n", + "\n", + "`GroupedLinear` reaches a 3.89x speedup over the baseline, or **289%**." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## [Improvement 3] Precision Optimization and Fused MLP (Precision: `MXFP8`)\n", + "\n", + "With EP and grouped expert GEMMs in place, the next improvement lowers precision from `BF16` to `MXFP8`. `MXFP8` converts the weight and activation values to 8 bits instead of 16 bits. To preserve dynamic range, it keeps one `E8M0` scale factor for every 32 values; applying that scale recovers a wider numerical range. On Blackwell GPUs, `MXFP8` is native and hardware accelerated, so `MXFP8` GEMMs can run through specialized Tensor Core instructions. Read more about `MXFP8` and block scaling in the [Transformer Engine FP8 primer](https://docs.nvidia.com/deeplearning/transformer-engine/user-guide/examples/fp8_primer.html#MXFP8-and-block-scaling).\n", + "\n", + "The model still keeps its master weights in `BF16`, so using `MXFP8` adds `Quantization` and `De-Quantization` work around the GEMMs. `Quantization` converts `BF16` weights and activations to `MXFP8` before the low-precision GEMM; `De-Quantization` converts the result back to the higher-precision format. The naive path performs these as separate operations. This motivates the fused MLP path shown below.\n", + "\n", + "
\n", + "\n", + "
Fig 4: The MXFP8 path fuses multiple operations into one kernel before the down projection.
\n", + "
\n", + "\n", + "To use `MXFP8`, we simply define a recipe and pass it to the model. \n", + "\n", + "```python\n", + "fp8_recipe = te_recipe.MXFP8BlockScaling()\n", + "model = TEMixtralMXFP8ForCausalLM(config, fp8_recipe=fp8_recipe, dispatcher=dispatcher)\n", + "```\n", + "\n", + "Now, the model's forward and backward passes run under `MXFP8` precision which is enabled through TE's `autocast` API:\n", + "\n", + "```python\n", + "with te.autocast(enabled=True, recipe=self._fp8_recipe):\n", + " for decoder_layer in self.layers:\n", + " hidden_states = decoder_layer(hidden_states)\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "To use fused MLP, we import TE's `Sequential` API to chain together `gate_up`, `ScaledSwiGLU`, and `down`. It also folds the `De-Quantization` step into the fused path. `ScaledSwiGLU` is chosen to combine the routing probabilities (\"scales\") with the expert FFN computations. \n", + "\n", + "```python\n", + "from transformer_engine.pytorch.ops import GroupedLinear, ScaledSwiGLU, Sequential\n", + "\n", + "experts_ffn = Sequential(GroupedLinear(gate_up), ScaledSwiGLU(), GroupedLinear(down))\n", + "```\n", + "\n", + "TE's `Sequential` scans the ops and, if the pattern matches, replaces the `GroupedLinear -> ScaledSwiGLU -> GroupedLinear` pattern with a fused operation object: `ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8` for forward and a matching fused backward op. It reduces framework overhead, fuses the SwiGLU/probability-scaling work into the grouped MLP path, and avoids some intermediate materialization.\n", + "\n", + "
\n", + "\n", + "Note\n", + "\n", + "`NVTE_CUTEDSL_FUSED_GROUPED_MLP=1` must be set before TE imports the fused op registration. In this tutorial, `run_finetune_ep.py` already does that automatically for improvement 3.\n", + "\n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Execute the following in the terminal:\n", + "\n", + "```bash\n", + "torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 3 --ep-size 2 --batch-size 12 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n", + "```\n", + "\n", + "Here is the expected result:\n", + "```\n", + "30 fine-tuning steps complete!\n", + "Median time per step: 542 ms\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "With MXFP8 fused MLP included, the final comparison is:\n", + "\n", + "| Models | Precision | Step Time | Speedup (over baseline) |\n", + "|---|---|---:|---:|\n", + "| HF baseline | BF16 | 2472 ms | 1 |\n", + "| TE EP Python loop | BF16 | 747 ms | 3.31 |\n", + "| TE with `GroupedLinear` | BF16 | 635 ms | 3.89 |\n", + "| TE with MXFP8 fused MLP | MXFP8 | 542 ms | 4.56 |\n", + "\n", + "For Mixtral-8x7B, we get the largest speedup with MXFP8 fused MLP: 4.56x faster than the baseline, or **356%**." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Conclusion\n", + "\n", + "This tutorial walks through three progressive optimization improvements that speed up the fine-tuning of the Mixtral-8x7B model by replacing building blocks in a Hugging Face baseline with TE-native blocks like MXFP8 and fused grouped MLP. The tutorial uses **global batch 48, seq 8192** on 8x B300 to demonstrate the speedups.\n", + "\n", + "To run all improvements together, execute the following in a terminal. All four runs use the same global batch size of 48. The TE runs use DP=4, so the per-rank batch size is 12.\n", + "\n", + "```bash\n", + "python3 run_finetune_ep.py --improvement 0 --batch-size 48 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n", + "\n", + "torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 1 --ep-size 2 --batch-size 12 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n", + "\n", + "torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 2 --ep-size 2 --batch-size 12 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n", + "\n", + "torchrun --standalone --nproc_per_node=8 run_finetune_ep.py --improvement 3 --ep-size 2 --batch-size 12 --max-seq-length 8192 --warmup-steps 5 --train-steps 30\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "
\n", + "\n", + "Note on Scaling\n", + "\n", + "For large-scale training, check out [Megatron's performance summary](https://docs.nvidia.com/nemo/megatron-bridge/latest/performance-summary.html).\n", + "\n", + "
" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Appendix: Dependencies\n", + "\n", + "| File | Purpose |\n", + "|---|---|\n", + "| `te_mixtral.py` | BF16 TE Mixtral implementation |\n", + "| `te_mixtral_mxfp8.py` | MXFP8 implementation |\n", + "| `te_moe_dispatch.py` | Token dispatch/combine for MXFP8 |\n", + "| `hf_to_te_weights.py` | Converts Hugging Face weights to Transformer Engine format |\n", + "| `utils.py` | Training loop |\n", + "| `run_finetune_ep.py` | CLI launcher |\n", + "| `requirements.txt` | Python package versions |\n", + "| `collator.py` | Input sequence preparation |" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python", + "version": "3.12.0" + }, + "nbsphinx": { + "execute": "never" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} diff --git a/docs/examples/te_mixtral/utils.py b/docs/examples/te_mixtral/utils.py new file mode 100644 index 000000000..559bab885 --- /dev/null +++ b/docs/examples/te_mixtral/utils.py @@ -0,0 +1,485 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import os + +import torch +import torch.distributed as dist +from torch.optim import AdamW +from torch.utils.data import DataLoader +from torch.utils.data.distributed import DistributedSampler + +from transformers import ( + AutoModelForCausalLM, + AutoTokenizer, + AutoConfig, + get_linear_schedule_with_warmup, + DataCollatorForLanguageModeling, +) +from datasets import load_dataset +from accelerate import Accelerator + + +class HyperParameters: + def __init__(self): + # "bf16" (improvements 1-2) or "mxfp8" (improvement 3). + self.mixed_precision = "bf16" + + self.model_name = "mistralai/Mixtral-8x7B-v0.1" + self.dataset_name = "timdettmers/openassistant-guanaco" + self.dataset_text_field = "text" + self.learning_rate = 1.41e-5 + self.batch_size = 4 + self.max_seq_length = 2048 + self.gradient_accumulation_steps = 1 + self.num_warmup_steps = 5 + self.num_training_steps = 3 + + self.weights_cache_dir = "" + self.hf_access_token = "" + self.expert_parallel_size = 8 + + # "loop" (improvement 1) or "grouped_op" (improvements 2-3). + self.expert_ffn_mode = "grouped_op" + + # "te_mixtral" (improvements 1-2, BF16) or "te_mixtral_mxfp8" (improvement 3). + self.model_impl = "te_mixtral" + + +def get_dataloaders(accelerator: Accelerator, hyperparams: HyperParameters): + from collator import DataCollatorWithFlattening + + dataset = load_dataset(hyperparams.dataset_name, split="train") + tokenizer = AutoTokenizer.from_pretrained(hyperparams.model_name) + if getattr(tokenizer, "pad_token", None) is None: + tokenizer.pad_token = tokenizer.eos_token + + def tokenize(element): + outputs = tokenizer( + element["text"], + truncation=True, + padding=False, + max_length=hyperparams.max_seq_length, + return_overflowing_tokens=False, + return_length=False, + ) + return {"input_ids": outputs["input_ids"], "attention_mask": outputs["attention_mask"]} + + with accelerator.main_process_first(): + dataset = dataset.map(tokenize, batched=True, remove_columns=dataset.column_names) + + pad_multiple = 32 if hyperparams.mixed_precision == "mxfp8" else 16 + bshd_collator = DataCollatorForLanguageModeling( + tokenizer=tokenizer, + mlm=False, + pad_to_multiple_of=pad_multiple, + ) + data_collator = DataCollatorWithFlattening( + collator=bshd_collator, + pad_to_multiple_of=pad_multiple, + separator_id=-100, + ) + + sampler = None + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if hyperparams.expert_parallel_size > 1 and world_size > 1: + ep_size = hyperparams.expert_parallel_size + dp_size = world_size // ep_size + global_rank = dist.get_rank() if dist.is_initialized() else int(os.environ.get("RANK", "0")) + sampler = DistributedSampler( + dataset, + num_replicas=dp_size, + rank=global_rank // ep_size, + shuffle=True, + drop_last=True, + ) + + train_dataloader = DataLoader( + dataset, + batch_size=hyperparams.batch_size, + sampler=sampler, + collate_fn=data_collator, + drop_last=True, + ) + return train_dataloader + + +def ensure_model_is_downloaded(hyperparams: HyperParameters): + assert hyperparams.model_name in [ + "mistralai/Mixtral-8x7B-v0.1", + "mistralai/Mixtral-8x22B-v0.1", + ], "Only Mixtral-8x7B-v0.1 and Mixtral-8x22B-v0.1 are supported." + + from huggingface_hub import login, snapshot_download + + try: + login(hyperparams.hf_access_token) + except Exception as e: + if "Invalid token passed!" in str(e): + print( + "Please provide a valid HF Access Token. " + "See: https://huggingface.co/docs/hub/en/security-tokens" + ) + else: + print(f"Login exception: {e}") + + hyperparams.weights_cache_dir = snapshot_download( + repo_id=hyperparams.model_name, + cache_dir=hyperparams.weights_cache_dir or None, + ) + print(f"Model cache directory: {hyperparams.weights_cache_dir}") + + +def init_baseline_model(hyperparams: HyperParameters): + """Load the vanilla HuggingFace Mixtral model in BF16.""" + ensure_model_is_downloaded(hyperparams) + + config = AutoConfig.from_pretrained(hyperparams.weights_cache_dir) + config._attn_implementation = "flash_attention_2" + load_kwargs = {"config": config, "torch_dtype": torch.bfloat16} + if int(os.environ.get("WORLD_SIZE", "1")) == 1 and torch.cuda.device_count() > 1: + load_kwargs["device_map"] = "auto" + + model = AutoModelForCausalLM.from_pretrained(hyperparams.weights_cache_dir, **load_kwargs) + if not hasattr(model, "hf_device_map"): + model = model.cuda() + model.config.use_cache = False + return model + + +def _enable_fused_mxfp8_grouped_mlp() -> None: + """Improvement 3: enable the fused ``ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8`` and + backward kernel in the installed TE without recompiling. + + ``NVTE_CUTEDSL_FUSED_GROUPED_MLP=1`` must be set *before* + ``transformer_engine.pytorch.ops`` is imported — the fusion is registered + at TE module-import-time. ``run_finetune_ep.py`` sniffs ``--improvement 3`` + and sets the env var before importing ``utils``. + + We also (a) relax the SM-version check from ``!= 10`` to ``>= 10`` so + SM>=11 successors of B300 fire the kernel, and (b) wrap the cudnn-frontend + grouped-GEMM wrappers so the installed TE's ``c_dtype`` kwarg (dropped by + cudnn-frontend 1.23.0) is silently filtered out. + """ + os.environ["NVTE_CUTEDSL_FUSED_GROUPED_MLP"] = "1" + + import inspect + import cudnn # type: ignore + from transformer_engine.pytorch.ops.fused import forward_grouped_mlp as _fwd_mod + from transformer_engine.pytorch.ops.fused import backward_grouped_mlp as _bwd_mod + from transformer_engine.pytorch.utils import get_device_compute_capability + + def _make_is_supported(kernel_method_names): + def _is_supported(cls) -> bool: + if int(os.environ.get("NVTE_CUTEDSL_FUSED_GROUPED_MLP", "0")) <= 0: + return False + if get_device_compute_capability()[0] < 10: + return False + try: + for method_name in kernel_method_names: + getattr(cls, method_name)() + except ImportError: + return False + return True + + return _is_supported + + def _make_compat_kernel(real_callable): + accepted = set(inspect.signature(real_callable).parameters) + + def _compat(**kwargs): + for k in list(kwargs): + if k not in accepted: + kwargs.pop(k) + return real_callable(**kwargs) + + return _compat + + def _patch_kernel_method(cls, method_name, wrapper_name): + compat = _make_compat_kernel(getattr(cudnn, wrapper_name)) + + def _kernel_classmethod(_cls): + return compat + + setattr(cls, method_name, classmethod(_kernel_classmethod)) + + fwd_cls = _fwd_mod.ForwardGroupedMLP_CuTeGEMMSwiGLU_MXFP8 + bwd_cls = _bwd_mod.BackwardGroupedMLP_CuTeGEMMDSwiGLU_MXFP8 + fwd_cls.is_supported = classmethod( + _make_is_supported(("grouped_gemm_glu_kernel", "grouped_gemm_quant_kernel")) + ) + bwd_cls.is_supported = classmethod( + _make_is_supported(("grouped_gemm_dglu_kernel", "grouped_gemm_quant_kernel")) + ) + _patch_kernel_method(fwd_cls, "grouped_gemm_glu_kernel", "grouped_gemm_glu_wrapper_sm100") + _patch_kernel_method(fwd_cls, "grouped_gemm_quant_kernel", "grouped_gemm_quant_wrapper_sm100") + _patch_kernel_method(bwd_cls, "grouped_gemm_dglu_kernel", "grouped_gemm_dglu_wrapper_sm100") + _patch_kernel_method(bwd_cls, "grouped_gemm_quant_kernel", "grouped_gemm_quant_wrapper_sm100") + + +def init_te_mixtral_model(hyperparams: HyperParameters): + """Load Mixtral with TE-optimised MoE blocks.""" + ensure_model_is_downloaded(hyperparams) + + import transformer_engine.common.recipe as te_recipe + + if hyperparams.model_impl == "te_mixtral_mxfp8": + if hyperparams.mixed_precision != "mxfp8": + raise ValueError("model_impl='te_mixtral_mxfp8' requires mixed_precision='mxfp8'.") + _enable_fused_mxfp8_grouped_mlp() + from te_mixtral_mxfp8 import TEMixtralMXFP8ForCausalLM as ForCausalLM + from te_mixtral_mxfp8 import replace_params + else: + from te_mixtral import TEMixtralForCausalLM as ForCausalLM + from te_mixtral import replace_params + + base_config = AutoConfig.from_pretrained(hyperparams.weights_cache_dir) + base_config._attn_implementation = "flash_attention_2" + te_config = ForCausalLM.config_class(**base_config.to_dict()) + te_config.expert_parallel_size = hyperparams.expert_parallel_size + if hasattr(te_config, "expert_ffn_mode"): + te_config.expert_ffn_mode = hyperparams.expert_ffn_mode + + fp8_recipe = None + if hyperparams.mixed_precision == "mxfp8": + fp8_recipe = te_recipe.MXFP8BlockScaling(fp8_format=te_recipe.Format.E4M3) + te_config.layer_precision = ["fp8"] * te_config.num_hidden_layers + elif hyperparams.mixed_precision != "bf16": + raise ValueError( + f"Unsupported mixed_precision={hyperparams.mixed_precision!r}; use 'bf16' or 'mxfp8'." + ) + + local_rank = int(os.environ.get("LOCAL_RANK", "0")) + world_size = int(os.environ.get("WORLD_SIZE", "1")) + if world_size > 1: + torch.cuda.set_device(local_rank) + if not dist.is_initialized(): + dist.init_process_group(backend="nccl", init_method="env://") + if world_size % hyperparams.expert_parallel_size != 0: + raise ValueError( + f"WORLD_SIZE ({world_size}) must be a multiple of " + f"expert_parallel_size ({hyperparams.expert_parallel_size})." + ) + elif hyperparams.expert_parallel_size != 1: + raise ValueError("expert_parallel_size > 1 requires torchrun distributed launch.") + + hf_model = AutoModelForCausalLM.from_pretrained( + hyperparams.weights_cache_dir, + config=base_config, + torch_dtype=torch.bfloat16, + device_map="cpu", + ) + model = ForCausalLM(te_config, fp8_recipe=fp8_recipe).to( + device=f"cuda:{local_rank}", + dtype=torch.bfloat16, + ) + te_state_dict = model.state_dict() + replace_params(hf_model.state_dict(), te_state_dict, model.config) + missing, unexpected = model.load_state_dict(te_state_dict, strict=False) + if unexpected: + raise RuntimeError(f"Unexpected keys when loading TE state dict: {unexpected}") + non_extra_missing = [key for key in missing if not key.endswith("_extra_state")] + if non_extra_missing: + raise RuntimeError(f"Missing non-extra-state keys in TE model: {non_extra_missing}") + del hf_model + + model._te_mixtral_dp_group = None + model._te_mixtral_dp_size = 1 + + if hyperparams.expert_parallel_size > 1: + ep_size = hyperparams.expert_parallel_size + dp_size = world_size // ep_size + global_rank = dist.get_rank() + dp_group = None + if dp_size == 1: + ep_group = dist.group.WORLD + else: + # Rank layout is [DP, EP]. EP groups are contiguous ranks; DP groups + # contain the same local expert shard across DP replicas. + ep_group = None + for dp_rank in range(dp_size): + ranks = list(range(dp_rank * ep_size, (dp_rank + 1) * ep_size)) + group = dist.new_group(ranks=ranks) + if global_rank in ranks: + ep_group = group + for ep_rank in range(ep_size): + ranks = [dp_rank * ep_size + ep_rank for dp_rank in range(dp_size)] + group = dist.new_group(ranks=ranks) + if global_rank in ranks: + dp_group = group + if ep_group is None: + raise RuntimeError(f"Rank {global_rank} was not assigned to an EP group.") + model.model.set_ep_groups(ep_group=ep_group) + model._te_mixtral_dp_group = dp_group + model._te_mixtral_dp_size = dp_size + + model.config.use_cache = False + return model + + +def build_adamw(model, hyperparams: HyperParameters): + params = [param for param in model.parameters() if param.requires_grad] + use_fused = hyperparams.expert_parallel_size == 1 + return AdamW( + params=params, + lr=hyperparams.learning_rate, + fused=use_fused, + foreach=False, + ) + + +def sync_data_parallel_gradients(model) -> None: + dp_group = getattr(model, "_te_mixtral_dp_group", None) + if dp_group is None: + return + + dp_size = getattr(model, "_te_mixtral_dp_size", dist.get_world_size(dp_group)) + for param in model.parameters(): + if param.grad is None: + continue + dist.all_reduce(param.grad, op=dist.ReduceOp.SUM, group=dp_group) + param.grad.div_(dp_size) + + +def move_batch_to_device(batch, device): + if torch.is_tensor(batch): + return batch.to(device=device, non_blocking=True) + if isinstance(batch, dict): + return {key: move_batch_to_device(value, device) for key, value in batch.items()} + if isinstance(batch, tuple): + return tuple(move_batch_to_device(value, device) for value in batch) + if isinstance(batch, list): + return [move_batch_to_device(value, device) for value in batch] + return batch + + +def wrap_with_accelerator(model, hyperparams: HyperParameters): + # The TE-native MXFP8 model handles its own FP8 autocast; keep + # Accelerate's mixed_precision on bf16 to avoid double-wrapping the recipe. + use_te_mxfp8 = hyperparams.mixed_precision == "mxfp8" + accelerator_mixed_precision = "bf16" if use_te_mxfp8 else hyperparams.mixed_precision + + accelerator = Accelerator( + gradient_accumulation_steps=hyperparams.gradient_accumulation_steps, + mixed_precision=accelerator_mixed_precision, + ) + + train_dataloader = get_dataloaders(accelerator, hyperparams) + optimizer = build_adamw(model, hyperparams) + lr_scheduler = get_linear_schedule_with_warmup( + optimizer=optimizer, + num_warmup_steps=hyperparams.num_warmup_steps, + num_training_steps=hyperparams.num_warmup_steps + hyperparams.num_training_steps, + ) + + if hyperparams.expert_parallel_size > 1: + # EP path: keep the DP-aware sampler intact and manually sync DP gradients. + # The dataloader is intentionally not prepared by Accelerate, so keep + # the scheduler as a plain PyTorch scheduler to avoid stepping it once + # per process. + optimizer = accelerator.prepare(optimizer) + return accelerator, model, optimizer, train_dataloader, lr_scheduler + + if hasattr(model, "hf_device_map"): + optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + optimizer, train_dataloader, lr_scheduler + ) + return accelerator, model, optimizer, train_dataloader, lr_scheduler + + model, optimizer, train_dataloader, lr_scheduler = accelerator.prepare( + model, optimizer, train_dataloader, lr_scheduler + ) + return accelerator, model, optimizer, train_dataloader, lr_scheduler + + +def finetune_model(model, hyperparams, accelerator, train_dataloader, optimizer, lr_scheduler): + """Run a short fine-tuning loop and report median step time.""" + model.train() + total_loss = 0 + optimizer.zero_grad() + + # Cycle the dataloader so long sweeps don't hit StopIteration when + # batch * world_size * num_steps exceeds the dataset size. + def _cycle(loader): + while True: + for x in loader: + yield x + + train_dataloader = enumerate(_cycle(train_dataloader)) + + for _ in range(hyperparams.num_warmup_steps): + _, batch = next(train_dataloader) + if hyperparams.expert_parallel_size > 1: + batch = move_batch_to_device(batch, accelerator.device) + with accelerator.accumulate(model): + outputs = model(**batch) + loss = outputs.loss + total_loss += loss.detach().float() + accelerator.backward(loss) + sync_data_parallel_gradients(model) + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + step_times_ms: list[float] = [] + is_printer = int(os.environ.get("LOCAL_RANK", "0")) == 0 + torch.cuda.synchronize() + + for step_idx in range(hyperparams.num_training_steps): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + + _, batch = next(train_dataloader) + if hyperparams.expert_parallel_size > 1: + batch = move_batch_to_device(batch, accelerator.device) + + with accelerator.accumulate(model): + outputs = model(**batch) + loss = outputs.loss + total_loss += loss.detach().float() + + accelerator.backward(loss) + sync_data_parallel_gradients(model) + + optimizer.step() + lr_scheduler.step() + optimizer.zero_grad() + + end.record() + end.synchronize() + step_ms = start.elapsed_time(end) + step_times_ms.append(step_ms) + if is_printer: + print( + f"[step {step_idx + 1}/{hyperparams.num_training_steps}] {step_ms:.1f} ms", + flush=True, + ) + + accelerator.end_training() + + n = len(step_times_ms) + median_ms = sorted(step_times_ms)[n // 2] + last_ms = step_times_ms[-1] + print( + f"{n} fine-tuning steps complete!\n" + f"Median time per step: {median_ms:.0f} ms\n" + f"Last step time: {last_ms:.0f} ms" + ) + + +def run_te_mixtral_finetune(hyperparams: HyperParameters): + model = init_te_mixtral_model(hyperparams) + accelerator, model, optimizer, train_dataloader, lr_scheduler = wrap_with_accelerator( + model, hyperparams + ) + finetune_model(model, hyperparams, accelerator, train_dataloader, optimizer, lr_scheduler) + + +def run_hf_baseline_finetune(hyperparams: HyperParameters): + model = init_baseline_model(hyperparams) + accelerator, model, optimizer, train_dataloader, lr_scheduler = wrap_with_accelerator( + model, hyperparams + ) + finetune_model(model, hyperparams, accelerator, train_dataloader, optimizer, lr_scheduler) diff --git a/docs/index.rst b/docs/index.rst index 53c4b0e37..52a61b960 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -56,6 +56,7 @@ Transformer Engine documentation examples/advanced_optimizations.ipynb examples/te_llama/tutorial_accelerate_hf_llama_with_te.ipynb examples/te_gemma/tutorial_generation_gemma_with_te.ipynb + examples/te_mixtral/tutorial_accelerate_hf_mixtral_with_te.ipynb examples/onnx/onnx_export.ipynb examples/te_jax_integration.rst examples/op_fuser/op_fuser.rst From 4442134dd725ed7ae36d34e7c1f4eabb72b182b3 Mon Sep 17 00:00:00 2001 From: harry zhou <67385896+harryzhou2000@users.noreply.github.com> Date: Wed, 27 May 2026 07:54:23 +0800 Subject: [PATCH 083/180] [Common] Fix fused MoE aux loss for sequence aux loss (#3018) * [Common] Allow expanded columns in fused MoE aux loss Signed-off-by: Harry Zhou * [PyTorch] Cover expanded columns in fused MoE aux loss test Signed-off-by: Harry Zhou * [Common] Document sequence aux loss column expansion Signed-off-by: Harry Zhou * [PyTorch] Scale fused aux loss tolerance by column count Signed-off-by: Harry Zhou --------- Signed-off-by: Harry Zhou --- tests/pytorch/test_fused_router.py | 13 ++++++++----- .../common/fused_router/fused_moe_aux_loss.cu | 9 ++++++--- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/tests/pytorch/test_fused_router.py b/tests/pytorch/test_fused_router.py index 274a35b81..f54d16abe 100644 --- a/tests/pytorch/test_fused_router.py +++ b/tests/pytorch/test_fused_router.py @@ -414,17 +414,20 @@ def test_fused_scores_for_aux_loss(dtype, num_tokens, num_experts, topk, score_f @pytest.mark.parametrize("num_tokens", [2048, 7168, 14234]) @pytest.mark.parametrize("num_experts", [1024, 256, 128, 32]) @pytest.mark.parametrize("topk", [4, 32]) -def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk): +@pytest.mark.parametrize("expert_multiplier", [1, 2]) +def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk, expert_multiplier): if topk >= num_experts: pytest.skip(f"topk ({topk}) >= num_experts ({num_experts})") + # Sequence aux loss batches independent sequences along the expert dimension. + num_cols = num_experts * expert_multiplier # Construct the special probs to avoid inf in the sigmoid function offset = torch.arange(-num_tokens // 2, num_tokens // 2, dtype=dtype, device="cuda") * 1e-4 - probs = torch.arange(-num_experts // 2, num_experts // 2, device="cuda", dtype=dtype) * 1e-2 + probs = torch.arange(-num_cols // 2, num_cols // 2, device="cuda", dtype=dtype) * 1e-2 probs = probs.unsqueeze(0).repeat(num_tokens, 1) + offset.unsqueeze(1) - probs = probs.view(num_tokens, num_experts) + probs = probs.view(num_tokens, num_cols) probs.requires_grad = True - tokens_per_expert = torch.randint(1, 1000, (num_experts,), device="cuda", dtype=torch.int32) + tokens_per_expert = torch.randint(1, 1000, (num_cols,), device="cuda", dtype=torch.int32) coeff = 0.01 probs_clone = deepcopy(probs) @@ -448,7 +451,7 @@ def test_fused_moe_aux_loss(dtype, num_tokens, num_experts, topk): coeff=coeff, ) - atol, rtol = _get_tolerances(dtype, num_experts) + atol, rtol = _get_tolerances(dtype, num_cols) torch.testing.assert_close(aux_loss, aux_loss_fused, atol=atol, rtol=rtol) # Backward diff --git a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu index 7e516af97..cc5e5e3bc 100644 --- a/transformer_engine/common/fused_router/fused_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_moe_aux_loss.cu @@ -87,8 +87,11 @@ void fused_moe_aux_loss_forward_kernel_launcher(const DataType* probs, int num_cols, int topk, float coeff, DataType* aux_loss, float* Coeff_buf, cudaStream_t stream) { - NVTE_CHECK(num_experts == num_cols, "Number of experts (", num_experts, - ") must be equal to number of input columns (", num_cols, ")."); + NVTE_CHECK(num_cols > 0, "num_cols must be positive, got ", num_cols); + NVTE_CHECK(num_experts > 0, "num_experts must be positive, got ", num_experts); + // Sequence aux loss batches independent sequences along the expert dimension. + NVTE_CHECK(num_cols % num_experts == 0, "Number of input columns (", num_cols, + ") must be a multiple of number of experts (", num_experts, ")."); // Round up to a multiple of warp size for correct warp shuffles. const int block_size = ((std::min(1024, num_cols) + static_cast(kThreadsPerWarp) - 1) / @@ -98,7 +101,7 @@ void fused_moe_aux_loss_forward_kernel_launcher(const DataType* probs, // One CompType per thread in shared memory. const size_t smem_size = block_size * sizeof(CompType); - check_shared_memory_capacity_num_experts(smem_size, num_experts); + check_shared_memory_capacity_num_experts(smem_size, num_cols); // Compute final coefficient and zero the float accumulator (Coeff_buf[1]) before launch. const float C_coeff = (num_experts * coeff) / topk / total_num_tokens / total_num_tokens; From be37e9b74d484576bca8147a1e6fbe3299da8537 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Wed, 27 May 2026 02:01:59 +0200 Subject: [PATCH 084/180] [common] Grouped gemm update - nvfp4 for blackwell and fp8 blockwise hopper (#2971) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * code init Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * code drop Signed-off-by: Pawel Gadzinski * Remove redundant nvte_set/get_grouped_tensor_swizzled_scales Use existing nvte_set_grouped_tensor_param with kNVTEGroupedWithGEMMSwizzledScales instead of the dedicated set/get functions. Signed-off-by: Pawel Gadzinski * Add Hopper support for grouped GEMM and refactor cuBLAS version checks - Add CUBLAS_NVFP4_GROUPED_GEMM_VERSION and CUBLAS_FP8_BLOCK_GROUPED_GEMM_VERSION macros (13.4+) - Update check_grouped_gemm_requirements to allow SM90 with cuBLAS 13.4+ - Refactor execute_grouped_gemm to use GroupedGemmConfig struct - Add divisibility-by-128 validation for FP8 block scaling in setup kernel and quantizer - Support scalar alpha/beta for Hopper (no per-group alpha/beta) - Expose get_grouped_gemm_setup_workspace_size to PyTorch via pybind - Update PyTorch tests to run grouped GEMM on Hopper with cuBLAS 13.4+ Signed-off-by: Pawel Gadzinski Made-with: Cursor * Add NVFP4 support for discrete-input grouped GEMM and skip FP8 tensor scaling tests on Hopper Extend nvte_grouped_gemm_with_discrete_inputA to handle NVFP4 (Float4E2M1) inputs: accept kFloat4E2M1 dtype, propagate scale_inv pointers, collect contiguous amax from discrete tensors, and enforce swizzled-scales checks for NVFP4 alongside MXFP8. Also add GTEST_SKIP for FP8 tensor scaling grouped GEMM on Hopper since cuBLAS does not support it there. Signed-off-by: Pawel Gadzinski Made-with: Cursor * Add alignment assertions for MXFP8/NVFP4 scale offsets in grouped GEMM tests The setup kernel computes per-tensor scale pointers as data_offset / block_size, which assumes no padding in the scale buffer. This is only correct when first_dim % 128 == 0 and last_dim % 128 == 0 (MXFP8) or last_dim % 64 == 0 (NVFP4). Add explicit assertions in build_grouped_tensor to catch any future test shapes that violate this. Signed-off-by: Pawel Gadzinski Made-with: Cursor * Fix grouped GEMM: NVFP4 columnwise transa=N + relax MXFP8 alignment for swizzle tests cublaslt_grouped_gemm.cu: - Fix incorrect handling of NVFP4/MXFP8 columnwise data in build_grouped_gemm_multi_inputA_args by adding a swap_dims flag consistent with choose_grouped_operand_storage. Use A_sel.trans (post-flip) for gemm_config.avg_k so K is selected from the correct dim with discrete A_list. tests/cpp/test_common.{h,cu}: - Add enforce_grouped_gemm_alignment parameter (default true) to build_grouped_tensor; the MXFP8/NVFP4 first/last_dim 128/64 alignment asserts are only relevant for the grouped GEMM setup kernel, so callers that bypass it (swizzle/unswizzle) opt out. tests/cpp/operator/test_swizzle.cu: - Pass enforce_grouped_gemm_alignment=false to build_grouped_tensor in MXFP8 swizzle/unswizzle/roundtrip tests, which intentionally exercise non-padded shapes. tests/cpp/operator/test_grouped_gemm.cu: - Sync GPU/cuBLAS skip rules across all 3 sub-tests, add cudaDeviceSynchronize() after nvte_multi_tensor_gemm reference for defensive sync, and skip NVFP4 + AllDifferent in all 3 sub-tests due to a known flaky bug in the nvte_multi_tensor_gemm reference. Signed-off-by: Pawel Gadzinski Made-with: Cursor * Clarify swap_dims comment in build_grouped_gemm_multi_inputA_args Signed-off-by: Pawel Gadzinski Made-with: Cursor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix grouped GEMM scale_inv offsets for NVFP4 and FP8 block scaling Apply the same fix as upstream PR #2954 (MXFP8 unaligned dims) to the analogous NVFP4 / FP8 block scaling paths in setup_grouped_gemm_kernel. Background: cuBLAS grouped GEMM expects each expert's scale_inv to live at a specific offset in the contiguous grouped buffer. The quantizer allocates each per-expert scale_inv tensor padded to the layout cuBLAS needs (swizzled 128x4 for MX/NV; ceildiv(., 128) x roundup(., 4) for block scaling). The setup kernel was computing these offsets as data_offset / block_size for everything except MXFP8 — silently correct when dims align to 128, but pointing at the middle of the previous expert's scale tile when they do not. In MoE forward this is reachable through variable per-expert token counts. Add three device helpers mirroring compute_grouped_tensor_mxfp8_- scale_inv_offset: - compute_grouped_tensor_nvfp4_scale_inv_offset - compute_grouped_tensor_block_1d_scale_inv_offset - compute_grouped_tensor_block_2d_scale_inv_offset Each sums the same padded per-tensor sizes the quantizer uses at alloc time (Float8BlockQuantizer::get_scale_shape, NVFP4Quantizer::get_scale_- shape). NVFP4 columnwise data is set up via use_columnwise(swap_dims=true), so sel.shape is already pre-transposed for that recipe — the rowwise formula on (first, last) recovers the colwise alloc. For block scaling the formula depends on the canonical orientation, so propagate a new swap_dims field on GroupedOperandSelection and pass effective_rowwise (sel.rowwise || sel.swap_dims) into the kernel. MXFP8 is invariant under this change because swap_dims is always false there and its helper's byte count is invariant under the rowwise flag anyway. Test: add ShapeCase::kUnalignedAllSame with (M, N, K) = (160, 288, 416) — all multiples of 32/16 (per-recipe block size) but none multiples of 128, so each expert's scale tile is padded. Exercise it across MXFP8 / NVFP4 / FP8 block scaling and the three transpose configs that match the existing parameter grid. Relax build_grouped_tensor's defensive %128 / %64 alignment assertions to %32 / %16 (block-size only), which is the actual quantizer requirement now that the offset arithmetic no longer assumes zero padding. Co-authored-by: Claude Opus 4.7 (1M context) Signed-off-by: Pawel Gadzinski * Relax NVFP4 amax contiguity; consolidate scale_inv offset helpers; test cleanup Production: - nvte_grouped_gemm_with_discrete_inputA no longer requires per-expert amax buffers to be contiguous. Add `amax_ptrs[kMaxGroups]` to MultiTensorGroupGemmInputArgs and read each tensor's amax via indirection in setup_grouped_gemm_kernel (mirrors the existing scale_inv_ptrs pattern). The launcher enables the NVFP4 alpha computation when amax is available from either source. - Consolidate four near-identical compute_grouped_tensor_{mxfp8,nvfp4,block_1d,block_2d}_scale_inv_offset into a single template `compute_grouped_scale_inv_offset` and collapse the A/B recipe-switch in setup_grouped_gemm_kernel into a local `fill_scale_ptr` lambda. Tests: - Drop the per-test amax staging workaround in run_grouped_gemm_discrete_in_case (no longer needed after the contiguity relax). - Fix amax management in make_nvfp4_operand: copy values into result's own amax buffers instead of aliasing pointers (prevents double-free). - Extract the three duplicated cuBLAS-version/compute-capability skip blocks into a shared `grouped_gemm_skip_reason` helper. Signed-off-by: Pawel Gadzinski Co-authored-by: Cursor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Remove unused float_size in GroupedGemmSetupWorkspace::from_buffers Silences -Wunused-variable (#177-D in nvcc). Signed-off-by: Pawel Gadzinski Co-authored-by: Cursor * Fix Hopper grouped GEMM alpha beta handling Signed-off-by: Pawel Gadzinski * fix Signed-off-by: Pawel Gadzinski * Address code review: NVFP4 amax check, swap_dims default, test refactor - nvte_grouped_gemm and nvte_grouped_gemm_with_discrete_out now validate per-operand amax for NVFP4 (previously silently dropped the global-scale factor when amax was missing). discrete_inputA path also checks B's amax. - Remove unused ShapeCase::kUnalignedAllSameNVFP4 enum and its comment. - OperandStorageChoice::swap_dims now defaults to false; rowwise returns no longer pass spurious swap_dims=true. - Unify GroupedGemmSetupWorkspace layout: from_buffers(nullptr, n) returns the total byte count, and required_setup_size derives its result from it so the layout cannot drift between the two. - test_common.cu: consolidate the three gather_*_scales lambdas into a single gather_scale_inv(bytes_per_elem, get_shape, get_cpu_ptr) helper. - test_grouped_gemm.cu: extract make_grouped_gemm_ref / make_alpha_beta / compare_grouped_d_to_multi helpers; the three run_* variants drop from ~1029 to 774 lines with no behavior change. Signed-off-by: Pawel Gadzinski Co-authored-by: Cursor * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address code review feedback (#2971) - discrete-A path: force non_tn_fp8_ok=false for FP8 block scaling to match select_grouped_operand logic for B. - Setup workspace: explicit 16-byte base alignment NVTE_CHECK before GroupedGemmSetupWorkspace::from_buffers; matches contract every standard allocator (cudaMalloc / PyTorch / XLA) already satisfies. - Tests: remove compile-time cuBLAS gating, run-time check via cuda::cublas_version() suffices since the test doesn't call cuBLAS directly. - Tests: replace InputCase enum with std::optional; nullopt = BF16, otherwise scaling mode drives dispatch. - Tests: NVFP4 operand uses one shared BF16 input transposed via nvte_transpose for the columnwise direction (no more duplicate fillUniform with different seeds across row/col). - Tests: parametrize output dtype (BF16 default, plus FP16 and FP32 cases on BF16/FP8/NVFP4 recipes); implementation already accepts all three. - Tests: add NVFP4 2D-quantization coverage (verifies VEC16 scale layout fed to cuBLAS is unchanged vs 1D). - Tests: tighten NVFP4 alignment check from %16 to %32 (TMA requirement of the optimized BF16 quantize path), fix misleading kUnalignedAllSame comment ({160,288,416} are multiples of 32, common to MXFP8 and NVFP4). Signed-off-by: Pawel Gadzinski * Simplify make_*_operand signatures to take use_rowwise directly cuBLAS scaled-GEMM kernels all run in TN, so each operand only needs the single direction matching its transpose flag. Mapping (is_A, transposed) -> use_rowwise is uniform across MXFP8 / NVFP4 / FP8 block scaling; move it to make_grouped_gemm_ref once instead of duplicating the if/else in each helper. For NVFP4 specifically: drop nvte_transpose + the two-step intermediate tensor assembly. nvte_quantize_v2 with a single-direction NVFP4 output (rowwise XOR columnwise) directly produces what cuBLAS needs — columnwise output goes via fallback quantize kernel (rowwise-only NVFP4 quantize kernel hard-fails when output has no rowwise data, see quantize.cuh:111). Matches the production pattern used by PyTorch and JAX bindings: never allocate both NVFP4 directions on a single tensor (swizzle hard-fails when both scale_invs are set, see swizzle.cu:985). Net change: -112 lines. Signed-off-by: Pawel Gadzinski * Test infrastructure fixes for NVFP4 columnwise + skip unsupported combos - test_common.cu: move NVFP4_1D_SCALING to the DELAYED/BLOCK_SCALING switch arm that allocates the columnwise data buffer as transpose(passed_shape). Required by TE's NVFP4 convention (column-wise data is transposed-then- quantized; see Tensor::shape() in common/common.h). Was incorrectly grouped with MXFP8 (whose columnwise data shape matches the logical shape). Without this, allocating columnwise-only NVFP4 in a test wires up a scale_inv shape that disagrees with what TE's CheckScaleTensorShape expects. - libtransformer_engine.version: export transformer_engine::cuda::cublas_version so the test can read the run-time cuBLAS version (analogous to the already-exported sm_arch, sm_count, current_device). - test_grouped_gemm.cu: simplify make_*_operand signatures to take use_rowwise directly, with the (is_A, transposed) -> use_rowwise mapping centralized in make_grouped_gemm_ref. NVFP4 now uses single-direction allocation (no nvte_transpose, no two-step assembly) — same pattern as MXFP8 / FP8 block scaling, working thanks to the test_common.cu fix. - test_grouped_gemm.cu: tighten skip logic via grouped_gemm_skip_reason (now takes TestParams): * Skip NVFP4 + FP16 output (cuBLAS hard-errors in cublaslt_gemm.cu) * Skip NVFP4 2D quantization in non-TN layouts (fallback quantize path doesn't support 2D; production weight quantization always allocates both directions and hits the optimized path). Drop the BF16 + FP16 output test case (cuBLAS grouped GEMM has no algorithm for that combination, even in TN). Signed-off-by: Pawel Gadzinski * Trim verbose comments in grouped GEMM test Drop multi-line rationale comments that were conversation residue (linking back to specific other files/lines, restating GEMM TN convention, listing PR numbers) and keep just the short ones that name what the code does. Signed-off-by: Pawel Gadzinski * Fix FP8 block scaling NT/NN failures on Hopper Two related bugs in the grouped GEMM path were causing FP8 block scaling NT and NN cases to fail on Hopper for non-Mul128 dims (Mul32 tests). 1. build_grouped_gemm_multi_inputA_args (discrete-A path) used t->shape() (LOGICAL) for cuBLAS rows/cols, while the symmetric grouped path (build_grouped_tensor / select_grouped_operand) uses tensors[i]->{rowwise,columnwise}_shape() (PHYSICAL). For FP8 block columnwise the two differ (physical is transposed of logical), so args.rows became N instead of K. Switch the discrete-A path to use data.shape directly — this aligns both paths and makes swap_dims redundant for this function. 2. padded_block_{1d,2d}_scale_inv_floats had a columnwise branch that swapped first/last in the formula, but the quantizer (test_common.cu:get_scales) always uses logical dims regardless of direction. Combined with grouped meta passing transposed dims for columnwise data, the formula produced wrong per-expert scale stride — only visible for dims not divisible by 128. The unified formula ceil(last/128) * roundup(first, 4) is correct for both directions because the meta swap and quantizer swap cancel out. Also remove now-redundant check_fp4_output_compat duplication and trim stale comments. Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Use logical grouped GEMM shapes Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix mixed FP8 block grouped GEMM scaling Handle per-operand FP8 block scaling modes when computing grouped scale offsets, and tighten NVFP4 validation to reject unsupported mixed or non-per-group-alpha paths. Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Trim unsupported FP8 block grouped GEMM tests Avoid test cases that require direct columnwise 2D block quantization, which is not a supported test setup path. Signed-off-by: Pawel Gadzinski * Address review comment Added validation for FP8 block scaling support in grouped GEMM. Signed-off-by: vthumbe1503 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Fix tests to relax alignment requirements for swizzling tests Signed-off-by: vthumbe1503 * Fix alignment calculation in required_setup_size Adjust required_setup_size to account for alignment. Signed-off-by: vthumbe1503 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * use nvfp4 alpha only for nvfp4 Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 * Clean up grouped GEMM tests and document alignment - Clarify scale_inv padding comments in test_grouped_gemm.cu so it is obvious that kAllSameMul32 exercises per-expert scale_inv offsets only for recipes that pad grouped scale_inv storage. - Centralize the 32 MiB cuBLAS workspace size into a single kCublasWorkspaceBytes constant in test_grouped_gemm.cu. - Drop the test-only enforce_grouped_gemm_alignment flag from build_grouped_tensor and remove all its call-site overrides; alignment is already validated inside the grouped GEMM implementation. - Add a short comment in cublaslt_grouped_gemm.cu explaining why the int and float arrays in GroupedGemmSetupWorkspace do not require an explicit align_ptr() call (cuBLAS only mandates 16-byte alignment for the pointer arrays; 4-byte natural alignment is sufficient for the int/float arrays and is preserved by the layout). Signed-off-by: Pawel Gadzinski * Post-merge cleanup in grouped GEMM helpers and swizzle tests - test_swizzle.cu: drop the enforce_grouped_gemm_alignment=false argument introduced upstream after the flag was removed from build_grouped_tensor. - cublaslt_grouped_gemm.cu (_with_discrete_inputA): remove redundant B swizzled-scales check; the same condition is already enforced inside validate_grouped_gemm_inputs({inputB}, ...) called just above. - build_grouped_gemm_multi_inputA_args: drop the requires_scale_inv parameter and derive it per-tensor from t->dtype() / t->scaling_mode. validate_grouped_gemm_multi_inputA_list already enforces uniform scaling mode, so per-iteration evaluation is safe and removes the duplicated computation at the call site. Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski Signed-off-by: vthumbe1503 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Cursor Co-authored-by: vthumbe1503 Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/cpp/operator/test_grouped_gemm.cu | 943 ++++++++------- tests/cpp/test_common.cu | 208 +++- tests/cpp/test_common.h | 4 + tests/pytorch/test_numerics.py | 7 +- transformer_engine/common/common.h | 4 + .../common/gemm/cublaslt_grouped_gemm.cu | 1051 ++++++++++++----- .../common/libtransformer_engine.version | 1 + .../common/transformer_engine.cpp | 2 + .../pytorch/cpp_extensions/gemm.py | 28 +- .../pytorch/csrc/extensions/gemm.cpp | 11 +- .../pytorch/csrc/extensions/pybind.cpp | 2 + 11 files changed, 1393 insertions(+), 868 deletions(-) diff --git a/tests/cpp/operator/test_grouped_gemm.cu b/tests/cpp/operator/test_grouped_gemm.cu index bcacb2f80..12b470346 100644 --- a/tests/cpp/operator/test_grouped_gemm.cu +++ b/tests/cpp/operator/test_grouped_gemm.cu @@ -14,6 +14,7 @@ #include #include #include +#include #include #include @@ -24,35 +25,74 @@ #include #include "../test_common.h" +#include "util/cuda_runtime.h" using namespace transformer_engine; using namespace test; namespace { -enum class InputCase { - kFP8Current, +enum class InputRecipe { kBF16, + kFP8Current, kMXFP8, + kNVFP4, + kFP8BlockScaling1D1D, + kFP8BlockScaling2D1D, + kFP8BlockScaling1D2D, }; -enum class ShapeCase { - kAllSame, - kSameFirst, - kSameLast, - kAllDifferent, -}; +inline const char* recipe_name(InputRecipe recipe) { + switch (recipe) { + case InputRecipe::kBF16: return "BF16"; + case InputRecipe::kFP8Current: return "FP8Current"; + case InputRecipe::kMXFP8: return "MXFP8"; + case InputRecipe::kNVFP4: return "NVFP4"; + case InputRecipe::kFP8BlockScaling1D1D: return "FP8BlockScaling1D1D"; + case InputRecipe::kFP8BlockScaling2D1D: return "FP8BlockScaling2D1D"; + case InputRecipe::kFP8BlockScaling1D2D: return "FP8BlockScaling1D2D"; + default: return "Unknown"; + } +} -size_t grouped_setup_workspace_size(const size_t num_tensors) { - const size_t ptr_bytes = num_tensors * sizeof(void*); - const size_t int_bytes = num_tensors * sizeof(int); - // Layout: 8 pointer arrays (A, B, C, D, alpha, beta, a_scale, b_scale) + 6 int arrays - size_t size = 8 * ptr_bytes + 6 * int_bytes; - const size_t alignment = 256; - size = ((size + alignment - 1) / alignment) * alignment; - return size; +inline bool is_fp8_block_recipe(InputRecipe recipe) { + return recipe == InputRecipe::kFP8BlockScaling1D1D || + recipe == InputRecipe::kFP8BlockScaling2D1D || + recipe == InputRecipe::kFP8BlockScaling1D2D; } +inline NVTEScalingMode a_scaling_mode(InputRecipe recipe) { + switch (recipe) { + case InputRecipe::kFP8Current: return NVTE_DELAYED_TENSOR_SCALING; + case InputRecipe::kMXFP8: return NVTE_MXFP8_1D_SCALING; + case InputRecipe::kNVFP4: return NVTE_NVFP4_1D_SCALING; + case InputRecipe::kFP8BlockScaling2D1D: return NVTE_BLOCK_SCALING_2D; + case InputRecipe::kFP8BlockScaling1D1D: + case InputRecipe::kFP8BlockScaling1D2D: return NVTE_BLOCK_SCALING_1D; + case InputRecipe::kBF16: return NVTE_DELAYED_TENSOR_SCALING; + default: return NVTE_DELAYED_TENSOR_SCALING; + } +} + +inline NVTEScalingMode b_scaling_mode(InputRecipe recipe) { + switch (recipe) { + case InputRecipe::kFP8BlockScaling1D2D: return NVTE_BLOCK_SCALING_2D; + case InputRecipe::kFP8BlockScaling2D1D: return NVTE_BLOCK_SCALING_1D; + default: return a_scaling_mode(recipe); + } +} + +// Mul128 cases use dims that are multiples of 128 - full functionality across all recipes. +// kAllSameMul32 uses dims that are multiples of 32 but not 128; for recipes with padded +// grouped scale_inv storage, it exercises per-expert scale_inv offsets. +enum class ShapeCase { + kAllSameMul128, + kSameFirstMul128, + kSameLastMul128, + kAllDifferentMul128, + kAllSameMul32, +}; + Tensor make_fp8_operand(const std::string& name, const std::vector& shape) { Tensor input_fp32(name + "_fp32", shape, DType::kFloat32); @@ -88,69 +128,86 @@ Tensor make_bf16_operand(const std::string& name, const std::vector& sha return t; } -// Creates an MXFP8 operand with the correct data layout for GEMM. -// MXFP8 GEMM requirements (scales are along K dimension): -// A transposed -> needs rowwise data/scales -// A non-transposed -> needs columnwise data/scales -// B transposed -> needs columnwise data/scales -// B non-transposed -> needs rowwise data/scales +// Creates an MXFP8 operand with the given single direction (scales along K dimension). Tensor make_mxfp8_operand(const std::string& name, const std::vector& shape, - bool is_A, bool transposed) { - // Determine which data layout we need - bool use_rowwise, use_colwise; - if (is_A) { - // A: transposed -> rowwise, non-transposed -> columnwise - use_rowwise = transposed; - use_colwise = !transposed; - } else { - // B: transposed -> columnwise, non-transposed -> rowwise (opposite of A!) - use_rowwise = !transposed; - use_colwise = transposed; - } - - // Create BF16 input with random data + bool use_rowwise) { Tensor input_bf16(name + "_bf16", shape, DType::kBFloat16); fillUniform(&input_bf16); - // Create MXFP8 tensor with only the required data layout - Tensor mxfp8(name, shape, TypeInfo::dtype, use_rowwise, use_colwise, + Tensor mxfp8(name, shape, TypeInfo::dtype, use_rowwise, !use_rowwise, NVTE_MXFP8_1D_SCALING); - - // Quantize BF16 -> MXFP8 nvte_quantize(input_bf16.data(), mxfp8.data(), 0); - // Create output tensor for swizzled scales (same data shape, same layout) Tensor mxfp8_swizzled(name + "_swizzled", shape, TypeInfo::dtype, - use_rowwise, use_colwise, NVTE_MXFP8_1D_SCALING); + use_rowwise, !use_rowwise, NVTE_MXFP8_1D_SCALING); mxfp8_swizzled.set_with_gemm_swizzled_scales(true); // Must be set BEFORE swizzle call - // Copy quantized data from mxfp8 to mxfp8_swizzled - if (use_rowwise) { - size_t data_bytes = test::bytes(mxfp8.rowwise_shape(), mxfp8.dtype()); - NVTE_CHECK_CUDA(cudaMemcpy(mxfp8_swizzled.rowwise_dptr(), mxfp8.rowwise_dptr(), - data_bytes, cudaMemcpyDeviceToDevice)); - } - if (use_colwise) { - size_t data_bytes = test::bytes(mxfp8.columnwise_shape(), mxfp8.dtype()); - NVTE_CHECK_CUDA(cudaMemcpy(mxfp8_swizzled.columnwise_dptr(), mxfp8.columnwise_dptr(), - data_bytes, cudaMemcpyDeviceToDevice)); - } + const size_t data_bytes = test::bytes( + use_rowwise ? mxfp8.rowwise_shape() : mxfp8.columnwise_shape(), mxfp8.dtype()); + void* dst = use_rowwise ? mxfp8_swizzled.rowwise_dptr() : mxfp8_swizzled.columnwise_dptr(); + void* src = use_rowwise ? mxfp8.rowwise_dptr() : mxfp8.columnwise_dptr(); + NVTE_CHECK_CUDA(cudaMemcpy(dst, src, data_bytes, cudaMemcpyDeviceToDevice)); - // Swizzle scales for GEMM nvte_swizzle_scaling_factors(mxfp8.data(), mxfp8_swizzled.data(), 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + return mxfp8_swizzled; +} + +// Creates an NVFP4 operand with the given single direction, swizzled scales. +Tensor make_nvfp4_operand(const std::string& name, const std::vector& shape, + bool use_rowwise) { + Tensor input_bf16(name + "_bf16", shape, DType::kBFloat16); + fillUniform(&input_bf16); - // Sync to ensure operations are complete + Tensor nvfp4(name, shape, DType::kFloat4E2M1, use_rowwise, !use_rowwise, + NVTE_NVFP4_1D_SCALING); + QuantizationConfigWrapper quant_config; + nvte_quantize_v2(input_bf16.data(), nvfp4.data(), quant_config, 0); + + Tensor nvfp4_sw(name + "_sw", shape, DType::kFloat4E2M1, use_rowwise, !use_rowwise, + NVTE_NVFP4_1D_SCALING); + nvfp4_sw.set_with_gemm_swizzled_scales(true); + + // Copy quantized data + amax to swizzled tensor (swizzle only rewrites scale_inv). + const auto amax_kind = use_rowwise ? kNVTEAmax : kNVTEColumnwiseAmax; + const NVTEBasicTensor src_amax = nvte_get_tensor_param(nvfp4.data(), amax_kind); + const NVTEBasicTensor dst_amax = nvte_get_tensor_param(nvfp4_sw.data(), amax_kind); + NVTE_CHECK_CUDA(cudaMemcpy(dst_amax.data_ptr, src_amax.data_ptr, sizeof(float), + cudaMemcpyDeviceToDevice)); + const size_t data_bytes = test::bytes( + use_rowwise ? nvfp4.rowwise_shape() : nvfp4.columnwise_shape(), nvfp4.dtype()); + void* dst_data = use_rowwise ? nvfp4_sw.rowwise_dptr() : nvfp4_sw.columnwise_dptr(); + void* src_data = use_rowwise ? nvfp4.rowwise_dptr() : nvfp4.columnwise_dptr(); + NVTE_CHECK_CUDA(cudaMemcpy(dst_data, src_data, data_bytes, cudaMemcpyDeviceToDevice)); + + nvte_swizzle_scaling_factors(nvfp4.data(), nvfp4_sw.data(), 0); NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + return nvfp4_sw; +} - return mxfp8_swizzled; +// Creates an FP8 block-scaling operand with the given single direction (TN-only on Hopper). +Tensor make_fp8_block_scaling_operand(const std::string& name, const std::vector& shape, + bool use_rowwise, + NVTEScalingMode scaling_mode = NVTE_BLOCK_SCALING_1D) { + Tensor input_bf16(name + "_bf16", shape, DType::kBFloat16); + fillUniform(&input_bf16); + + Tensor fp8_bs(name, shape, TypeInfo::dtype, use_rowwise, !use_rowwise, + scaling_mode); + QuantizationConfigWrapper quant_config; + quant_config.set_force_pow_2_scales(true); + nvte_quantize_v2(input_bf16.data(), fp8_bs.data(), quant_config, 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + return fp8_bs; } struct TestParams { - InputCase input_case; + InputRecipe recipe; bool transa; bool transb; ShapeCase shape_case; bool use_null_c = false; // When true, pass nullptr for C (valid when beta=0) + DType output_dtype = DType::kBFloat16; // Implementation also accepts FP16 / FP32. }; // Returns a vector of (M, N, K) tuples for each GEMM in the group. @@ -159,113 +216,227 @@ struct TestParams { // K - reduction dimension shared between A and B std::vector> make_shapes(ShapeCase scase) { switch (scase) { - case ShapeCase::kAllSame: + case ShapeCase::kAllSameMul128: return {{128, 256, 384}, {128, 256, 384}, {128, 256, 384}}; - case ShapeCase::kSameFirst: + case ShapeCase::kSameFirstMul128: // Same M (first dim), varying N and K return {{128, 256, 384}, {128, 384, 512}, {128, 512, 640}}; - case ShapeCase::kSameLast: + case ShapeCase::kSameLastMul128: // Same N (last dim), varying M and K return {{128, 256, 384}, {256, 256, 512}, {384, 256, 640}}; - case ShapeCase::kAllDifferent: - default: + case ShapeCase::kAllDifferentMul128: return {{128, 256, 384}, {256, 384, 512}, {384, 512, 640}}; + case ShapeCase::kAllSameMul32: + default: + return {{160, 288, 416}, {160, 288, 416}, {160, 288, 416}}; } } -void run_grouped_gemm_case(const TestParams& params) { -#if CUBLAS_VERSION < 130300 - GTEST_SKIP() << "Grouped GEMM requires cuBLAS 13.3+, but compile-time cuBLAS version is " - << CUBLAS_VERSION << "."; -#else - if (getDeviceComputeCapability() < blackwellComputeCapability) { - GTEST_SKIP() << "Grouped GEMM requires Blackwell (SM100) or newer."; - } +constexpr size_t kCublasGroupedGemmVersion = 130300; // Blackwell-only grouped GEMM +constexpr size_t kCublasGroupedGemmHopperVersion = 130400; // adds Hopper support +constexpr size_t kCublasWorkspaceBytes = 32ull * 1024 * 1024; - const std::vector> shapes = make_shapes(params.shape_case); +inline std::string grouped_gemm_skip_reason(const TestParams& params) { + const size_t cublas_ver = transformer_engine::cuda::cublas_version(); + if (cublas_ver < kCublasGroupedGemmVersion) { + return "Grouped GEMM requires cuBLAS 13.3+, but run-time cuBLAS version is " + + std::to_string(cublas_ver) + "."; + } + const int32_t cc = getDeviceComputeCapability(); + const std::string cc_suffix = + "but device compute capability is " + std::to_string(cc) + "."; + if (cc < hopperComputeCapability) { + return "Grouped GEMM requires Hopper (SM90) or newer, " + cc_suffix; + } + if (cc < blackwellComputeCapability && cublas_ver < kCublasGroupedGemmHopperVersion) { + return "Grouped GEMM on Hopper (SM90) requires cuBLAS 13.4+, but run-time cuBLAS " + "version is " + std::to_string(cublas_ver) + "."; + } + if (params.recipe != InputRecipe::kBF16) { + const bool is_blackwell_plus = cc >= blackwellComputeCapability; + const bool fp8_block = is_fp8_block_recipe(params.recipe); + if (!is_blackwell_plus && !fp8_block) { + return std::string(recipe_name(params.recipe)) + + " grouped GEMM requires Blackwell (SM100) or newer, " + cc_suffix; + } + if (is_blackwell_plus && fp8_block) { + return "FP8 block scaling grouped GEMM is only supported on Hopper (SM90), " + cc_suffix; + } + if (params.recipe == InputRecipe::kNVFP4 && params.output_dtype == DType::kFloat16) { + return "NVFP4 grouped GEMM does not support FP16 output."; + } + } + return ""; +} - const size_t num_gemms = shapes.size(); +// Reference setup shared by the three run_* variants: builds A/B/D tensors per recipe, +// runs nvte_multi_tensor_gemm to fill D_multi with reference results, and keeps the +// workspaces alive (returned in the struct so callers don't have to track them). +// Output dtype comes from TestParams::output_dtype (BF16 / FP16 / FP32). +struct GroupedGemmRefSetup { + std::vector> shapes; + size_t num_gemms = 0; std::vector A_tensors; std::vector B_tensors; std::vector D_multi; + std::vector workspaces; + bool use_split_accum = false; +}; - A_tensors.reserve(num_gemms); - B_tensors.reserve(num_gemms); - D_multi.reserve(num_gemms); +inline GroupedGemmRefSetup make_grouped_gemm_ref(const TestParams& params) { + GroupedGemmRefSetup s; + s.shapes = make_shapes(params.shape_case); + s.num_gemms = s.shapes.size(); + s.A_tensors.reserve(s.num_gemms); + s.B_tensors.reserve(s.num_gemms); + s.D_multi.reserve(s.num_gemms); + + for (size_t i = 0; i < s.num_gemms; ++i) { + const auto [M, N, K] = s.shapes[i]; + const std::vector a_shape = + params.transa ? std::vector{N, K} : std::vector{K, N}; + const std::vector b_shape = + params.transb ? std::vector{K, M} : std::vector{M, K}; + if (params.recipe == InputRecipe::kBF16) { + s.A_tensors.emplace_back(make_bf16_operand("A" + std::to_string(i), a_shape)); + s.B_tensors.emplace_back(make_bf16_operand("B" + std::to_string(i), b_shape)); + } else { + const bool a_use_rowwise = params.transa; + const bool b_use_rowwise = !params.transb; + switch (params.recipe) { + case InputRecipe::kFP8Current: + s.A_tensors.emplace_back(make_fp8_operand("A" + std::to_string(i), a_shape)); + s.B_tensors.emplace_back(make_fp8_operand("B" + std::to_string(i), b_shape)); + break; + case InputRecipe::kMXFP8: + s.A_tensors.emplace_back(make_mxfp8_operand("A" + std::to_string(i), a_shape, + a_use_rowwise)); + s.B_tensors.emplace_back(make_mxfp8_operand("B" + std::to_string(i), b_shape, + b_use_rowwise)); + break; + case InputRecipe::kNVFP4: + s.A_tensors.emplace_back(make_nvfp4_operand("A" + std::to_string(i), a_shape, + a_use_rowwise)); + s.B_tensors.emplace_back(make_nvfp4_operand("B" + std::to_string(i), b_shape, + b_use_rowwise)); + break; + case InputRecipe::kFP8BlockScaling1D1D: + case InputRecipe::kFP8BlockScaling2D1D: + case InputRecipe::kFP8BlockScaling1D2D: + s.A_tensors.emplace_back(make_fp8_block_scaling_operand("A" + std::to_string(i), + a_shape, a_use_rowwise, + a_scaling_mode(params.recipe))); + s.B_tensors.emplace_back(make_fp8_block_scaling_operand("B" + std::to_string(i), + b_shape, b_use_rowwise, + b_scaling_mode(params.recipe))); + break; + default: + NVTE_ERROR("Unsupported scaling mode in grouped GEMM test: " + + std::string(recipe_name(params.recipe))); + } + } + s.D_multi.emplace_back(Tensor("D_multi" + std::to_string(i), + std::vector{M, N}, params.output_dtype)); + } - for (size_t i = 0; i < num_gemms; ++i) { - const auto [M, N, K] = shapes[i]; - const std::vector a_shape = params.transa ? std::vector{N, K} - : std::vector{K, N}; - const std::vector b_shape = params.transb ? std::vector{K, M} - : std::vector{M, K}; - switch (params.input_case) { - case InputCase::kFP8Current: { - A_tensors.emplace_back(make_fp8_operand("A" + std::to_string(i), a_shape)); - B_tensors.emplace_back(make_fp8_operand("B" + std::to_string(i), b_shape)); + // FP8 block scaling requires split accumulator (no fast accumulation). + s.use_split_accum = is_fp8_block_recipe(params.recipe); + + std::vector A_ptrs(s.num_gemms), B_ptrs(s.num_gemms), D_ptrs(s.num_gemms); + std::vector workspace_ptrs(s.num_gemms, nullptr); + std::vector bias_ptrs(s.num_gemms, nullptr), gelu_ptrs(s.num_gemms, nullptr); + s.workspaces.reserve(s.num_gemms); + for (size_t i = 0; i < s.num_gemms; ++i) { + A_ptrs[i] = s.A_tensors[i].data(); + B_ptrs[i] = s.B_tensors[i].data(); + D_ptrs[i] = s.D_multi[i].data(); + s.workspaces.emplace_back(Tensor("workspace" + std::to_string(i), + std::vector{kCublasWorkspaceBytes}, DType::kByte)); + workspace_ptrs[i] = s.workspaces.back().data(); + } + nvte_multi_tensor_gemm(A_ptrs.data(), B_ptrs.data(), D_ptrs.data(), bias_ptrs.data(), + gelu_ptrs.data(), static_cast(s.num_gemms), + params.transa, params.transb, false, workspace_ptrs.data(), + false, s.use_split_accum, 0, 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + return s; +} + +// Allocate and initialize alpha/beta tensors for grouped GEMM. +// Hopper requires a single shared scalar; Blackwell+ uses per-matrix scalars. +struct AlphaBetaTensors { + Tensor alpha; + Tensor beta; +}; + +inline AlphaBetaTensors make_alpha_beta(size_t num_gemms) { + const int32_t cc = getDeviceComputeCapability(); + const size_t n = cc < blackwellComputeCapability ? 1 : num_gemms; + AlphaBetaTensors ab{Tensor("alpha", std::vector{n}, DType::kFloat32), + Tensor("beta", std::vector{n}, DType::kFloat32)}; + std::vector a(n, 1.f); + std::vector b(n, 0.f); + NVTE_CHECK_CUDA(cudaMemcpy(ab.alpha.rowwise_dptr(), a.data(), n * sizeof(float), + cudaMemcpyHostToDevice)); + NVTE_CHECK_CUDA(cudaMemcpy(ab.beta.rowwise_dptr(), b.data(), n * sizeof(float), + cudaMemcpyHostToDevice)); + return ab; +} + +// Compare each tensor inside a grouped D buffer (with per-tensor offsets) against the +// reference D_multi[i] tensors. +inline void compare_grouped_d_to_multi( + const GroupedBuffers& grouped_D, + const std::vector>& shapes, + std::vector& D_multi, const char* tag) { + for (size_t i = 0; i < shapes.size(); ++i) { + Tensor grouped_split("grouped_D" + std::to_string(i), + std::vector{static_cast(std::get<0>(shapes[i])), + static_cast(std::get<1>(shapes[i]))}, + D_multi[i].dtype()); + const size_t offset_bytes = + static_cast(grouped_D.offsets_host[i]) * grouped_D.elem_size; + NVTE_CHECK_CUDA(cudaMemcpy(grouped_split.rowwise_dptr(), + static_cast(grouped_D.get_data()) + offset_bytes, + grouped_D.tensor_bytes[i], cudaMemcpyDeviceToDevice)); + grouped_split.to_cpu(); + D_multi[i].to_cpu(); + auto [atol, rtol] = getTolerances(D_multi[i].dtype()); + switch (D_multi[i].dtype()) { + case DType::kBFloat16: + compareResults(tag, grouped_split, D_multi[i].rowwise_cpu_dptr(), true, atol, rtol); break; - } - case InputCase::kBF16: { - A_tensors.emplace_back(make_bf16_operand("A" + std::to_string(i), a_shape)); - B_tensors.emplace_back(make_bf16_operand("B" + std::to_string(i), b_shape)); + case DType::kFloat16: + compareResults(tag, grouped_split, D_multi[i].rowwise_cpu_dptr(), true, atol, rtol); break; - } - case InputCase::kMXFP8: { - A_tensors.emplace_back(make_mxfp8_operand("A" + std::to_string(i), a_shape, - /*is_A=*/true, params.transa)); - B_tensors.emplace_back(make_mxfp8_operand("B" + std::to_string(i), b_shape, - /*is_A=*/false, params.transb)); + case DType::kFloat32: + compareResults(tag, grouped_split, D_multi[i].rowwise_cpu_dptr(), true, atol, rtol); break; - } + default: + NVTE_ERROR("Unsupported D dtype in test: " + + std::to_string(static_cast(D_multi[i].dtype()))); } - D_multi.emplace_back(Tensor("D_multi" + std::to_string(i), - std::vector{M, N}, - DType::kBFloat16)); } +} - std::vector A_ptrs(num_gemms); - std::vector B_ptrs(num_gemms); - std::vector D_ptrs(num_gemms); - std::vector workspaces(num_gemms); - std::vector workspace_ptrs(num_gemms, nullptr); - std::vector A_views; - std::vector B_views; +void run_grouped_gemm_case(const TestParams& params) { + if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) { + GTEST_SKIP() << reason; + } + auto ref = make_grouped_gemm_ref(params); + const auto& shapes = ref.shapes; + const size_t num_gemms = ref.num_gemms; + + std::vector A_views, B_views; A_views.reserve(num_gemms); B_views.reserve(num_gemms); - - // Empty bias/gelu arrays for nvte_multi_tensor_gemm (no epilogues) - std::vector bias_ptrs(num_gemms, nullptr); - std::vector gelu_ptrs(num_gemms, nullptr); - - const size_t cublas_ws_bytes = 32ull * 1024 * 1024; - for (size_t i = 0; i < num_gemms; ++i) { - A_ptrs[i] = A_tensors[i].data(); - B_ptrs[i] = B_tensors[i].data(); - D_ptrs[i] = D_multi[i].data(); - workspaces[i] = Tensor("workspace" + std::to_string(i), std::vector{cublas_ws_bytes}, DType::kByte); - workspace_ptrs[i] = workspaces[i].data(); - A_views.push_back(&A_tensors[i]); - B_views.push_back(&B_tensors[i]); + A_views.push_back(&ref.A_tensors[i]); + B_views.push_back(&ref.B_tensors[i]); } - nvte_multi_tensor_gemm(A_ptrs.data(), - B_ptrs.data(), - D_ptrs.data(), - bias_ptrs.data(), - gelu_ptrs.data(), - static_cast(num_gemms), - params.transa, - params.transb, - false, // grad - workspace_ptrs.data(), - false, // accumulate - false, // use_split_accumulator - 0, // sm_count - 0); - - GroupedBuffers grouped_A = build_grouped_tensor(A_views, A_tensors[0].scaling_mode()); - GroupedBuffers grouped_B = build_grouped_tensor(B_views, B_tensors[0].scaling_mode()); + GroupedBuffers grouped_A = build_grouped_tensor(A_views, ref.A_tensors[0].scaling_mode()); + GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode()); std::vector C_tensors; std::vector D_group_tensors; @@ -277,11 +448,11 @@ void run_grouped_gemm_case(const TestParams& params) { if (!params.use_null_c) { C_tensors.emplace_back(Tensor("C" + std::to_string(i), std::vector{static_cast(M), static_cast(N)}, - DType::kBFloat16)); + params.output_dtype)); } D_group_tensors.emplace_back(Tensor("D_group" + std::to_string(i), std::vector{static_cast(M), static_cast(N)}, - DType::kBFloat16)); + params.output_dtype)); NVTE_CHECK_CUDA(cudaMemset(D_group_tensors.back().rowwise_dptr(), 0, bytes(D_group_tensors.back().rowwise_shape(), D_group_tensors.back().dtype()))); } @@ -299,152 +470,44 @@ void run_grouped_gemm_case(const TestParams& params) { } GroupedBuffers grouped_D = build_grouped_tensor(D_views, NVTE_DELAYED_TENSOR_SCALING); - // Per-matrix alpha/beta (all 1.0 and 0.0 respectively) - Tensor alpha_tensor("alpha", std::vector{num_gemms}, DType::kFloat32); - Tensor beta_tensor("beta", std::vector{num_gemms}, DType::kFloat32); - std::vector alpha_vals(num_gemms, 1.f); - std::vector beta_vals(num_gemms, 0.f); - NVTE_CHECK_CUDA(cudaMemcpy(alpha_tensor.rowwise_dptr(), alpha_vals.data(), - num_gemms * sizeof(float), cudaMemcpyHostToDevice)); - NVTE_CHECK_CUDA(cudaMemcpy(beta_tensor.rowwise_dptr(), beta_vals.data(), - num_gemms * sizeof(float), cudaMemcpyHostToDevice)); - - const size_t setup_ws_bytes = grouped_setup_workspace_size(num_gemms); - Tensor setup_ws("setup_ws", std::vector{setup_ws_bytes}, DType::kByte); - Tensor cublas_ws("cublas_ws", std::vector{cublas_ws_bytes}, DType::kByte); - - nvte_grouped_gemm(grouped_A.get_handle(), - params.transa, - grouped_B.get_handle(), - params.transb, - params.use_null_c ? nullptr : grouped_C->get_handle(), - grouped_D.get_handle(), - alpha_tensor.data(), - beta_tensor.data(), - setup_ws.data(), - cublas_ws.data(), - nullptr, // config (use defaults) - 0); - NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + AlphaBetaTensors ab = make_alpha_beta(num_gemms); - // Compare results - for (size_t i = 0; i < num_gemms; ++i) { - Tensor grouped_split("grouped_D" + std::to_string(i), - std::vector{static_cast(std::get<0>(shapes[i])), - static_cast(std::get<1>(shapes[i]))}, - D_multi[i].dtype()); - const size_t offset_bytes = static_cast(grouped_D.offsets_host[i]) * grouped_D.elem_size; - NVTE_CHECK_CUDA(cudaMemcpy(grouped_split.rowwise_dptr(), - static_cast(grouped_D.get_data()) + offset_bytes, - grouped_D.tensor_bytes[i], - cudaMemcpyDeviceToDevice)); - grouped_split.to_cpu(); - D_multi[i].to_cpu(); - auto [atol, rtol] = getTolerances(D_multi[i].dtype()); - compareResults("grouped_vs_multi", - grouped_split, - D_multi[i].rowwise_cpu_dptr(), - true, - atol, - rtol); - } -#endif // CUBLAS_VERSION >= 130300 -} + const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms); + Tensor setup_ws("setup_ws", std::vector{setup_ws_bytes}, DType::kByte); + Tensor cublas_ws("cublas_ws", std::vector{kCublasWorkspaceBytes}, DType::kByte); -void run_grouped_gemm_discrete_out_case(const TestParams& params) { -#if CUBLAS_VERSION < 130300 - GTEST_SKIP() << "Grouped GEMM requires cuBLAS 13.3+, but compile-time cuBLAS version is " - << CUBLAS_VERSION << "."; -#else - if (getDeviceComputeCapability() < blackwellComputeCapability) { - GTEST_SKIP() << "Grouped GEMM requires Blackwell (SM100) or newer."; + GroupedMatmulConfigWrapper grouped_config; + if (ref.use_split_accum) { + grouped_config.set_use_split_accumulator(true); } - const std::vector> shapes = make_shapes(params.shape_case); - - const size_t num_gemms = shapes.size(); - std::vector A_tensors; - std::vector B_tensors; - std::vector D_multi; + nvte_grouped_gemm(grouped_A.get_handle(), params.transa, grouped_B.get_handle(), params.transb, + params.use_null_c ? nullptr : grouped_C->get_handle(), grouped_D.get_handle(), + ab.alpha.data(), ab.beta.data(), setup_ws.data(), cublas_ws.data(), + grouped_config, 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); - A_tensors.reserve(num_gemms); - B_tensors.reserve(num_gemms); - D_multi.reserve(num_gemms); + compare_grouped_d_to_multi(grouped_D, shapes, ref.D_multi, "grouped_vs_multi"); +} - for (size_t i = 0; i < num_gemms; ++i) { - const auto [M, N, K] = shapes[i]; - const std::vector a_shape = params.transa ? std::vector{N, K} - : std::vector{K, N}; - const std::vector b_shape = params.transb ? std::vector{K, M} - : std::vector{M, K}; - switch (params.input_case) { - case InputCase::kFP8Current: { - A_tensors.emplace_back(make_fp8_operand("A" + std::to_string(i), a_shape)); - B_tensors.emplace_back(make_fp8_operand("B" + std::to_string(i), b_shape)); - break; - } - case InputCase::kBF16: { - A_tensors.emplace_back(make_bf16_operand("A" + std::to_string(i), a_shape)); - B_tensors.emplace_back(make_bf16_operand("B" + std::to_string(i), b_shape)); - break; - } - case InputCase::kMXFP8: { - A_tensors.emplace_back(make_mxfp8_operand("A" + std::to_string(i), a_shape, - /*is_A=*/true, params.transa)); - B_tensors.emplace_back(make_mxfp8_operand("B" + std::to_string(i), b_shape, - /*is_A=*/false, params.transb)); - break; - } - } - D_multi.emplace_back(Tensor("D_multi" + std::to_string(i), - std::vector{M, N}, - DType::kBFloat16)); +void run_grouped_gemm_discrete_out_case(const TestParams& params) { + if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) { + GTEST_SKIP() << reason; } + auto ref = make_grouped_gemm_ref(params); + const auto& shapes = ref.shapes; + const size_t num_gemms = ref.num_gemms; - std::vector A_ptrs(num_gemms); - std::vector B_ptrs(num_gemms); - std::vector D_ptrs(num_gemms); - std::vector workspaces(num_gemms); - std::vector workspace_ptrs(num_gemms, nullptr); - std::vector A_views; - std::vector B_views; + std::vector A_views, B_views; A_views.reserve(num_gemms); B_views.reserve(num_gemms); - - // Empty bias/gelu arrays for nvte_multi_tensor_gemm (no epilogues) - std::vector bias_ptrs(num_gemms, nullptr); - std::vector gelu_ptrs(num_gemms, nullptr); - - const size_t cublas_ws_bytes = 32ull * 1024 * 1024; - for (size_t i = 0; i < num_gemms; ++i) { - A_ptrs[i] = A_tensors[i].data(); - B_ptrs[i] = B_tensors[i].data(); - D_ptrs[i] = D_multi[i].data(); - workspaces[i] = - Tensor("workspace" + std::to_string(i), std::vector{cublas_ws_bytes}, DType::kByte); - workspace_ptrs[i] = workspaces[i].data(); - A_views.push_back(&A_tensors[i]); - B_views.push_back(&B_tensors[i]); + A_views.push_back(&ref.A_tensors[i]); + B_views.push_back(&ref.B_tensors[i]); } - nvte_multi_tensor_gemm(A_ptrs.data(), - B_ptrs.data(), - D_ptrs.data(), - bias_ptrs.data(), - gelu_ptrs.data(), - static_cast(num_gemms), - params.transa, - params.transb, - false, // grad - workspace_ptrs.data(), - false, // accumulate - false, // use_split_accumulator - 0, // sm_count - 0); - - GroupedBuffers grouped_A = build_grouped_tensor(A_views, A_tensors[0].scaling_mode()); - GroupedBuffers grouped_B = build_grouped_tensor(B_views, B_tensors[0].scaling_mode()); + GroupedBuffers grouped_A = build_grouped_tensor(A_views, ref.A_tensors[0].scaling_mode()); + GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode()); std::vector C_tensors; std::vector D_list_tensors; @@ -455,10 +518,10 @@ void run_grouped_gemm_discrete_out_case(const TestParams& params) { (void)K; if (!params.use_null_c) { C_tensors.emplace_back( - Tensor("C" + std::to_string(i), std::vector{M, N}, DType::kBFloat16)); + Tensor("C" + std::to_string(i), std::vector{M, N}, params.output_dtype)); } D_list_tensors.emplace_back( - Tensor("D_list" + std::to_string(i), std::vector{M, N}, DType::kBFloat16)); + Tensor("D_list" + std::to_string(i), std::vector{M, N}, params.output_dtype)); NVTE_CHECK_CUDA(cudaMemset(D_list_tensors.back().rowwise_dptr(), 0, bytes(D_list_tensors.back().rowwise_shape(), D_list_tensors.back().dtype()))); @@ -477,160 +540,74 @@ void run_grouped_gemm_discrete_out_case(const TestParams& params) { D_list_ptrs.push_back(D_list_tensors[i].data()); } - // Per-matrix alpha/beta (all 1.0 and 0.0 respectively) - Tensor alpha_tensor("alpha", std::vector{num_gemms}, DType::kFloat32); - Tensor beta_tensor("beta", std::vector{num_gemms}, DType::kFloat32); - std::vector alpha_vals(num_gemms, 1.f); - std::vector beta_vals(num_gemms, 0.f); - NVTE_CHECK_CUDA(cudaMemcpy(alpha_tensor.rowwise_dptr(), alpha_vals.data(), - num_gemms * sizeof(float), cudaMemcpyHostToDevice)); - NVTE_CHECK_CUDA(cudaMemcpy(beta_tensor.rowwise_dptr(), beta_vals.data(), - num_gemms * sizeof(float), cudaMemcpyHostToDevice)); - - const size_t setup_ws_bytes = grouped_setup_workspace_size(num_gemms); - Tensor setup_ws("setup_ws", std::vector{setup_ws_bytes}, DType::kByte); - Tensor cublas_ws("cublas_ws", std::vector{cublas_ws_bytes}, DType::kByte); - - nvte_grouped_gemm_with_discrete_out(grouped_A.get_handle(), - params.transa, - grouped_B.get_handle(), - params.transb, - params.use_null_c ? nullptr : C_list_ptrs.data(), - params.use_null_c ? 0 : num_gemms, - D_list_ptrs.data(), - num_gemms, - alpha_tensor.data(), - beta_tensor.data(), - setup_ws.data(), - cublas_ws.data(), - nullptr, // config (use defaults) - 0); - NVTE_CHECK_CUDA(cudaDeviceSynchronize()); + AlphaBetaTensors ab = make_alpha_beta(num_gemms); - // Compare results - for (size_t i = 0; i < num_gemms; ++i) { - D_list_tensors[i].to_cpu(); - D_multi[i].to_cpu(); - auto [atol, rtol] = getTolerances(D_multi[i].dtype()); - compareResults("grouped_list_vs_multi", - D_list_tensors[i], - D_multi[i].rowwise_cpu_dptr(), - true, - atol, - rtol); - } -#endif // CUBLAS_VERSION >= 130300 -} + const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms); + Tensor setup_ws("setup_ws", std::vector{setup_ws_bytes}, DType::kByte); + Tensor cublas_ws("cublas_ws", std::vector{kCublasWorkspaceBytes}, DType::kByte); -void run_grouped_gemm_discrete_in_case(const TestParams& params) { -#if CUBLAS_VERSION < 130300 - GTEST_SKIP() << "Grouped GEMM requires cuBLAS 13.3+, but compile-time cuBLAS version is " - << CUBLAS_VERSION << "."; -#else - if (getDeviceComputeCapability() < blackwellComputeCapability) { - GTEST_SKIP() << "Grouped GEMM requires Blackwell (SM100) or newer."; + GroupedMatmulConfigWrapper grouped_config; + if (ref.use_split_accum) { + grouped_config.set_use_split_accumulator(true); } - const std::vector> shapes = make_shapes(params.shape_case); - - const size_t num_gemms = shapes.size(); - std::vector A_tensors; - std::vector B_tensors; - std::vector D_multi; - - A_tensors.reserve(num_gemms); - B_tensors.reserve(num_gemms); - D_multi.reserve(num_gemms); + nvte_grouped_gemm_with_discrete_out( + grouped_A.get_handle(), params.transa, grouped_B.get_handle(), params.transb, + params.use_null_c ? nullptr : C_list_ptrs.data(), params.use_null_c ? 0 : num_gemms, + D_list_ptrs.data(), num_gemms, ab.alpha.data(), ab.beta.data(), setup_ws.data(), + cublas_ws.data(), grouped_config, 0); + NVTE_CHECK_CUDA(cudaDeviceSynchronize()); for (size_t i = 0; i < num_gemms; ++i) { - const auto [M, N, K] = shapes[i]; - const std::vector a_shape = params.transa ? std::vector{N, K} - : std::vector{K, N}; - const std::vector b_shape = params.transb ? std::vector{K, M} - : std::vector{M, K}; - switch (params.input_case) { - case InputCase::kFP8Current: { - A_tensors.emplace_back(make_fp8_operand("A" + std::to_string(i), a_shape)); - B_tensors.emplace_back(make_fp8_operand("B" + std::to_string(i), b_shape)); + D_list_tensors[i].to_cpu(); + ref.D_multi[i].to_cpu(); + auto [atol, rtol] = getTolerances(ref.D_multi[i].dtype()); + switch (ref.D_multi[i].dtype()) { + case DType::kBFloat16: + compareResults("grouped_list_vs_multi", D_list_tensors[i], + ref.D_multi[i].rowwise_cpu_dptr(), true, atol, rtol); break; - } - case InputCase::kBF16: { - A_tensors.emplace_back(make_bf16_operand("A" + std::to_string(i), a_shape)); - B_tensors.emplace_back(make_bf16_operand("B" + std::to_string(i), b_shape)); + case DType::kFloat16: + compareResults("grouped_list_vs_multi", D_list_tensors[i], + ref.D_multi[i].rowwise_cpu_dptr(), true, atol, rtol); break; - } - case InputCase::kMXFP8: { - A_tensors.emplace_back(make_mxfp8_operand("A" + std::to_string(i), a_shape, - /*is_A=*/true, params.transa)); - B_tensors.emplace_back(make_mxfp8_operand("B" + std::to_string(i), b_shape, - /*is_A=*/false, params.transb)); + case DType::kFloat32: + compareResults("grouped_list_vs_multi", D_list_tensors[i], + ref.D_multi[i].rowwise_cpu_dptr(), true, atol, rtol); break; - } + default: + NVTE_ERROR("Unsupported D dtype in test: " + + std::to_string(static_cast(ref.D_multi[i].dtype()))); } - D_multi.emplace_back(Tensor("D_multi" + std::to_string(i), - std::vector{M, N}, - DType::kBFloat16)); } +} + +void run_grouped_gemm_discrete_in_case(const TestParams& params) { + if (auto reason = grouped_gemm_skip_reason(params); !reason.empty()) { + GTEST_SKIP() << reason; + } + auto ref = make_grouped_gemm_ref(params); + const auto& shapes = ref.shapes; + const size_t num_gemms = ref.num_gemms; - std::vector A_ptrs(num_gemms); - std::vector B_ptrs(num_gemms); - std::vector D_ptrs(num_gemms); - std::vector workspaces(num_gemms); - std::vector workspace_ptrs(num_gemms, nullptr); - std::vector A_views; std::vector B_views; - A_views.reserve(num_gemms); B_views.reserve(num_gemms); + for (size_t i = 0; i < num_gemms; ++i) B_views.push_back(&ref.B_tensors[i]); - // Empty bias/gelu arrays for nvte_multi_tensor_gemm (no epilogues) - std::vector bias_ptrs(num_gemms, nullptr); - std::vector gelu_ptrs(num_gemms, nullptr); - - const size_t cublas_ws_bytes = 32ull * 1024 * 1024; + GroupedBuffers grouped_B = build_grouped_tensor(B_views, ref.B_tensors[0].scaling_mode()); - for (size_t i = 0; i < num_gemms; ++i) { - A_ptrs[i] = A_tensors[i].data(); - B_ptrs[i] = B_tensors[i].data(); - D_ptrs[i] = D_multi[i].data(); - workspaces[i] = - Tensor("workspace" + std::to_string(i), std::vector{cublas_ws_bytes}, DType::kByte); - workspace_ptrs[i] = workspaces[i].data(); - A_views.push_back(&A_tensors[i]); - B_views.push_back(&B_tensors[i]); - } - - nvte_multi_tensor_gemm(A_ptrs.data(), - B_ptrs.data(), - D_ptrs.data(), - bias_ptrs.data(), - gelu_ptrs.data(), - static_cast(num_gemms), - params.transa, - params.transb, - false, // grad - workspace_ptrs.data(), - false, // accumulate - false, // use_split_accumulator - 0, // sm_count - 0); - - GroupedBuffers grouped_B = build_grouped_tensor(B_views, B_tensors[0].scaling_mode()); - - std::vector C_tensors; - std::vector D_group_tensors; + std::vector C_tensors, D_group_tensors; C_tensors.reserve(num_gemms); D_group_tensors.reserve(num_gemms); for (size_t i = 0; i < num_gemms; ++i) { const auto [M, N, K] = shapes[i]; (void)K; if (!params.use_null_c) { - C_tensors.emplace_back(Tensor("C" + std::to_string(i), - std::vector{M, N}, - DType::kBFloat16)); + C_tensors.emplace_back(Tensor("C" + std::to_string(i), std::vector{M, N}, + params.output_dtype)); } D_group_tensors.emplace_back(Tensor("D_group" + std::to_string(i), - std::vector{M, N}, - DType::kBFloat16)); + std::vector{M, N}, params.output_dtype)); NVTE_CHECK_CUDA(cudaMemset(D_group_tensors.back().rowwise_dptr(), 0, bytes(D_group_tensors.back().rowwise_shape(), D_group_tensors.back().dtype()))); @@ -638,9 +615,7 @@ void run_grouped_gemm_discrete_in_case(const TestParams& params) { std::vector C_views, D_views; for (size_t i = 0; i < num_gemms; ++i) { - if (!params.use_null_c) { - C_views.push_back(&C_tensors[i]); - } + if (!params.use_null_c) C_views.push_back(&C_tensors[i]); D_views.push_back(&D_group_tensors[i]); } @@ -650,63 +625,28 @@ void run_grouped_gemm_discrete_in_case(const TestParams& params) { } GroupedBuffers grouped_D = build_grouped_tensor(D_views, NVTE_DELAYED_TENSOR_SCALING); - // Per-matrix alpha/beta (all 1.0 and 0.0 respectively) - Tensor alpha_tensor("alpha", std::vector{num_gemms}, DType::kFloat32); - Tensor beta_tensor("beta", std::vector{num_gemms}, DType::kFloat32); - std::vector alpha_vals(num_gemms, 1.f); - std::vector beta_vals(num_gemms, 0.f); - NVTE_CHECK_CUDA(cudaMemcpy(alpha_tensor.rowwise_dptr(), alpha_vals.data(), - num_gemms * sizeof(float), cudaMemcpyHostToDevice)); - NVTE_CHECK_CUDA(cudaMemcpy(beta_tensor.rowwise_dptr(), beta_vals.data(), - num_gemms * sizeof(float), cudaMemcpyHostToDevice)); - - const size_t setup_ws_bytes = grouped_setup_workspace_size(num_gemms); + AlphaBetaTensors ab = make_alpha_beta(num_gemms); + + const size_t setup_ws_bytes = nvte_get_grouped_gemm_setup_workspace_size(num_gemms); Tensor setup_ws("setup_ws", std::vector{setup_ws_bytes}, DType::kByte); - Tensor cublas_ws("cublas_ws", std::vector{cublas_ws_bytes}, DType::kByte); + Tensor cublas_ws("cublas_ws", std::vector{kCublasWorkspaceBytes}, DType::kByte); std::vector A_list_ptrs; A_list_ptrs.reserve(num_gemms); - for (size_t i = 0; i < num_gemms; ++i) { - A_list_ptrs.push_back(A_tensors[i].data()); + for (size_t i = 0; i < num_gemms; ++i) A_list_ptrs.push_back(ref.A_tensors[i].data()); + + GroupedMatmulConfigWrapper grouped_config; + if (ref.use_split_accum) { + grouped_config.set_use_split_accumulator(true); } - nvte_grouped_gemm_with_discrete_inputA(A_list_ptrs.data(), - num_gemms, - params.transa, - grouped_B.get_handle(), - params.transb, - params.use_null_c ? nullptr : grouped_C->get_handle(), - grouped_D.get_handle(), - alpha_tensor.data(), - beta_tensor.data(), - setup_ws.data(), - cublas_ws.data(), - nullptr, // config (use defaults) - 0); + nvte_grouped_gemm_with_discrete_inputA( + A_list_ptrs.data(), num_gemms, params.transa, grouped_B.get_handle(), params.transb, + params.use_null_c ? nullptr : grouped_C->get_handle(), grouped_D.get_handle(), + ab.alpha.data(), ab.beta.data(), setup_ws.data(), cublas_ws.data(), grouped_config, 0); NVTE_CHECK_CUDA(cudaDeviceSynchronize()); - // Compare results - for (size_t i = 0; i < num_gemms; ++i) { - Tensor grouped_split("grouped_D" + std::to_string(i), - std::vector{static_cast(std::get<0>(shapes[i])), - static_cast(std::get<1>(shapes[i]))}, - D_multi[i].dtype()); - const size_t offset_bytes = static_cast(grouped_D.offsets_host[i]) * grouped_D.elem_size; - NVTE_CHECK_CUDA(cudaMemcpy(grouped_split.rowwise_dptr(), - static_cast(grouped_D.get_data()) + offset_bytes, - grouped_D.tensor_bytes[i], - cudaMemcpyDeviceToDevice)); - grouped_split.to_cpu(); - D_multi[i].to_cpu(); - auto [atol, rtol] = getTolerances(D_multi[i].dtype()); - compareResults("grouped_discrete_in_vs_multi", - grouped_split, - D_multi[i].rowwise_cpu_dptr(), - true, - atol, - rtol); - } -#endif // CUBLAS_VERSION >= 130300 + compare_grouped_d_to_multi(grouped_D, shapes, ref.D_multi, "grouped_discrete_in_vs_multi"); } class GroupedGemmTest : public ::testing::TestWithParam {}; @@ -724,38 +664,87 @@ TEST_P(GroupedGemmTest, CompareWithMultiTensorGemmDiscreteIn) { } std::string MakeGroupedGemmTestName(const testing::TestParamInfo& info) { - constexpr const char* kInputNames[] = {"FP8Current", "BF16", "MXFP8"}; - constexpr const char* kShapeNames[] = {"AllSame", "SameM", "SameN", "AllDiff"}; + constexpr const char* kShapeNames[] = {"AllSameMul128", "SameMMul128", "SameNMul128", + "AllDiffMul128", "AllSameMul32"}; const std::string layout = std::string("ta") + (info.param.transa ? "T" : "N") + "tb" + (info.param.transb ? "T" : "N"); const std::string null_c = info.param.use_null_c ? "_NullC" : ""; - return std::string(kInputNames[static_cast(info.param.input_case)]) + "_" + - kShapeNames[static_cast(info.param.shape_case)] + "_" + layout + null_c; + std::string out_suffix; + switch (info.param.output_dtype) { + case DType::kBFloat16: break; // default, no suffix + case DType::kFloat16: out_suffix = "_outFP16"; break; + case DType::kFloat32: out_suffix = "_outFP32"; break; + default: out_suffix = "_outUnknown"; break; + } + return std::string(recipe_name(info.param.recipe)) + "_" + + kShapeNames[static_cast(info.param.shape_case)] + "_" + layout + null_c + out_suffix; } -// TestParams: {input_case, transa, transb, shape_case, use_null_c} +// TestParams: {recipe, transa, transb, shape_case, use_null_c} const std::vector kTestParams = { // FP8 tests (each tensor has random mean/stddev -> different scales) - {InputCase::kFP8Current, true, false, ShapeCase::kAllDifferent, false}, - {InputCase::kFP8Current, false, true, ShapeCase::kAllDifferent, false}, - {InputCase::kFP8Current, false, false, ShapeCase::kAllSame, false}, + {InputRecipe::kFP8Current, true, false, ShapeCase::kAllDifferentMul128, false}, + {InputRecipe::kFP8Current, false, true, ShapeCase::kAllDifferentMul128, false}, + {InputRecipe::kFP8Current, false, false, ShapeCase::kAllSameMul128, false}, // BF16 tests - {InputCase::kBF16, true, false, ShapeCase::kSameFirst, false}, - {InputCase::kBF16, false, true, ShapeCase::kSameLast, false}, - {InputCase::kBF16, false, false, ShapeCase::kAllSame, false}, - {InputCase::kBF16, true, true, ShapeCase::kAllDifferent, false}, + {InputRecipe::kBF16, true, false, ShapeCase::kSameFirstMul128, false}, + {InputRecipe::kBF16, false, true, ShapeCase::kSameLastMul128, false}, + {InputRecipe::kBF16, false, false, ShapeCase::kAllSameMul128, false}, + {InputRecipe::kBF16, true, true, ShapeCase::kAllDifferentMul128, false}, // Test NULL C (valid when beta=0) - {InputCase::kBF16, false, false, ShapeCase::kAllSame, true}, + {InputRecipe::kBF16, false, false, ShapeCase::kAllSameMul128, true}, // MXFP8 tests - {InputCase::kMXFP8, true, false, ShapeCase::kAllSame, false}, - {InputCase::kMXFP8, true, false, ShapeCase::kAllDifferent, false}, - {InputCase::kMXFP8, false, true, ShapeCase::kAllSame, false}, - {InputCase::kMXFP8, false, true, ShapeCase::kAllDifferent, false}, - {InputCase::kMXFP8, false, false, ShapeCase::kAllSame, false}, - {InputCase::kMXFP8, false, false, ShapeCase::kAllDifferent, false}, - {InputCase::kMXFP8, false, false, ShapeCase::kSameFirst, false}, + {InputRecipe::kMXFP8, true, false, ShapeCase::kAllSameMul128, false}, + {InputRecipe::kMXFP8, true, false, ShapeCase::kAllDifferentMul128, false}, + {InputRecipe::kMXFP8, false, true, ShapeCase::kAllSameMul128, false}, + {InputRecipe::kMXFP8, false, true, ShapeCase::kAllDifferentMul128, false}, + {InputRecipe::kMXFP8, false, false, ShapeCase::kAllSameMul128, false}, + {InputRecipe::kMXFP8, false, false, ShapeCase::kAllDifferentMul128, false}, + {InputRecipe::kMXFP8, false, false, ShapeCase::kSameFirstMul128, false}, // MXFP8 with NULL C - {InputCase::kMXFP8, true, false, ShapeCase::kAllSame, true}, + {InputRecipe::kMXFP8, true, false, ShapeCase::kAllSameMul128, true}, + // NVFP4 tests (all transpose combinations - GEMM internally forces TN) + {InputRecipe::kNVFP4, true, false, ShapeCase::kAllSameMul128, false}, + {InputRecipe::kNVFP4, true, false, ShapeCase::kAllDifferentMul128, false}, + {InputRecipe::kNVFP4, true, false, ShapeCase::kSameFirstMul128, false}, + {InputRecipe::kNVFP4, true, false, ShapeCase::kSameLastMul128, false}, + {InputRecipe::kNVFP4, false, true, ShapeCase::kAllSameMul128, false}, + {InputRecipe::kNVFP4, false, true, ShapeCase::kAllDifferentMul128, false}, + {InputRecipe::kNVFP4, false, false, ShapeCase::kAllSameMul128, false}, + {InputRecipe::kNVFP4, false, false, ShapeCase::kAllDifferentMul128, false}, + // NVFP4 with NULL C + {InputRecipe::kNVFP4, true, false, ShapeCase::kAllSameMul128, true}, + // Non-default output dtypes (BF16 covered everywhere else). + {InputRecipe::kBF16, false, false, ShapeCase::kAllSameMul128, false, + /*output_dtype=*/DType::kFloat32}, + {InputRecipe::kFP8Current, true, false, ShapeCase::kAllSameMul128, false, + /*output_dtype=*/DType::kFloat16}, + // FP8 Block Scaling tests (TN layout on Hopper, block size 128) + {InputRecipe::kFP8BlockScaling1D1D, true, false, ShapeCase::kAllSameMul128, false}, + {InputRecipe::kFP8BlockScaling1D1D, true, false, ShapeCase::kAllDifferentMul128, false}, + {InputRecipe::kFP8BlockScaling1D1D, true, false, ShapeCase::kSameFirstMul128, false}, + {InputRecipe::kFP8BlockScaling1D1D, true, false, ShapeCase::kSameLastMul128, false}, + {InputRecipe::kFP8BlockScaling1D1D, false, true, ShapeCase::kAllSameMul128, false}, + {InputRecipe::kFP8BlockScaling1D1D, false, false, ShapeCase::kAllSameMul128, false}, + // FP8 Block Scaling with NULL C + {InputRecipe::kFP8BlockScaling1D1D, true, false, ShapeCase::kAllSameMul128, true}, + // Dims multiples of 32 but not 128 exercise padded scale_inv offsets for recipes that use + // padded grouped scale_inv storage. + {InputRecipe::kMXFP8, true, false, ShapeCase::kAllSameMul32, false}, + {InputRecipe::kMXFP8, false, true, ShapeCase::kAllSameMul32, false}, + {InputRecipe::kMXFP8, false, false, ShapeCase::kAllSameMul32, false}, + {InputRecipe::kNVFP4, true, false, ShapeCase::kAllSameMul32, false}, + {InputRecipe::kNVFP4, false, true, ShapeCase::kAllSameMul32, false}, + {InputRecipe::kNVFP4, false, false, ShapeCase::kAllSameMul32, false}, + {InputRecipe::kFP8BlockScaling1D1D, true, false, ShapeCase::kAllSameMul32, false}, + {InputRecipe::kFP8BlockScaling1D1D, false, true, ShapeCase::kAllSameMul32, false}, + {InputRecipe::kFP8BlockScaling1D1D, false, false, ShapeCase::kAllSameMul32, false}, + // Mixed FP8 block scaling modes supported by cuBLASLt: 2D x 1D and 1D x 2D. + {InputRecipe::kFP8BlockScaling2D1D, true, false, ShapeCase::kAllSameMul128, false}, + {InputRecipe::kFP8BlockScaling2D1D, true, false, ShapeCase::kAllDifferentMul128, false}, + {InputRecipe::kFP8BlockScaling1D2D, true, false, ShapeCase::kAllSameMul128, false}, + {InputRecipe::kFP8BlockScaling1D2D, true, false, ShapeCase::kAllDifferentMul128, false}, + {InputRecipe::kFP8BlockScaling1D2D, false, false, ShapeCase::kAllSameMul32, false}, }; INSTANTIATE_TEST_SUITE_P(OperatorTest, diff --git a/tests/cpp/test_common.cu b/tests/cpp/test_common.cu index e35f5e029..fc41d4472 100644 --- a/tests/cpp/test_common.cu +++ b/tests/cpp/test_common.cu @@ -315,7 +315,8 @@ Tensor::Tensor(const std::string& name, switch (scaling_mode) { case NVTE_DELAYED_TENSOR_SCALING: case NVTE_BLOCK_SCALING_1D: - case NVTE_BLOCK_SCALING_2D: { + case NVTE_BLOCK_SCALING_2D: + case NVTE_NVFP4_1D_SCALING: { // Column-wise data shape is transposed if (shape.ndim > 0) { columnwise_shape_vec.emplace_back(shape.data[shape.ndim - 1]); @@ -325,8 +326,7 @@ Tensor::Tensor(const std::string& name, } break; } - case NVTE_MXFP8_1D_SCALING: - case NVTE_NVFP4_1D_SCALING: { + case NVTE_MXFP8_1D_SCALING: { // Column-wise data matches shape for (size_t i = 0; i < shape.ndim; ++i) { columnwise_shape_vec.emplace_back(shape.data[i]); @@ -1072,13 +1072,18 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, const bool has_columnwise = tensors[0]->columnwise(); NVTE_CHECK(has_rowwise || has_columnwise, "Tensors must have at least one data layout."); - const NVTEShape shape = has_rowwise ? tensors[0]->rowwise_shape() - : tensors[0]->columnwise_shape(); const DType dtype = tensors[0]->dtype(); const size_t num_tensors = tensors.size(); - const size_t elem_size = typeToNumBits(dtype) / 8; + const size_t bits_per_elem = typeToNumBits(dtype); + const bool is_sub_byte = (bits_per_elem < 8); + const size_t elem_size = is_sub_byte ? 0 : bits_per_elem / 8; GroupedBuffers grouped; - grouped.elem_size = elem_size; + grouped.elem_size = elem_size; // Only used for D output extraction (always >= 1 byte dtype) + + // Helper: convert element count to byte count (handles sub-byte types like FP4) + auto elems_to_bytes = [bits_per_elem](int64_t elems) -> size_t { + return static_cast((elems * static_cast(bits_per_elem)) / 8); + }; grouped.num_tensors = num_tensors; grouped.dtype = dtype; grouped.scaling_mode = scaling_mode; @@ -1088,12 +1093,13 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, std::vector first_dims(num_tensors); std::vector last_dims(num_tensors); for (size_t i = 0; i < num_tensors; ++i) { - const auto s = has_rowwise ? tensors[i]->rowwise_shape() - : tensors[i]->columnwise_shape(); + const auto s = tensors[i]->shape(); NVTE_CHECK(s.ndim == 2, "Grouped tensor build expects 2D tensors."); first_dims[i] = static_cast(s.data[0]); last_dims[i] = static_cast(s.data[1]); - grouped.tensor_bytes[i] = bytes(s, dtype); + const auto storage_shape = has_rowwise ? tensors[i]->rowwise_shape() + : tensors[i]->columnwise_shape(); + grouped.tensor_bytes[i] = bytes(storage_shape, dtype); } const bool same_first = std::all_of(first_dims.begin(), first_dims.end(), @@ -1107,9 +1113,14 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, // cuBLAS requires aligned pointers for vectorized loads static std::mt19937 gen(12345); std::uniform_int_distribution dist(0, 3); - // Calculate elements needed for 16-byte alignment in bytes, rounded up - const size_t align_elements = - std::max(1, (16 + elem_size - 1) / elem_size); // 16 bytes / element_size + // Calculate elements needed for 16-byte alignment + size_t align_elements; + if (is_sub_byte) { + // Sub-byte types (e.g. FP4): 16 bytes = 16*8/bits_per_elem elements + align_elements = (16 * 8) / bits_per_elem; + } else { + align_elements = std::max(1, (16 + elem_size - 1) / elem_size); + } return dist(gen) * static_cast(align_elements); }; @@ -1157,7 +1168,7 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, const int64_t total_elems = need_offsets ? (offsets[last_idx] + numel(last_idx)) : (logical_first * logical_last); - const size_t total_bytes = static_cast(total_elems) * elem_size; + const size_t total_bytes = elems_to_bytes(total_elems); NVTEGroupedTensor h = grouped.handle.get(); @@ -1167,8 +1178,8 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, if (has_rowwise) { grouped.data = cuda_alloc(total_bytes); for (size_t i = 0; i < num_tensors; ++i) { - const size_t offset_bytes = static_cast(offsets[i]) * elem_size; - NVTE_CHECK_CUDA(cudaMemcpy(static_cast(grouped.data.get()) + offset_bytes, + const size_t offset_bytes_i = elems_to_bytes(offsets[i]); + NVTE_CHECK_CUDA(cudaMemcpy(static_cast(grouped.data.get()) + offset_bytes_i, tensors[i]->rowwise_dptr(), grouped.tensor_bytes[i], cudaMemcpyDeviceToDevice)); @@ -1181,8 +1192,8 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, if (has_columnwise) { grouped.columnwise_data = cuda_alloc(total_bytes); for (size_t i = 0; i < num_tensors; ++i) { - const size_t offset_bytes = static_cast(offsets[i]) * elem_size; - NVTE_CHECK_CUDA(cudaMemcpy(static_cast(grouped.columnwise_data.get()) + offset_bytes, + const size_t offset_bytes_i = elems_to_bytes(offsets[i]); + NVTE_CHECK_CUDA(cudaMemcpy(static_cast(grouped.columnwise_data.get()) + offset_bytes_i, tensors[i]->columnwise_dptr(), grouped.tensor_bytes[i], cudaMemcpyDeviceToDevice)); @@ -1221,6 +1232,33 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, nvte_set_grouped_tensor_param(h, kNVTEGroupedTensorOffsets, &off_tensor, sizeof(off_tensor)); } + // Shared gather of per-tensor scale_inv buffers into a contiguous device buffer. + // Returns (device buffer, total element count). Used by all block-scaling recipes + // (MXFP8 / NVFP4 / FP8 block) — they only differ in element size and CPU getter. + auto gather_scale_inv = [&](size_t bytes_per_elem, auto get_shape_fn, + auto get_cpu_ptr_fn) -> std::pair, size_t> { + size_t total_elems = 0; + std::vector elem_offsets(num_tensors); + std::vector numels(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + elem_offsets[i] = total_elems; + const NVTEShape sshape = get_shape_fn(tensors[i]); + size_t numel = 1; + for (size_t d = 0; d < sshape.ndim; ++d) numel *= sshape.data[d]; + numels[i] = numel; + total_elems += numel; + } + CudaPtr<> buffer = cuda_alloc(total_elems * bytes_per_elem); + for (size_t i = 0; i < num_tensors; ++i) { + tensors[i]->to_cpu(); + NVTE_CHECK_CUDA(cudaGetLastError()); + void* dst = static_cast(buffer.get()) + elem_offsets[i] * bytes_per_elem; + NVTE_CHECK_CUDA(cudaMemcpy(dst, get_cpu_ptr_fn(tensors[i]), + numels[i] * bytes_per_elem, cudaMemcpyHostToDevice)); + } + return {std::move(buffer), total_elems}; + }; + if (isFp8Type(dtype) && scaling_mode == NVTE_DELAYED_TENSOR_SCALING) { // FP8 tensor scaling: one float scale_inv per tensor // For delayed scaling, rowwise and columnwise share the same scale @@ -1243,67 +1281,113 @@ GroupedBuffers build_grouped_tensor(const std::vector& tensors, nvte_set_grouped_tensor_param(h, kNVTEGroupedColumnwiseScaleInv, &scale_tensor, sizeof(scale_tensor)); } else if (scaling_mode == NVTE_MXFP8_1D_SCALING) { - // MXFP8: E8M0 scale_inv per block of 32 elements - // Helper to gather scale_inv from individual tensors into a contiguous buffer - auto gather_scales = [&]( - auto get_shape_fn, - auto get_cpu_ptr_fn) -> std::pair, size_t> { - // Compute total size and offsets - size_t total_bytes = 0; - std::vector scale_offsets(num_tensors); - std::vector numels(num_tensors); - - for (size_t i = 0; i < num_tensors; ++i) { - scale_offsets[i] = total_bytes; - const NVTEShape shape = get_shape_fn(tensors[i]); - size_t numel = 1; - for (size_t d = 0; d < shape.ndim; ++d) { - numel *= shape.data[d]; - } - numels[i] = numel; - total_bytes += numel; // E8M0 is 1 byte per element - } - - // Allocate and copy - CudaPtr<> buffer = cuda_alloc(total_bytes); - for (size_t i = 0; i < num_tensors; ++i) { - tensors[i]->to_cpu(); - NVTE_CHECK_CUDA(cudaGetLastError()); - void* dst = static_cast(buffer.get()) + scale_offsets[i]; - const void* src = get_cpu_ptr_fn(tensors[i]); - NVTE_CHECK_CUDA(cudaMemcpy(dst, src, numels[i], cudaMemcpyHostToDevice)); - } - return {std::move(buffer), total_bytes}; - }; - - // Gather rowwise scale_inv if available + // MXFP8: E8M0 scale_inv per block of 32 elements (1 byte per scale element). if (has_rowwise) { - auto [row_buffer, row_total] = gather_scales( + auto [row_buffer, row_total] = gather_scale_inv( + /*bytes_per_elem=*/1, [](Tensor* t) { return t->rowwise_scale_inv_shape(); }, - [](Tensor* t) { return t->rowwise_cpu_scale_inv_ptr(); }); + [](Tensor* t) -> const void* { return t->rowwise_cpu_scale_inv_ptr(); }); grouped.scale_inv = std::move(row_buffer); - NVTEShape row_shape = nvte_make_shape(&row_total, 1); NVTEBasicTensor row_tensor{grouped.scale_inv.get(), kNVTEFloat8E8M0, row_shape}; nvte_set_grouped_tensor_param(h, kNVTEGroupedRowwiseScaleInv, &row_tensor, sizeof(row_tensor)); } - - // Gather columnwise scale_inv if available if (has_columnwise) { - auto [col_buffer, col_total] = gather_scales( + auto [col_buffer, col_total] = gather_scale_inv( + /*bytes_per_elem=*/1, [](Tensor* t) { return t->columnwise_scale_inv_shape(); }, - [](Tensor* t) { return t->columnwise_cpu_scale_inv_ptr(); }); + [](Tensor* t) -> const void* { return t->columnwise_cpu_scale_inv_ptr(); }); grouped.columnwise_scale_inv = std::move(col_buffer); - NVTEShape col_shape = nvte_make_shape(&col_total, 1); NVTEBasicTensor col_tensor{grouped.columnwise_scale_inv.get(), kNVTEFloat8E8M0, col_shape}; nvte_set_grouped_tensor_param(h, kNVTEGroupedColumnwiseScaleInv, &col_tensor, sizeof(col_tensor)); } - - // Mark as having swizzled scales (required for GEMM) const uint8_t swizzled = 1; nvte_set_grouped_tensor_param(h, kNVTEGroupedWithGEMMSwizzledScales, &swizzled, sizeof(swizzled)); + } else if (scaling_mode == NVTE_BLOCK_SCALING_1D || scaling_mode == NVTE_BLOCK_SCALING_2D) { + // FP8 block scaling: float32 scale_inv per block of 128 elements. + if (has_rowwise) { + auto [row_buffer, row_total] = gather_scale_inv( + /*bytes_per_elem=*/sizeof(float), + [](Tensor* t) { return t->rowwise_scale_inv_shape(); }, + [](Tensor* t) -> const void* { return t->rowwise_cpu_scale_inv_ptr(); }); + grouped.scale_inv = std::move(row_buffer); + NVTEShape row_shape = nvte_make_shape(&row_total, 1); + NVTEBasicTensor row_tensor{grouped.scale_inv.get(), kNVTEFloat32, row_shape}; + nvte_set_grouped_tensor_param(h, kNVTEGroupedRowwiseScaleInv, &row_tensor, sizeof(row_tensor)); + } + if (has_columnwise) { + auto [col_buffer, col_total] = gather_scale_inv( + /*bytes_per_elem=*/sizeof(float), + [](Tensor* t) { return t->columnwise_scale_inv_shape(); }, + [](Tensor* t) -> const void* { return t->columnwise_cpu_scale_inv_ptr(); }); + grouped.columnwise_scale_inv = std::move(col_buffer); + NVTEShape col_shape = nvte_make_shape(&col_total, 1); + NVTEBasicTensor col_tensor{grouped.columnwise_scale_inv.get(), kNVTEFloat32, col_shape}; + nvte_set_grouped_tensor_param(h, kNVTEGroupedColumnwiseScaleInv, &col_tensor, sizeof(col_tensor)); + } + } else if (scaling_mode == NVTE_NVFP4_1D_SCALING) { + // NVFP4: E4M3 scale_inv per block of 16 elements (swizzled for GEMM, 1 byte per scale). + if (has_rowwise) { + auto [row_buffer, row_total] = gather_scale_inv( + /*bytes_per_elem=*/1, + [](Tensor* t) { return t->rowwise_scale_inv_shape(); }, + [](Tensor* t) -> const void* { return t->rowwise_cpu_scale_inv_ptr(); }); + grouped.scale_inv = std::move(row_buffer); + NVTEShape row_shape = nvte_make_shape(&row_total, 1); + NVTEBasicTensor row_tensor{grouped.scale_inv.get(), kNVTEFloat8E4M3, row_shape}; + nvte_set_grouped_tensor_param(h, kNVTEGroupedRowwiseScaleInv, &row_tensor, sizeof(row_tensor)); + } + if (has_columnwise) { + auto [col_buffer, col_total] = gather_scale_inv( + /*bytes_per_elem=*/1, + [](Tensor* t) { return t->columnwise_scale_inv_shape(); }, + [](Tensor* t) -> const void* { return t->columnwise_cpu_scale_inv_ptr(); }); + grouped.columnwise_scale_inv = std::move(col_buffer); + NVTEShape col_shape = nvte_make_shape(&col_total, 1); + NVTEBasicTensor col_tensor{grouped.columnwise_scale_inv.get(), kNVTEFloat8E4M3, col_shape}; + nvte_set_grouped_tensor_param(h, kNVTEGroupedColumnwiseScaleInv, &col_tensor, sizeof(col_tensor)); + } + + // Mark as having swizzled scales (required for NVFP4 GEMM) + uint8_t swizzled = 1; + nvte_set_grouped_tensor_param(h, kNVTEGroupedWithGEMMSwizzledScales, &swizzled, sizeof(swizzled)); + + // Gather per-tensor amax values for NVFP4 global scale computation + auto gather_amax = [&](NVTETensorParam param) -> CudaPtr<> { + // Check if first tensor has this amax + NVTEBasicTensor first_amax = nvte_get_tensor_param(tensors[0]->data(), param); + if (first_amax.data_ptr == nullptr) return CudaPtr<>(); + + std::vector amax_cpu(num_tensors); + for (size_t i = 0; i < num_tensors; ++i) { + NVTEBasicTensor amax_bt = nvte_get_tensor_param(tensors[i]->data(), param); + NVTE_CHECK(amax_bt.data_ptr != nullptr, "Tensor ", i, " is missing amax"); + float val; + NVTE_CHECK_CUDA(cudaMemcpy(&val, amax_bt.data_ptr, sizeof(float), cudaMemcpyDeviceToHost)); + amax_cpu[i] = val; + } + CudaPtr<> dev = cuda_alloc(sizeof(float) * num_tensors); + NVTE_CHECK_CUDA(cudaMemcpy(dev.get(), amax_cpu.data(), + sizeof(float) * num_tensors, cudaMemcpyHostToDevice)); + return dev; + }; + + grouped.amax_dev = gather_amax(kNVTEAmax); + if (grouped.amax_dev.get()) { + NVTEShape amax_shape = nvte_make_shape(&num_tensors, 1); + NVTEBasicTensor amax_tensor{grouped.amax_dev.get(), kNVTEFloat32, amax_shape}; + nvte_set_grouped_tensor_param(h, kNVTEGroupedAmax, &amax_tensor, sizeof(amax_tensor)); + } + + grouped.columnwise_amax_dev = gather_amax(kNVTEColumnwiseAmax); + if (grouped.columnwise_amax_dev.get()) { + NVTEShape amax_shape = nvte_make_shape(&num_tensors, 1); + NVTEBasicTensor amax_tensor{grouped.columnwise_amax_dev.get(), kNVTEFloat32, amax_shape}; + nvte_set_grouped_tensor_param(h, kNVTEGroupedColumnwiseAmax, &amax_tensor, sizeof(amax_tensor)); + } + } return grouped; diff --git a/tests/cpp/test_common.h b/tests/cpp/test_common.h index fd03d283d..11d96c2e6 100644 --- a/tests/cpp/test_common.h +++ b/tests/cpp/test_common.h @@ -177,6 +177,8 @@ class Tensor { NVTEShape columnwise_shape() const noexcept { return tensor_.get_columnwise_data().shape; } + NVTEShape shape() const noexcept { return tensor_.shape(); } + NVTEShape rowwise_scale_inv_shape() const { NVTE_CHECK(rowwise_, "Tensor does not have rowwise data!"); return tensor_.get_rowwise_scale_inv().shape; @@ -596,6 +598,8 @@ struct GroupedBuffers { CudaPtr last_dims_dev; CudaPtr offsets_dev; CudaPtr<> columnwise_data; + CudaPtr<> amax_dev; // Per-tensor amax for NVFP4 grouped GEMM + CudaPtr<> columnwise_amax_dev; // Per-tensor columnwise amax for NVFP4 grouped GEMM NVTEShape logical_shape{}; std::vector offsets_host; std::vector tensor_bytes; diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 5f82bfcba..368c95e27 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -2935,10 +2935,13 @@ def _apply_grouped_bias_ref( @pytest.mark.parametrize("accumulate", [False, True]) @pytest.mark.parametrize("use_bias_scale", [False, True]) def test_grouped_gemm_grouped_tensor(z, m, n, k, case, layout, accumulate, use_bias_scale) -> None: + if torch.cuda.get_device_capability() < (9, 0): + pytest.skip("Grouped GEMM requires Hopper (SM90) or newer.") + if torch.cuda.get_device_capability() < (10, 0): + if tex.get_cublasLt_version() < 130400: + pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") if tex.get_cublasLt_version() < 130300: pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") - if torch.cuda.get_device_capability() < (10, 0): - pytest.skip("Grouped GEMM requires Blackwell (SM100) or newer.") if not is_bf16_available(): pytest.skip("bfloat16 is required for grouped GEMM test.") diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 5b6a9bf41..6aa8798b7 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -91,6 +91,10 @@ inline bool is_mxfp_scaling(const NVTEScalingMode &mode) { return mode == NVTE_M inline bool is_nvfp_scaling(const NVTEScalingMode &mode) { return mode == NVTE_NVFP4_1D_SCALING; } +inline bool is_fp8_block_scaling(const NVTEScalingMode &mode) { + return mode == NVTE_BLOCK_SCALING_1D || mode == NVTE_BLOCK_SCALING_2D; +} + inline size_t product(const std::vector &shape, const size_t begin, const size_t end) { NVTE_CHECK(begin <= end && end <= shape.size(), "Attempted to access entries ", begin, " to ", end, " in a vector with ", shape.size(), " entries"); diff --git a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu index 6a7af158e..f064af247 100644 --- a/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu @@ -33,6 +33,16 @@ inline void CreateCublasHandle(cublasLtHandle_t *handle) { // MXFP8 support for grouped GEMM requires cuBLAS 13.3+ #define CUBLAS_MXFP8_GROUPED_GEMM_VERSION 130300 + +// Hopper (SM90) support for grouped GEMM requires cuBLAS 13.4+ +#define CUBLAS_GROUPED_GEMM_HOPPER_VERSION 130400 + +// NVFP4 support for grouped GEMM requires cuBLAS 13.4+ +#define CUBLAS_NVFP4_GROUPED_GEMM_VERSION 130400 + +// FP8 block scaling support for grouped GEMM requires cuBLAS 13.4+ +#define CUBLAS_FP8_BLOCK_GROUPED_GEMM_VERSION 130400 + // BF16 support for grouped GEMM requires cuBLAS 13.3+ #define CUBLAS_GROUPED_GEMM_VERSION 130300 @@ -132,124 +142,126 @@ inline int64_t compute_avg_last_dim(const transformer_engine::GroupedTensor *t) static constexpr size_t kGroupedGemmAlignment = 256; static constexpr size_t kGroupedGemmCublasWorkspaceSize = 32ull * 1024 * 1024; // 32 MiB -// Workspace layout for grouped GEMM +// Workspace layout for grouped GEMM. +// Layout described once in `from_buffers`; `required_setup_size` runs the same walker +// with base=nullptr to derive the total byte count, so the two stay in sync by construction. struct GroupedGemmSetupWorkspace { - void **A_ptrs; - void **B_ptrs; - void **C_ptrs; - void **D_ptrs; - float **alpha_ptrs; - float **beta_ptrs; - void ** - a_scale_inv_ptrs; // Per-tensor FP8 scale pointers for A (float* for tensor scaling, E8M0* for MXFP8) - void ** - b_scale_inv_ptrs; // Per-tensor FP8 scale pointers for B (float* for tensor scaling, E8M0* for MXFP8) + void **A_ptrs = nullptr; + void **B_ptrs = nullptr; + void **C_ptrs = nullptr; + void **D_ptrs = nullptr; + float **alpha_ptrs = nullptr; + float **beta_ptrs = nullptr; + // Per-tensor scale_inv pointers (float* for tensor/FP8 block scaling, E8M0* for MXFP8, + // E4M3* for NVFP4) + void **a_scale_inv_ptrs = nullptr; + void **b_scale_inv_ptrs = nullptr; // Storage dimensions for cuBLAS matrix layouts - int *a_rows; - int *a_cols; - int *b_rows; - int *b_cols; - int *d_rows; // M (first dim) - also used for C - int *d_cols; // N (last dim) - also used for C - - // Initialize from workspace buffer - // Layout: all pointer arrays first (16-byte aligned for cuBLAS), then int arrays - static GroupedGemmSetupWorkspace from_buffers(char *setup_ws_ptr, size_t num_tensors) { + int *a_rows = nullptr; + int *a_cols = nullptr; + int *b_rows = nullptr; + int *b_cols = nullptr; + int *d_rows = nullptr; // M (first dim) - also used for C + int *d_cols = nullptr; // N (last dim) - also used for C + // NVFP4: per-group computed alpha values (alpha * amax_A * amax_B * factor_inv) + float *nvfp4_computed_alpha = nullptr; + // End-of-layout offset in bytes (unaligned). required_setup_size rounds this up. + size_t total_bytes = 0; + + // Walk the layout once. If `base` is non-null, fields are populated; otherwise + // only `total_bytes` is meaningful (used by required_setup_size). + static GroupedGemmSetupWorkspace from_buffers(char *base, size_t num_tensors) { GroupedGemmSetupWorkspace ws; - size_t offset = 0; + constexpr size_t kPtrAlignment = 16; // cuBLAS requires 16-byte alignment for pointer arrays const size_t ptr_size = num_tensors * sizeof(void *); const size_t int_size = num_tensors * sizeof(int); - constexpr size_t kPtrAlignment = 16; // cuBLAS requires 16-byte alignment for pointer arrays + const size_t float_size = num_tensors * sizeof(float); + size_t offset = 0; - // Helper to align offset to kPtrAlignment - auto align_offset = [&]() { + auto align_ptr = [&]() { offset = (offset + kPtrAlignment - 1) / kPtrAlignment * kPtrAlignment; }; + auto place = [&](auto *&field, size_t size_bytes) { + using Field = std::remove_reference_t; + if (base != nullptr) field = reinterpret_cast(base + offset); + offset += size_bytes; + }; - // Pointer arrays first (all 16-byte aligned for cuBLAS grouped GEMM) - align_offset(); - ws.A_ptrs = reinterpret_cast(setup_ws_ptr + offset); - offset += ptr_size; - align_offset(); - ws.B_ptrs = reinterpret_cast(setup_ws_ptr + offset); - offset += ptr_size; - align_offset(); - ws.C_ptrs = reinterpret_cast(setup_ws_ptr + offset); - offset += ptr_size; - align_offset(); - ws.D_ptrs = reinterpret_cast(setup_ws_ptr + offset); - offset += ptr_size; - align_offset(); - ws.alpha_ptrs = reinterpret_cast(setup_ws_ptr + offset); - offset += ptr_size; - align_offset(); - ws.beta_ptrs = reinterpret_cast(setup_ws_ptr + offset); - offset += ptr_size; - align_offset(); - ws.a_scale_inv_ptrs = reinterpret_cast(setup_ws_ptr + offset); - offset += ptr_size; - align_offset(); - ws.b_scale_inv_ptrs = reinterpret_cast(setup_ws_ptr + offset); - offset += ptr_size; - - // Int arrays for storage dimensions (4-byte aligned is fine) - align_offset(); - ws.a_rows = reinterpret_cast(setup_ws_ptr + offset); - offset += int_size; - ws.a_cols = reinterpret_cast(setup_ws_ptr + offset); - offset += int_size; - ws.b_rows = reinterpret_cast(setup_ws_ptr + offset); - offset += int_size; - ws.b_cols = reinterpret_cast(setup_ws_ptr + offset); - offset += int_size; - ws.d_rows = reinterpret_cast(setup_ws_ptr + offset); - offset += int_size; - ws.d_cols = reinterpret_cast(setup_ws_ptr + offset); - + // 8 pointer arrays (each 16-byte aligned), then 6 int arrays, then 1 float array. + align_ptr(); + place(ws.A_ptrs, ptr_size); + align_ptr(); + place(ws.B_ptrs, ptr_size); + align_ptr(); + place(ws.C_ptrs, ptr_size); + align_ptr(); + place(ws.D_ptrs, ptr_size); + align_ptr(); + place(ws.alpha_ptrs, ptr_size); + align_ptr(); + place(ws.beta_ptrs, ptr_size); + align_ptr(); + place(ws.a_scale_inv_ptrs, ptr_size); + align_ptr(); + place(ws.b_scale_inv_ptrs, ptr_size); + // Int/float arrays follow without extra align_ptr(): cuBLAS only requires 16-byte + // alignment for the pointer arrays above; int and float need just their natural + // 4-byte alignment. The offset is 16-byte aligned after the last align_ptr() and + // each subsequent place() adds N*4 bytes, so it stays a multiple of 4. + place(ws.a_rows, int_size); + place(ws.a_cols, int_size); + place(ws.b_rows, int_size); + place(ws.b_cols, int_size); + place(ws.d_rows, int_size); + place(ws.d_cols, int_size); + place(ws.nvfp4_computed_alpha, float_size); + + ws.total_bytes = offset; return ws; } - // Calculate required size for setup workspace static size_t required_setup_size(size_t num_tensors, size_t alignment) { - const size_t ptr_size = num_tensors * sizeof(void *); - const size_t int_size = num_tensors * sizeof(int); - constexpr size_t kPtrAlignment = 16; // Must match from_buffers - - // Layout: 8 ptr arrays (each 16-byte aligned), then 6 int arrays - // Each ptr array takes ptr_size bytes but needs to start at 16-byte boundary - auto aligned_ptr_size = ((ptr_size + kPtrAlignment - 1) / kPtrAlignment) * kPtrAlignment; - size_t size = 8 * aligned_ptr_size + 6 * int_size; - size = ((size + alignment - 1) / alignment) * alignment; - return size; + const size_t raw = from_buffers(nullptr, num_tensors).total_bytes; + // Additional alignment bytes is to take care of the case where the buffer + // is not already aligned. + return raw + alignment; } }; +inline bool grouped_gemm_supports_per_group_alpha_beta(int sm) { return sm >= 100; } + inline size_t validate_grouped_gemm_inputs( size_t num_tensors, std::initializer_list inputs, - const transformer_engine::Tensor *alpha_tensor, const transformer_engine::Tensor *beta_tensor) { + const transformer_engine::Tensor *alpha_tensor, const transformer_engine::Tensor *beta_tensor, + bool supports_per_group_alpha_beta) { NVTE_CHECK(num_tensors >= 1, "Grouped GEMM: number of tensors must be at least 1"); for (const auto *tensor : inputs) { NVTE_CHECK(tensor->num_tensors == num_tensors, "Grouped GEMM: inputs must have the same number of tensors"); } + // Hopper currently requires a uniform alpha/beta scalar for the whole grouped GEMM, + // while Blackwell+ supports per-matrix alpha/beta. const size_t alpha_numel = alpha_tensor->data.numel(); const size_t beta_numel = beta_tensor->data.numel(); - NVTE_CHECK(alpha_numel == num_tensors, "Grouped GEMM: alpha must have num_tensors (", num_tensors, - ") elements, got ", alpha_numel); - NVTE_CHECK(beta_numel == num_tensors, "Grouped GEMM: beta must have num_tensors (", num_tensors, - ") elements, got ", beta_numel); + const size_t expected_alphabeta_numel = supports_per_group_alpha_beta ? num_tensors : 1; + const char *alphabeta_desc = supports_per_group_alpha_beta ? "num_tensors" : "1"; + NVTE_CHECK(alpha_numel == expected_alphabeta_numel, "Grouped GEMM: alpha must have ", + alphabeta_desc, " element(s), got ", alpha_numel); + NVTE_CHECK(beta_numel == expected_alphabeta_numel, "Grouped GEMM: beta must have ", + alphabeta_desc, " element(s), got ", beta_numel); auto is_supported_input_dtype = [](transformer_engine::DType dtype) { return dtype == transformer_engine::DType::kFloat8E4M3 || dtype == transformer_engine::DType::kFloat8E5M2 || dtype == transformer_engine::DType::kBFloat16 || - dtype == transformer_engine::DType::kFloat16; + dtype == transformer_engine::DType::kFloat16 || + dtype == transformer_engine::DType::kFloat4E2M1; }; for (const auto *tensor : inputs) { if (tensor->has_data() || tensor->has_columnwise_data()) { NVTE_CHECK(is_supported_input_dtype(tensor->dtype()), - "Grouped GEMM inputs must be FP8, BF16, or FP16, got ", + "Grouped GEMM inputs must be FP8, NVFP4, BF16, or FP16, got ", transformer_engine::to_string(tensor->dtype()), "."); } } @@ -263,37 +275,54 @@ inline size_t validate_grouped_gemm_inputs( } if (ref != nullptr) { const bool ref_is_fp8 = is_fp8_dtype(ref->dtype()); + const bool ref_is_fp4 = is_fp4_dtype(ref->dtype()); const bool ref_is_mxfp8 = transformer_engine::is_mxfp_scaling(ref->scaling_mode); + const bool ref_is_nvfp4 = transformer_engine::is_nvfp_scaling(ref->scaling_mode); + const bool ref_is_fp8_block = transformer_engine::is_fp8_block_scaling(ref->scaling_mode); for (const auto *tensor : inputs) { if (!(tensor->has_data() || tensor->has_columnwise_data())) continue; NVTE_CHECK(is_fp8_dtype(tensor->dtype()) == ref_is_fp8, "Grouped GEMM: A and B must both be FP8 or both be non-FP8."); + NVTE_CHECK(is_fp4_dtype(tensor->dtype()) == ref_is_fp4, + "Grouped GEMM: A and B must both be NVFP4 or both be non-NVFP4."); NVTE_CHECK(transformer_engine::is_mxfp_scaling(tensor->scaling_mode) == ref_is_mxfp8, - "Grouped GEMM: A and B must both use MXFP8 scaling or both use tensor scaling."); - if (ref_is_mxfp8) { + "Grouped GEMM: A and B must both use MXFP8 scaling or both not."); + NVTE_CHECK(transformer_engine::is_nvfp_scaling(tensor->scaling_mode) == ref_is_nvfp4, + "Grouped GEMM: A and B must both use NVFP4 scaling or both not."); + NVTE_CHECK(transformer_engine::is_fp8_block_scaling(tensor->scaling_mode) == ref_is_fp8_block, + "Grouped GEMM: A and B must both use FP8 block scaling or both not."); + if (ref_is_mxfp8 || transformer_engine::is_nvfp_scaling(tensor->scaling_mode)) { NVTE_CHECK(tensor->with_gemm_swizzled_scales, - "MXFP8 grouped GEMM: scales must be swizzled for GEMM."); + "Grouped GEMM: scales must be swizzled for GEMM (MXFP8/NVFP4)."); } } } return num_tensors; } +inline void validate_grouped_gemm_output_dtype(transformer_engine::DType a_dtype, + transformer_engine::DType b_dtype, + transformer_engine::DType output_dtype, + const char *name) { + const bool is_output_dtype = output_dtype == transformer_engine::DType::kBFloat16 || + output_dtype == transformer_engine::DType::kFloat16 || + output_dtype == transformer_engine::DType::kFloat32; + NVTE_CHECK(is_output_dtype, "Grouped GEMM: ", name, " must be BF16, FP16, or FP32."); + if (!is_fp4_dtype(a_dtype) && !is_fp4_dtype(b_dtype)) return; + NVTE_CHECK(!is_fp4_dtype(output_dtype), "FP4 GEMM output is not supported!"); + NVTE_CHECK(get_cuda_dtype(output_dtype) != CUDA_R_16F, "FP4 GEMM does not support FP16 output!"); +} + inline void validate_grouped_gemm_outputs( - size_t num_tensors, std::initializer_list outputs) { - auto is_output_dtype = [](transformer_engine::DType dtype) { - return dtype == transformer_engine::DType::kBFloat16 || - dtype == transformer_engine::DType::kFloat16 || - dtype == transformer_engine::DType::kFloat32; - }; + size_t num_tensors, transformer_engine::DType a_dtype, transformer_engine::DType b_dtype, + std::initializer_list outputs) { for (const auto *tensor : outputs) { if (tensor == nullptr) { continue; } NVTE_CHECK(tensor->num_tensors == num_tensors, "Grouped GEMM: outputs must have the same number of tensors as inputs"); - NVTE_CHECK(is_output_dtype(tensor->dtype()), - "Grouped GEMM: outputs must be BF16, FP16, or FP32."); + validate_grouped_gemm_output_dtype(a_dtype, b_dtype, tensor->dtype(), "outputs"); } } @@ -303,11 +332,22 @@ inline size_t grouped_gemm_setup_workspace_size(size_t num_tensors) { inline void check_grouped_gemm_requirements(const char *api_name) { const int current_device = transformer_engine::cuda::current_device(); - NVTE_CHECK(transformer_engine::cuda::sm_arch(current_device) >= 100, api_name, - " requires Blackwell (SM100) or newer architecture."); - NVTE_CHECK(transformer_engine::cuda::cublas_version() >= CUBLAS_GROUPED_GEMM_VERSION, api_name, - " requires cuBLAS 13.3+, but run-time cuBLAS version is ", - transformer_engine::cuda::cublas_version()); + const int sm = transformer_engine::cuda::sm_arch(current_device); + const int cublas_ver = transformer_engine::cuda::cublas_version(); +#if CUBLAS_VERSION >= CUBLAS_GROUPED_GEMM_HOPPER_VERSION + NVTE_CHECK(sm >= 90, api_name, " requires Hopper (SM90) or newer architecture."); + NVTE_CHECK(cublas_ver >= CUBLAS_GROUPED_GEMM_VERSION, api_name, + " requires cuBLAS 13.3+, but run-time cuBLAS version is ", cublas_ver); + if (sm < 100) { + NVTE_CHECK(cublas_ver >= CUBLAS_GROUPED_GEMM_HOPPER_VERSION, api_name, + " on Hopper (SM90) requires cuBLAS 13.4+, but run-time cuBLAS version is ", + cublas_ver); + } +#else + NVTE_CHECK(sm >= 100, api_name, " requires Blackwell (SM100) or newer architecture."); + NVTE_CHECK(cublas_ver >= CUBLAS_GROUPED_GEMM_VERSION, api_name, + " requires cuBLAS 13.3+, but run-time cuBLAS version is ", cublas_ver); +#endif } inline transformer_engine::GroupedMatmulConfig parse_grouped_gemm_config( @@ -319,19 +359,90 @@ inline transformer_engine::GroupedMatmulConfig parse_grouped_gemm_config( return config_; } -// Select row-wise vs column-wise storage and adjust transpose flag for grouped GEMM. -// Mirrors the non-grouped GEMM logic for FP8 layout handling (TN-only on Hopper) and -// fallback to column-wise data when row-wise is absent. -// Contains all information needed for GEMM setup - shape already accounts for storage layout. +// Contains all information needed for one tensor operand for GEMM setup. struct GroupedOperandSelection { - TensorShapeInfo shape; // Shape info with dims already swapped for columnwise if needed + TensorShapeInfo logical_tensor_shape; char *dptr = nullptr; void *scale_inv = nullptr; // Contiguous array of scales (input) + void *amax = nullptr; // Per-tensor amax values (NVFP4 only) transformer_engine::DType dtype = transformer_engine::DType::kNumTypes; NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING; bool with_gemm_swizzled_scales = false; bool trans = false; bool rowwise = true; + // Whether selected storage is physically transposed relative to logical shape. + bool storage_transposed = false; +}; + +inline void validate_nvfp4_grouped_gemm_support(const GroupedOperandSelection &A_sel, + const GroupedOperandSelection &B_sel, + bool use_per_group_alpha_beta) { + const bool nvfp4 = transformer_engine::is_nvfp_scaling(A_sel.scaling_mode) || + transformer_engine::is_nvfp_scaling(B_sel.scaling_mode); + if (!nvfp4) return; + + NVTE_CHECK(transformer_engine::is_nvfp_scaling(A_sel.scaling_mode) && + transformer_engine::is_nvfp_scaling(B_sel.scaling_mode), + "Grouped GEMM: A and B must both use NVFP4 scaling or both not."); + NVTE_CHECK(use_per_group_alpha_beta, + "Grouped GEMM: NVFP4 requires per-group alpha/beta support because each group " + "has its own amax-derived global scale."); +} + +// FP8 block scaling grouped GEMM is only supported on Hopper (SM90). +inline void validate_fp8_block_grouped_gemm_support(const GroupedOperandSelection &A_sel, + const GroupedOperandSelection &B_sel, int sm) { + const bool a_fp8_block = transformer_engine::is_fp8_block_scaling(A_sel.scaling_mode); + const bool b_fp8_block = transformer_engine::is_fp8_block_scaling(B_sel.scaling_mode); + if (!a_fp8_block && !b_fp8_block) return; + + NVTE_CHECK(a_fp8_block && b_fp8_block, + "Grouped GEMM: A and B must both use FP8 block scaling or both not."); + NVTE_CHECK(sm == 90, + "Grouped GEMM: FP8 block scaling is only supported on Hopper (SM90); " + "use MXFP8 on Blackwell (SM100) or newer."); +} + +inline bool is_compatible_grouped_scaling_mode(NVTEScalingMode a_mode, NVTEScalingMode b_mode) { + const bool a_fp8_block = transformer_engine::is_fp8_block_scaling(a_mode); + const bool b_fp8_block = transformer_engine::is_fp8_block_scaling(b_mode); + if (a_fp8_block || b_fp8_block) { + return a_fp8_block && b_fp8_block && + !(a_mode == NVTE_BLOCK_SCALING_2D && b_mode == NVTE_BLOCK_SCALING_2D); + } + return a_mode == b_mode; +} + +// Validates A/B scaling-mode pairing and Hopper-only FP8 block scaling support. +// Call from every grouped GEMM entry point after operand scaling modes are known. +inline void validate_grouped_gemm_scaling_modes(NVTEScalingMode a_mode, NVTEScalingMode b_mode, + int sm, const char *api_name) { + if (a_mode == NVTE_BLOCK_SCALING_2D && b_mode == NVTE_BLOCK_SCALING_2D) { + NVTE_CHECK(false, api_name, + ": Only 1D by 1D, 1D by 2D, and 2D by 1D FP8 block scaling grouped GEMM is " + "supported, but got 2D by 2D."); + } + NVTE_CHECK(is_compatible_grouped_scaling_mode(a_mode, b_mode), api_name, + ": incompatible A/B scaling modes."); + if (transformer_engine::is_fp8_block_scaling(a_mode) || + transformer_engine::is_fp8_block_scaling(b_mode)) { + NVTE_CHECK(sm >= 90 && sm < 100, api_name, + ": FP8 block scaling grouped GEMM is only supported on Hopper (SM90-SM99), " + "not SM", + sm, "."); + } +} + +struct GroupedGemmConfig { + bool use_split_accumulator = false; + bool use_fp8 = false; + bool use_per_group_alpha_beta = false; + void *alpha_dptr = nullptr; + void *beta_dptr = nullptr; + int64_t avg_m = 0; + int64_t avg_n = 0; + int64_t avg_k = 0; + int sm_count = 0; }; constexpr int kMaxGroups = 64; @@ -346,6 +457,7 @@ struct MultiTensorGroupGemmOutputArgs { struct MultiTensorGroupGemmInputArgs { void *data_ptrs[kMaxGroups]; void *scale_inv_ptrs[kMaxGroups]; + void *amax_ptrs[kMaxGroups]; int rows[kMaxGroups]; int cols[kMaxGroups]; }; @@ -360,12 +472,15 @@ struct MultiTensorListInfo { struct OperandStorageChoice { bool use_rowwise = true; - bool swap_dims = true; + // Only meaningful when use_rowwise == false (columnwise storage). Indicates that the + // columnwise buffer is physically transposed relative to logical shape. + bool storage_transposed = false; bool trans = false; }; inline OperandStorageChoice choose_grouped_operand_storage(bool trans, bool is_A, bool is_mxfp8, - bool is_fp8, bool non_tn_fp8_ok, + bool is_fp8, bool is_nvfp4, + bool is_fp8_block, bool non_tn_fp8_ok, bool has_row, bool has_col, const char *name) { NVTE_CHECK(has_row || has_col, "Grouped GEMM: ", name, @@ -374,7 +489,7 @@ inline OperandStorageChoice choose_grouped_operand_storage(bool trans, bool is_A if (is_A) { if (trans) { NVTE_CHECK(has_row, "Grouped GEMM: MXFP8 transposed ", name, " is missing row-wise data"); - return {true, true, trans}; + return {true, false, trans}; } NVTE_CHECK(has_col, "Grouped GEMM: MXFP8 non-transposed ", name, " is missing column-wise data"); @@ -385,19 +500,48 @@ inline OperandStorageChoice choose_grouped_operand_storage(bool trans, bool is_A return {false, false, trans}; } NVTE_CHECK(has_row, "Grouped GEMM: MXFP8 non-transposed ", name, " is missing row-wise data"); - return {true, true, trans}; + return {true, false, trans}; } - // Hopper-style TN-only FP8: force TN by switching layout and flipping transpose when needed. + // FP8 block scaling on Hopper: force TN by using transposed columnwise data. + if (is_fp8_block && !non_tn_fp8_ok) { + if (is_A && !trans) { + NVTE_CHECK(has_col, "Grouped GEMM: ", name, + " is missing column-wise data needed for TN layout"); + return {false, true, true}; + } + if (!is_A && trans) { + NVTE_CHECK(has_col, "Grouped GEMM: ", name, + " is missing column-wise data needed for TN layout"); + return {false, true, false}; + } + } + + // NVFP4: force TN by switching layout and flipping transpose. + // NVFP4 columnwise data is the transposed tensor quantized rowwise. + if (is_nvfp4) { + if (is_A && !trans) { + NVTE_CHECK(has_col, "Grouped GEMM: ", name, + " is missing column-wise data needed for TN layout"); + return {false, true, true}; + } + if (!is_A && trans) { + NVTE_CHECK(has_col, "Grouped GEMM: ", name, + " is missing column-wise data needed for TN layout"); + return {false, true, false}; + } + } + + // Hopper-style TN-only FP8 (tensor scaling): force TN by switching layout and flipping transpose. if (is_fp8 && !non_tn_fp8_ok) { if (is_A && !trans) { NVTE_CHECK(has_col, "Grouped GEMM: ", name, - " is missing column-wise data needed for FP8 TN layout"); + " is missing column-wise data needed for TN layout"); return {false, true, true}; } if (!is_A && trans) { NVTE_CHECK(has_col, "Grouped GEMM: ", name, - " is missing column-wise data needed for FP8 TN layout"); + " is missing column-wise data needed for TN layout"); return {false, true, false}; } } @@ -410,7 +554,7 @@ inline OperandStorageChoice choose_grouped_operand_storage(bool trans, bool is_A } NVTE_CHECK(has_row, "Grouped GEMM: ", name, " is missing row-wise data"); - return {true, true, trans}; + return {true, false, trans}; } // Build Kernel Arguments detailing out addresses and other metadata for list of C/D tensors @@ -450,7 +594,7 @@ inline MultiTensorGroupGemmOutputArgs build_grouped_gemm_multi_out_args( // Build Kernel Arguments detailing out addresses and other metadata for list of A tensors // passed to the grouped GEMM kernel. Use-case: A --> List of Expert weights inline MultiTensorGroupGemmInputArgs build_grouped_gemm_multi_inputA_args( - const NVTETensor *tensor_list, size_t list_size, bool use_rowwise, bool is_fp8, + const NVTETensor *tensor_list, size_t list_size, bool use_rowwise, bool storage_transposed, int64_t *avg_first_dim, int64_t *avg_last_dim, const char *name) { using namespace transformer_engine; MultiTensorGroupGemmInputArgs args{}; @@ -467,20 +611,33 @@ inline MultiTensorGroupGemmInputArgs build_grouped_gemm_multi_inputA_args( use_rowwise ? t->scale_inv : t->columnwise_scale_inv; NVTE_CHECK(data.has_data(), "Grouped GEMM: ", name, "_list tensor ", i, " is missing required data."); - NVTE_CHECK(data.shape.size() == 2, "Grouped GEMM: ", name, "_list tensor ", i, " must be 2D."); args.data_ptrs[i] = data.dptr; - args.rows[i] = static_cast(data.shape[1]); - args.cols[i] = static_cast(data.shape[0]); - *avg_first_dim += static_cast(data.shape[0]); - *avg_last_dim += static_cast(data.shape[1]); + const auto &shape = t->shape(); + NVTE_CHECK(shape.size() == 2, "Grouped GEMM: ", name, "_list tensor ", i, " must be 2D."); + const size_t first_dim = shape[0]; + const size_t last_dim = shape[1]; + if (storage_transposed) { + args.rows[i] = static_cast(first_dim); + args.cols[i] = static_cast(last_dim); + } else { + args.rows[i] = static_cast(last_dim); + args.cols[i] = static_cast(first_dim); + } + *avg_first_dim += static_cast(first_dim); + *avg_last_dim += static_cast(last_dim); - if (is_fp8) { + const bool scale_inv_needed = is_fp8_dtype(t->dtype()) || is_nvfp_scaling(t->scaling_mode) || + is_fp8_block_scaling(t->scaling_mode); + if (scale_inv_needed) { NVTE_CHECK(scale_inv.has_data(), "Grouped GEMM: ", name, "_list tensor ", i, - " requires scale_inv for FP8."); + " requires scale_inv."); args.scale_inv_ptrs[i] = scale_inv.dptr; } else { args.scale_inv_ptrs[i] = nullptr; } + + const transformer_engine::SimpleTensor &amax_src = use_rowwise ? t->amax : t->columnwise_amax; + args.amax_ptrs[i] = amax_src.has_data() ? amax_src.dptr : nullptr; } *avg_first_dim /= static_cast(list_size); *avg_last_dim /= static_cast(list_size); @@ -509,8 +666,11 @@ inline MultiTensorListInfo validate_grouped_gemm_multi_inputA_list(const NVTETen info.scaling_mode = t0->scaling_mode; info.with_gemm_swizzled_scales = t0->with_gemm_swizzled_scales; const bool mxfp8 = transformer_engine::is_mxfp_scaling(info.scaling_mode); - NVTE_CHECK(info.scaling_mode == NVTE_DELAYED_TENSOR_SCALING || mxfp8, - "Grouped GEMM: input list only supports tensor scaling or MXFP8."); + const bool nvfp4 = transformer_engine::is_nvfp_scaling(info.scaling_mode); + const bool fp8_block = transformer_engine::is_fp8_block_scaling(info.scaling_mode); + NVTE_CHECK(info.scaling_mode == NVTE_DELAYED_TENSOR_SCALING || mxfp8 || nvfp4 || fp8_block, + "Grouped GEMM: input list only supports tensor scaling, MXFP8, NVFP4, " + "or FP8 block scaling."); for (size_t i = 0; i < list_size; ++i) { const transformer_engine::Tensor *t = @@ -547,11 +707,9 @@ inline MultiTensorListInfo validate_grouped_gemm_multi_inputA_list(const NVTETen return info; } -// Helper to create TensorShapeInfo from a GroupedTensor, optionally swapping first/last dims. -// When swap_dims=true, first_dims and last_dims are swapped to account for columnwise storage. -// Note: tensor_offsets are the same for rowwise and columnwise data (same element count per tensor). -inline TensorShapeInfo create_shape_info(const transformer_engine::GroupedTensor *t, - bool swap_dims) { +// Helper to create TensorShapeInfo from a GroupedTensor. Grouped tensor metadata is logical +// shape; storage-specific transposes are handled when building cuBLAS matrix layouts. +inline TensorShapeInfo create_shape_info(const transformer_engine::GroupedTensor *t) { const bool has_first = t->first_dims.has_data(); const bool has_last = t->last_dims.has_data(); NVTE_CHECK(has_first || t->all_same_first_dim(), @@ -567,10 +725,6 @@ inline TensorShapeInfo create_shape_info(const transformer_engine::GroupedTensor const int64_t *offsets_ptr = t->tensor_offsets.has_data() ? static_cast(t->tensor_offsets.dptr) : nullptr; - if (swap_dims) { - // Swap first/last to account for columnwise (transposed) storage - return {last_ptr, first_ptr, offsets_ptr, uniform_last, uniform_first}; - } return {first_ptr, last_ptr, offsets_ptr, uniform_first, uniform_last}; } @@ -585,16 +739,19 @@ inline GroupedOperandSelection select_grouped_operand(const transformer_engine:: sel.trans = trans; sel.scaling_mode = t->scaling_mode; sel.dtype = t->dtype(); - sel.shape = create_shape_info(t, /*swap_dims=*/false); + sel.logical_tensor_shape = create_shape_info(t); return sel; } const auto sm = t->scaling_mode; const bool mxfp8 = is_mxfp_scaling(sm); + const bool nvfp4 = is_nvfp_scaling(sm); + const bool fp8_block = is_fp8_block_scaling(sm); // Validate scaling mode - NVTE_CHECK(sm == NVTE_DELAYED_TENSOR_SCALING || mxfp8, - "Grouped GEMM is only supported with bf16, fp8 tensor scaling and MXFP8"); + NVTE_CHECK(sm == NVTE_DELAYED_TENSOR_SCALING || mxfp8 || nvfp4 || fp8_block, + "Grouped GEMM is only supported with bf16, fp8 tensor scaling, MXFP8, NVFP4, " + "and FP8 block scaling"); const DType row_dtype = t->data.dtype; const DType col_dtype = t->columnwise_data.dtype; @@ -605,36 +762,40 @@ inline GroupedOperandSelection select_grouped_operand(const transformer_engine:: const DType rep_dtype = has_row ? row_dtype : col_dtype; const bool is_fp8 = is_fp8_dtype(rep_dtype); - const bool non_tn_fp8_ok = nvte_is_non_tn_fp8_gemm_supported(); + // FP8 block scaling on Hopper requires TN layout (same as tensor scaling) + const bool non_tn_fp8_ok = fp8_block ? false : nvte_is_non_tn_fp8_gemm_supported(); // Helper to select columnwise storage. - // swap_dims=true (default): swap first/last dims in shape info (used when columnwise == transposed). - // swap_dims=false: keep original dims (MXFP8: columnwise data has different scale direction, - // but the logical matrix shape and transpose flag remain unchanged). - auto use_columnwise = [&](bool swap_dims = true) { + // storage_transposed=true: columnwise data is physically transposed relative to logical shape. + // storage_transposed=false: columnwise data has logical shape (MXFP8). + auto use_columnwise = [&](bool storage_transposed = true) { sel.dptr = static_cast(t->columnwise_data.dptr); sel.scale_inv = t->columnwise_scale_inv.dptr; + sel.amax = t->columnwise_amax.dptr; sel.dtype = col_dtype; sel.rowwise = false; - sel.shape = create_shape_info(t, swap_dims); + sel.storage_transposed = storage_transposed; + sel.logical_tensor_shape = create_shape_info(t); }; // Helper to select row-wise storage auto use_rowwise = [&]() { sel.dptr = static_cast(t->data.dptr); sel.scale_inv = t->scale_inv.dptr; + sel.amax = t->amax.dptr; sel.dtype = row_dtype; sel.rowwise = true; - sel.shape = create_shape_info(t, /*swap_dims=*/false); + sel.logical_tensor_shape = create_shape_info(t); }; - const auto choice = choose_grouped_operand_storage(trans, is_A, mxfp8, is_fp8, non_tn_fp8_ok, - has_row, has_col, is_A ? "A" : "B"); + const auto choice = + choose_grouped_operand_storage(trans, is_A, mxfp8, is_fp8, nvfp4, fp8_block, non_tn_fp8_ok, + has_row, has_col, is_A ? "A" : "B"); sel.trans = choice.trans; if (choice.use_rowwise) { use_rowwise(); } else { - use_columnwise(choice.swap_dims); + use_columnwise(choice.storage_transposed); } return sel; } @@ -669,7 +830,8 @@ inline void init_matrix_layouts( } inline void init_matmul_desc(cublasLtMatmulDescOpaque_t &matmulDesc, cublasOperation_t op_A, - cublasOperation_t op_B, bool use_fp8, bool use_split_accumulator) { + cublasOperation_t op_B, bool use_fp8, bool use_split_accumulator, + bool use_per_group_alpha_beta) { NVTE_CHECK_CUBLAS(cublasLtMatmulDescInit(&matmulDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F)); NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_TRANSA, &op_A, @@ -681,13 +843,15 @@ inline void init_matmul_desc(cublasLtMatmulDescOpaque_t &matmulDesc, cublasOpera NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_POINTER_MODE, &pointer_mode, sizeof(pointer_mode))); - int64_t alphabeta_batch_stride = 1; - NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, - CUBLASLT_MATMUL_DESC_ALPHA_BATCH_STRIDE, - &alphabeta_batch_stride, sizeof(int64_t))); - NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, - CUBLASLT_MATMUL_DESC_BETA_BATCH_STRIDE, - &alphabeta_batch_stride, sizeof(int64_t))); + if (use_per_group_alpha_beta) { + int64_t alphabeta_batch_stride = 1; + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, + CUBLASLT_MATMUL_DESC_ALPHA_BATCH_STRIDE, + &alphabeta_batch_stride, sizeof(int64_t))); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, + CUBLASLT_MATMUL_DESC_BETA_BATCH_STRIDE, + &alphabeta_batch_stride, sizeof(int64_t))); + } // Fast accumulation is only supported for FP8 (mirrors non-grouped GEMM logic). int8_t fastAccuMode = use_split_accumulator ? 0 : static_cast(use_fp8); @@ -720,6 +884,74 @@ inline void set_mxfp8_scale_pointers(cublasLtMatmulDescOpaque_t &matmulDesc, #endif // CUBLAS_VERSION >= CUBLAS_MXFP8_GROUPED_GEMM_VERSION } +// Configures cuBLAS for NVFP4 grouped GEMM: sets VEC16_UE4M3 scale mode and scale pointers +// for both A and B. Requires cuBLAS 13.4+. +inline void set_nvfp4_scale_pointers(cublasLtMatmulDescOpaque_t &matmulDesc, + void **a_scale_inv_ptrs, void **b_scale_inv_ptrs) { +#if CUBLAS_VERSION >= CUBLAS_NVFP4_GROUPED_GEMM_VERSION + NVTE_CHECK(transformer_engine::cuda::cublas_version() >= CUBLAS_NVFP4_GROUPED_GEMM_VERSION, + "NVFP4 grouped GEMM requires cuBLAS 13.4+, but run-time cuBLAS version is ", + transformer_engine::cuda::cublas_version()); + const cublasLtMatmulMatrixScale_t scale_mode = CUBLASLT_MATMUL_MATRIX_SCALE_VEC16_UE4M3; + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_A_SCALE_MODE, + &scale_mode, sizeof(scale_mode))); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_B_SCALE_MODE, + &scale_mode, sizeof(scale_mode))); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, + CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, + &a_scale_inv_ptrs, sizeof(a_scale_inv_ptrs))); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, + CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, + &b_scale_inv_ptrs, sizeof(b_scale_inv_ptrs))); +#else + NVTE_CHECK(false, + "NVFP4 grouped GEMM requires cuBLAS 13.4+, but compile-time " + "cuBLAS version is ", + CUBLAS_VERSION); +#endif // CUBLAS_VERSION >= CUBLAS_NVFP4_GROUPED_GEMM_VERSION +} + +// Configures cuBLAS for FP8 block-scaling grouped GEMM: sets VEC128_32F or BLK128x128_32F +// scale mode and scale pointers for A and B. Requires cuBLAS 13.4+. +inline void set_fp8_block_scaling_scale_pointers(cublasLtMatmulDescOpaque_t &matmulDesc, + void **a_scale_inv_ptrs, void **b_scale_inv_ptrs, + NVTEScalingMode a_scaling_mode, + NVTEScalingMode b_scaling_mode) { +#if CUBLAS_VERSION >= CUBLAS_FP8_BLOCK_GROUPED_GEMM_VERSION + NVTE_CHECK( + transformer_engine::cuda::cublas_version() >= CUBLAS_FP8_BLOCK_GROUPED_GEMM_VERSION, + "FP8 block scaling grouped GEMM requires cuBLAS 13.4+, but run-time cuBLAS version is ", + transformer_engine::cuda::cublas_version()); + + NVTE_CHECK(!(a_scaling_mode == NVTE_BLOCK_SCALING_2D && b_scaling_mode == NVTE_BLOCK_SCALING_2D), + "Only 1D by 1D, 1D by 2D, and 2D by 1D block scaling grouped GEMM is supported, " + "but got 2D by 2D"); + + const cublasLtMatmulMatrixScale_t scale_mode_a = + a_scaling_mode == NVTE_BLOCK_SCALING_1D ? CUBLASLT_MATMUL_MATRIX_SCALE_VEC128_32F + : CUBLASLT_MATMUL_MATRIX_SCALE_BLK128x128_32F; + const cublasLtMatmulMatrixScale_t scale_mode_b = + b_scaling_mode == NVTE_BLOCK_SCALING_1D ? CUBLASLT_MATMUL_MATRIX_SCALE_VEC128_32F + : CUBLASLT_MATMUL_MATRIX_SCALE_BLK128x128_32F; + + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_A_SCALE_MODE, + &scale_mode_a, sizeof(scale_mode_a))); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, CUBLASLT_MATMUL_DESC_B_SCALE_MODE, + &scale_mode_b, sizeof(scale_mode_b))); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, + CUBLASLT_MATMUL_DESC_A_SCALE_POINTER, + &a_scale_inv_ptrs, sizeof(a_scale_inv_ptrs))); + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, + CUBLASLT_MATMUL_DESC_B_SCALE_POINTER, + &b_scale_inv_ptrs, sizeof(b_scale_inv_ptrs))); +#else + NVTE_CHECK(false, + "FP8 block scaling grouped GEMM requires cuBLAS 13.4+, but compile-time " + "cuBLAS version is ", + CUBLAS_VERSION); +#endif // CUBLAS_VERSION >= CUBLAS_FP8_BLOCK_GROUPED_GEMM_VERSION +} + // Configures cuBLAS for tensor-scaling FP8 grouped GEMM: sets PER_BATCH_SCALAR_32F scale mode // and scale pointers for A and B. Both operands are guaranteed FP8 by the caller. inline void set_fp8_scale_pointers(cublasLtMatmulDescOpaque_t &matmulDesc, void **a_scale_inv_ptrs, @@ -781,6 +1013,10 @@ inline GroupedGemmWorkspace setup_grouped_gemm_workspace(transformer_engine::Ten "Grouped GEMM setup workspace"); void *cublas_workspace_ptr = validate_and_get_workspace_ptr(wspace_cublas, cublas_workspace_size, "Grouped GEMM cuBLAS workspace"); + constexpr uintptr_t kSetupBaseAlignment = 16; + NVTE_CHECK(reinterpret_cast(setup_workspace_ptr) % kSetupBaseAlignment == 0, + "Grouped GEMM setup workspace must be ", kSetupBaseAlignment, + "-byte aligned (cuBLAS requires this for pointer arrays)."); auto setup_workspace = GroupedGemmSetupWorkspace::from_buffers( static_cast(setup_workspace_ptr), num_tensors); return {std::move(setup_workspace), cublas_workspace_ptr, num_tensors}; @@ -790,9 +1026,8 @@ inline void execute_grouped_gemm(const GroupedGemmSetupWorkspace &setup_workspac const GroupedOperandSelection &A_sel, const GroupedOperandSelection &B_sel, transformer_engine::DType d_dtype, size_t num_tensors, - bool use_split_accumulator, bool use_fp8, int64_t avg_m_val, - int64_t avg_n_val, int64_t avg_k_val, void *cublas_workspace_ptr, - cudaStream_t stream, int math_sm_count = 0) { + const GroupedGemmConfig &config, void *cublas_workspace_ptr, + cudaStream_t stream) { using cublasHandleManager = transformer_engine::detail::HandleManager; cublasLtHandle_t handle = cublasHandleManager::Instance().GetHandle(); @@ -805,26 +1040,42 @@ inline void execute_grouped_gemm(const GroupedGemmSetupWorkspace &setup_workspac num_tensors); cublasLtMatmulDescOpaque_t matmulDesc; - init_matmul_desc(matmulDesc, op_A, op_B, use_fp8, use_split_accumulator); + init_matmul_desc(matmulDesc, op_A, op_B, config.use_fp8, config.use_split_accumulator, + config.use_per_group_alpha_beta); if (transformer_engine::is_mxfp_scaling(A_sel.scaling_mode)) { set_mxfp8_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs, setup_workspace.b_scale_inv_ptrs); - } else if (use_fp8) { + } else if (transformer_engine::is_nvfp_scaling(A_sel.scaling_mode)) { + set_nvfp4_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs, + setup_workspace.b_scale_inv_ptrs); + } else if (transformer_engine::is_fp8_block_scaling(A_sel.scaling_mode)) { + set_fp8_block_scaling_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs, + setup_workspace.b_scale_inv_ptrs, A_sel.scaling_mode, + B_sel.scaling_mode); + } else if (config.use_fp8) { set_fp8_scale_pointers(matmulDesc, setup_workspace.a_scale_inv_ptrs, setup_workspace.b_scale_inv_ptrs); } - if (math_sm_count != 0) { - NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute( - &matmulDesc, CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET, &math_sm_count, sizeof(math_sm_count))); + if (config.sm_count != 0) { + NVTE_CHECK_CUBLAS(cublasLtMatmulDescSetAttribute(&matmulDesc, + CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET, + &config.sm_count, sizeof(config.sm_count))); } - cublasLtMatmulAlgo_t algo = select_grouped_gemm_algo(handle, matmulDesc, descA, descB, descC, - descD, avg_m_val, avg_n_val, avg_k_val); - - NVTE_CHECK_CUBLAS(cublasLtMatmul(handle, &matmulDesc, setup_workspace.alpha_ptrs, - setup_workspace.A_ptrs, &descA, setup_workspace.B_ptrs, &descB, - setup_workspace.beta_ptrs, setup_workspace.C_ptrs, &descC, - setup_workspace.D_ptrs, &descD, &algo, cublas_workspace_ptr, - kGroupedGemmCublasWorkspaceSize, stream)); + cublasLtMatmulAlgo_t algo = select_grouped_gemm_algo( + handle, matmulDesc, descA, descB, descC, descD, config.avg_m, config.avg_n, config.avg_k); + + // Hopper uses a single scalar alpha/beta for the whole grouped GEMM; + // Blackwell+ uses per-matrix alpha/beta arrays. + void *alpha_arg = config.use_per_group_alpha_beta + ? static_cast(setup_workspace.alpha_ptrs) + : config.alpha_dptr; + void *beta_arg = config.use_per_group_alpha_beta ? static_cast(setup_workspace.beta_ptrs) + : config.beta_dptr; + + NVTE_CHECK_CUBLAS(cublasLtMatmul(handle, &matmulDesc, alpha_arg, setup_workspace.A_ptrs, &descA, + setup_workspace.B_ptrs, &descB, beta_arg, setup_workspace.C_ptrs, + &descC, setup_workspace.D_ptrs, &descD, &algo, + cublas_workspace_ptr, kGroupedGemmCublasWorkspaceSize, stream)); } // Device helper: compute the element offset for tensor `idx` given shape metadata. @@ -871,22 +1122,64 @@ __forceinline__ __device__ int64_t padded_mxfp8_scale_inv_bytes(int64_t first, i return padded_scale_dim_y * padded_scale_dim_x; } -// Device helper: byte offset into a contiguous grouped MXFP8 scale_inv buffer for -// tensor `idx`. Each expert's scale_inv is expected to be padded -// to the 128x4 swizzled layout. -__forceinline__ __device__ int64_t compute_grouped_tensor_mxfp8_scale_inv_offset( - const TensorShapeInfo &meta, size_t idx, bool rowwise) { +__forceinline__ __device__ int64_t padded_nvfp4_scale_inv_bytes(int64_t first, int64_t last, + bool rowwise) { + namespace mxfp8_swizzle = transformer_engine::dispatch::mxfp8::swizzle; + constexpr int64_t kNvfp4BlockSize = 16; + const int64_t scale_tile_y = static_cast(mxfp8_swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_Y); + const int64_t scale_tile_x = static_cast(mxfp8_swizzle::GEMM_SWIZZLED_SCALE_TILE_DIM_X); + const int64_t scale_dim_y = rowwise ? first : last; + const int64_t data_dim_x = rowwise ? last : first; + const int64_t padded_scale_dim_y = + ((scale_dim_y + scale_tile_y - 1) / scale_tile_y) * scale_tile_y; + const int64_t scale_dim_x = (data_dim_x + kNvfp4BlockSize - 1) / kNvfp4BlockSize; + const int64_t padded_scale_dim_x = + ((scale_dim_x + scale_tile_x - 1) / scale_tile_x) * scale_tile_x; + // E4M3 scales are 1 byte per element. + return padded_scale_dim_y * padded_scale_dim_x; +} + +// FP8 block-scaling scale_inv layout matches the quantizer in get_scales() for logical dims. +__forceinline__ __device__ int64_t padded_block_1d_scale_inv_floats(int64_t first, int64_t last, + bool rowwise) { + constexpr int64_t kBlockLen = 128; + constexpr int64_t kRowAlign = 4; + const int64_t scale_dim_y = rowwise ? last : first; + const int64_t data_dim_x = rowwise ? first : last; + const int64_t y = (scale_dim_y + kBlockLen - 1) / kBlockLen; + const int64_t x = ((data_dim_x + kRowAlign - 1) / kRowAlign) * kRowAlign; + return y * x; +} + +__forceinline__ __device__ int64_t padded_block_2d_scale_inv_floats(int64_t first, int64_t last, + bool rowwise) { + constexpr int64_t kBlockLen = 128; + constexpr int64_t kRowAlign = 4; + const int64_t scale_dim_y = rowwise ? first : last; + const int64_t data_dim_x = rowwise ? last : first; + const int64_t y = (scale_dim_y + kBlockLen - 1) / kBlockLen; + const int64_t x_ceil = (data_dim_x + kBlockLen - 1) / kBlockLen; + const int64_t x = ((x_ceil + kRowAlign - 1) / kRowAlign) * kRowAlign; + return y * x; +} + +// Generic prefix-sum of per-tensor padded scale_inv sizes — used to locate where +// tensor `idx`'s scales start in a contiguous grouped scale_inv buffer. +// `PaddedFn` is a callable (int64_t first, int64_t last) -> int64_t returning the +// recipe-specific padded size (bytes for MXFP8/NVFP4, floats for FP8 block scaling). +template +__forceinline__ __device__ int64_t compute_grouped_scale_inv_offset(const TensorShapeInfo &meta, + size_t idx, PaddedFn padded) { if (meta.first_dims != nullptr || meta.last_dims != nullptr) { int64_t cumsum = 0; for (size_t i = 0; i < idx; i++) { const int64_t f = meta.first_dims ? meta.first_dims[i] : meta.uniform_first; const int64_t l = meta.last_dims ? meta.last_dims[i] : meta.uniform_last; - cumsum += padded_mxfp8_scale_inv_bytes(f, l, rowwise); + cumsum += padded(f, l); } return cumsum; } - return static_cast(idx) * - padded_mxfp8_scale_inv_bytes(meta.uniform_first, meta.uniform_last, rowwise); + return static_cast(idx) * padded(meta.uniform_first, meta.uniform_last); } // Linear scan to find which tensor contains the given row. @@ -1016,15 +1309,19 @@ __global__ void setup_grouped_gemm_kernel( void **a_scale_inv_ptrs, void **b_scale_inv_ptrs, // Inputs char *a_base, char *b_base, char *c_base, char *d_base, TensorShapeInfo A_meta, - TensorShapeInfo B_meta, TensorShapeInfo C_meta, TensorShapeInfo D_meta, size_t a_elem_size, - size_t b_elem_size, size_t c_elem_size, size_t d_elem_size, float *alpha_ptr, float *beta_ptr, + TensorShapeInfo B_meta, TensorShapeInfo C_meta, TensorShapeInfo D_meta, size_t a_bits_per_elem, + size_t b_bits_per_elem, size_t c_elem_size, size_t d_elem_size, float *alpha_ptr, + float *beta_ptr, bool use_per_group_alpha_beta, // Scale inputs: for tensor scaling, pass float* and set mxfp8_base to nullptr // For MXFP8, pass nullptr for tensor_scale and set mxfp8_base float *a_scale_base, float *b_scale_base, bool a_rowwise, bool b_rowwise, - NVTEScalingMode scaling_mode, size_t num_tensors, + bool a_storage_transposed, bool b_storage_transposed, NVTEScalingMode a_scaling_mode, + NVTEScalingMode b_scaling_mode, size_t num_tensors, MultiTensorGroupGemmInputArgs a_multi_tensor_args, MultiTensorGroupGemmOutputArgs c_multi_tensor_args, - MultiTensorGroupGemmOutputArgs d_multi_tensor_args) { + MultiTensorGroupGemmOutputArgs d_multi_tensor_args, + // NVFP4: per-group amax values and output buffer for computed alpha + float *a_amax, float *b_amax, float *nvfp4_computed_alpha) { size_t idx = blockIdx.x * blockDim.x + threadIdx.x; if (idx >= num_tensors) return; @@ -1034,10 +1331,7 @@ __global__ void setup_grouped_gemm_kernel( const bool has_d_multi_tensor = (d_base == nullptr); int64_t a_first = 0; int64_t a_last = 0; - if (has_a_multi_tensor) { - a_first = static_cast(a_multi_tensor_args.cols[idx]); - a_last = static_cast(a_multi_tensor_args.rows[idx]); - } else { + if (!has_a_multi_tensor) { a_first = A_meta.first_dims ? A_meta.first_dims[idx] : A_meta.uniform_first; a_last = A_meta.last_dims ? A_meta.last_dims[idx] : A_meta.uniform_last; } @@ -1053,21 +1347,34 @@ __global__ void setup_grouped_gemm_kernel( int64_t d_offset = compute_grouped_tensor_offset(D_meta, idx); // Compute data pointers - A_ptrs[idx] = - has_a_multi_tensor ? a_multi_tensor_args.data_ptrs[idx] : (a_base + a_offset * a_elem_size); - B_ptrs[idx] = b_base + b_offset * b_elem_size; + A_ptrs[idx] = has_a_multi_tensor ? a_multi_tensor_args.data_ptrs[idx] + : (a_base + (a_offset * a_bits_per_elem) / 8); + B_ptrs[idx] = b_base + (b_offset * b_bits_per_elem) / 8; C_ptrs[idx] = has_c_multi_tensor ? c_multi_tensor_args.data_ptrs[idx] : (c_base + c_offset * c_elem_size); D_ptrs[idx] = has_d_multi_tensor ? d_multi_tensor_args.data_ptrs[idx] : (d_base + d_offset * d_elem_size); - // Compute storage dimensions for cuBLAS matrix layouts. - // For INPUTS (A, B): Row-wise storage is seen as transposed column-major by cuBLAS, - // so rows=last, cols=first. For columnwise, dims are already swapped. - a_rows[idx] = static_cast(a_last); - a_cols[idx] = static_cast(a_first); - b_rows[idx] = static_cast(b_last); - b_cols[idx] = static_cast(b_first); + // Compute storage dimensions for cuBLAS matrix layouts from logical dims. + // Rowwise and MXFP8 columnwise storage use logical row-major layout, viewed as + // column-major rows=last, cols=first. Transposed columnwise storage reverses this. + if (has_a_multi_tensor) { + a_rows[idx] = a_multi_tensor_args.rows[idx]; + a_cols[idx] = a_multi_tensor_args.cols[idx]; + } else if (a_storage_transposed) { + a_rows[idx] = static_cast(a_first); + a_cols[idx] = static_cast(a_last); + } else { + a_rows[idx] = static_cast(a_last); + a_cols[idx] = static_cast(a_first); + } + if (b_storage_transposed) { + b_rows[idx] = static_cast(b_first); + b_cols[idx] = static_cast(b_last); + } else { + b_rows[idx] = static_cast(b_last); + b_cols[idx] = static_cast(b_first); + } if (has_d_multi_tensor) { d_rows[idx] = d_multi_tensor_args.rows[idx]; d_cols[idx] = d_multi_tensor_args.cols[idx]; @@ -1076,34 +1383,90 @@ __global__ void setup_grouped_gemm_kernel( d_cols[idx] = static_cast(d_first); } - // Fill alpha/beta pointers (per-matrix) - alpha_ptrs[idx] = alpha_ptr + idx; - beta_ptrs[idx] = beta_ptr + idx; + // Fill alpha/beta pointers. + // Hopper uses one shared alpha/beta scalar for all groups; Blackwell+ uses per-matrix scalars. + // For NVFP4 on Blackwell+: compute per-group alpha that includes global scale (amax). + // A's amax: grouped path indexes a_amax[idx]; discrete path reads amax_ptrs[idx]. + if (use_per_group_alpha_beta) { + float a_amax_val = 0.0f; + bool has_a_amax = false; + if (has_a_multi_tensor) { + auto *a_amax_p = static_cast(a_multi_tensor_args.amax_ptrs[idx]); + if (a_amax_p != nullptr) { + a_amax_val = *a_amax_p; + has_a_amax = true; + } + } else if (a_amax != nullptr) { + a_amax_val = a_amax[idx]; + has_a_amax = true; + } + if (has_a_amax && b_amax && nvfp4_computed_alpha) { + constexpr float factor_inv = 1.0f / (6.0f * 6.0f * 448.0f * 448.0f); + nvfp4_computed_alpha[idx] = alpha_ptr[idx] * a_amax_val * b_amax[idx] * factor_inv; + alpha_ptrs[idx] = &nvfp4_computed_alpha[idx]; + } else { + alpha_ptrs[idx] = alpha_ptr + idx; + } + beta_ptrs[idx] = beta_ptr + idx; + } else { + // Hopper: use single scalar for the whole grouped GEMM + alpha_ptrs[idx] = alpha_ptr; + beta_ptrs[idx] = beta_ptr; + } - // Fill scale pointers (per-matrix). - // The interpretation of the scale buffers depends on the shared scaling recipe: - // otherwise : one float per tensor, indexed by tensor index - if (a_scale_base) { - if (scaling_mode == NVTE_MXFP8_1D_SCALING) { - const int64_t a_scale_offset = - compute_grouped_tensor_mxfp8_scale_inv_offset(A_meta, idx, a_rowwise); - a_scale_inv_ptrs[idx] = reinterpret_cast( - static_cast(static_cast(a_scale_base)) + a_scale_offset); + // Fill scale pointers (per-matrix). For MXFP8/NVFP4 and FP8 block scaling, the per-expert + // scale_inv buffer is padded to a layout that depends on the recipe — offsets are computed + // from the same padded sizes that the quantizer uses at allocation, not from data_offset. + // NVTE_MXFP8_1D_SCALING : E8M0 byte stream; padded swizzled 128x4 tile, block_size=32. + // NVTE_NVFP4_1D_SCALING : E4M3 byte stream; padded swizzled 128x4 tile, block_size=16. + // NVTE_BLOCK_SCALING_1D : float32 array; ceildiv(./128) * roundup(./4) per tensor. + // NVTE_BLOCK_SCALING_2D : float32 array; ceildiv(./128) * roundup(ceildiv(./128), 4). + // otherwise (tensor) : one float per tensor, indexed by tensor index. + auto fill_scale_ptr = [&](void **ptrs, void *base, const TensorShapeInfo &meta, bool op_rowwise, + NVTEScalingMode op_scaling_mode) { + int64_t byte_offset = -1; + int64_t float_offset = -1; + switch (op_scaling_mode) { + case NVTE_MXFP8_1D_SCALING: + byte_offset = compute_grouped_scale_inv_offset(meta, idx, [=](int64_t f, int64_t l) { + return padded_mxfp8_scale_inv_bytes(f, l, op_rowwise); + }); + break; + case NVTE_NVFP4_1D_SCALING: + byte_offset = compute_grouped_scale_inv_offset(meta, idx, [=](int64_t f, int64_t l) { + return padded_nvfp4_scale_inv_bytes(f, l, op_rowwise); + }); + break; + case NVTE_BLOCK_SCALING_1D: + float_offset = compute_grouped_scale_inv_offset(meta, idx, [=](int64_t f, int64_t l) { + return padded_block_1d_scale_inv_floats(f, l, op_rowwise); + }); + break; + case NVTE_BLOCK_SCALING_2D: + float_offset = compute_grouped_scale_inv_offset(meta, idx, [=](int64_t f, int64_t l) { + return padded_block_2d_scale_inv_floats(f, l, op_rowwise); + }); + break; + default: + float_offset = static_cast(idx); + break; + } + if (byte_offset >= 0) { + ptrs[idx] = static_cast(base) + byte_offset; } else { - a_scale_inv_ptrs[idx] = static_cast(a_scale_base) + idx; + ptrs[idx] = static_cast(base) + float_offset; } + }; + + if (a_scale_base) { + fill_scale_ptr(a_scale_inv_ptrs, a_scale_base, A_meta, a_rowwise, a_scaling_mode); } else { a_scale_inv_ptrs[idx] = a_multi_tensor_args.scale_inv_ptrs[idx]; } if (b_scale_base) { - if (scaling_mode == NVTE_MXFP8_1D_SCALING) { - const int64_t b_scale_offset = - compute_grouped_tensor_mxfp8_scale_inv_offset(B_meta, idx, b_rowwise); - b_scale_inv_ptrs[idx] = reinterpret_cast( - static_cast(static_cast(b_scale_base)) + b_scale_offset); - } else { - b_scale_inv_ptrs[idx] = static_cast(b_scale_base) + idx; - } + fill_scale_ptr(b_scale_inv_ptrs, b_scale_base, B_meta, b_rowwise, b_scaling_mode); + } else { + b_scale_inv_ptrs[idx] = nullptr; } } @@ -1112,13 +1475,14 @@ inline void launch_grouped_gemm_setup( const GroupedGemmSetupWorkspace &ws, const GroupedOperandSelection &A_sel, const GroupedOperandSelection &B_sel, const transformer_engine::GroupedTensor *C, const transformer_engine::GroupedTensor *D, const transformer_engine::Tensor *alpha_tensor, - const transformer_engine::Tensor *beta_tensor, size_t num_tensors, cudaStream_t stream, + const transformer_engine::Tensor *beta_tensor, bool use_per_group_alpha_beta, + size_t num_tensors, cudaStream_t stream, const MultiTensorGroupGemmInputArgs &a_multi_tensor_args, const NVTETensor *C_list, const NVTETensor *D_list, char *a_base, transformer_engine::DType c_dtype, transformer_engine::DType d_dtype) { - // Use shape info from selection (already accounts for columnwise dimension swap) - TensorShapeInfo A_meta = A_sel.shape; - TensorShapeInfo B_meta = B_sel.shape; + // Use logical shape info from selection; storage transposes are tracked separately. + TensorShapeInfo A_meta = A_sel.logical_tensor_shape; + TensorShapeInfo B_meta = B_sel.logical_tensor_shape; TensorShapeInfo C_meta{}; TensorShapeInfo D_meta{}; @@ -1153,29 +1517,37 @@ inline void launch_grouped_gemm_setup( d_base = static_cast(D->data.dptr); } - const size_t a_elem_size = transformer_engine::typeToSize(A_sel.dtype); - const size_t b_elem_size = transformer_engine::typeToSize(B_sel.dtype); + const size_t a_bits_per_elem = transformer_engine::typeToNumBits(A_sel.dtype); + const size_t b_bits_per_elem = transformer_engine::typeToNumBits(B_sel.dtype); const size_t c_elem_size = transformer_engine::typeToSize(c_dtype); const size_t d_elem_size = transformer_engine::typeToSize(d_dtype); const int threads_per_block = 256; const int num_blocks = (num_tensors + threads_per_block - 1) / threads_per_block; - // A and B share the same scaling recipe (validated in validate_grouped_gemm_inputs). - // Pass scale buffers as void* and let the kernel interpret them via scaling_mode. + // Pass scale buffers as void* and let the kernel interpret them via each operand's scaling mode. - // Scale rowwise flag for MXFP8/NVFP4: to calculate scale_inv padding based offsets - // within kernel. Ignored for tensor scaling. const bool a_rowwise = A_sel.rowwise; const bool b_rowwise = B_sel.rowwise; + + // NVFP4 alpha needs A's amax from either A_sel.amax (grouped) or amax_ptrs (discrete). + const bool a_has_amax = (A_sel.amax != nullptr) || + (A_sel.dptr == nullptr && a_multi_tensor_args.amax_ptrs[0] != nullptr); + const bool needs_nvfp4_alpha = transformer_engine::is_nvfp_scaling(A_sel.scaling_mode) && + a_has_amax && (B_sel.amax != nullptr); + setup_grouped_gemm_kernel<<>>( ws.A_ptrs, ws.B_ptrs, ws.C_ptrs, ws.D_ptrs, ws.a_rows, ws.a_cols, ws.b_rows, ws.b_cols, ws.d_rows, ws.d_cols, ws.alpha_ptrs, ws.beta_ptrs, ws.a_scale_inv_ptrs, ws.b_scale_inv_ptrs, - A_sel.dptr, B_sel.dptr, c_base, d_base, A_meta, B_meta, C_meta, D_meta, a_elem_size, - b_elem_size, c_elem_size, d_elem_size, static_cast(alpha_tensor->data.dptr), - static_cast(beta_tensor->data.dptr), reinterpret_cast(A_sel.scale_inv), - reinterpret_cast(B_sel.scale_inv), a_rowwise, b_rowwise, A_sel.scaling_mode, - num_tensors, a_multi_tensor_args, c_multi_tensor_args, d_multi_tensor_args); + A_sel.dptr, B_sel.dptr, c_base, d_base, A_meta, B_meta, C_meta, D_meta, a_bits_per_elem, + b_bits_per_elem, c_elem_size, d_elem_size, static_cast(alpha_tensor->data.dptr), + static_cast(beta_tensor->data.dptr), use_per_group_alpha_beta, + reinterpret_cast(A_sel.scale_inv), reinterpret_cast(B_sel.scale_inv), + a_rowwise, b_rowwise, A_sel.storage_transposed, B_sel.storage_transposed, A_sel.scaling_mode, + B_sel.scaling_mode, num_tensors, a_multi_tensor_args, c_multi_tensor_args, + d_multi_tensor_args, A_sel.amax ? static_cast(A_sel.amax) : nullptr, + B_sel.amax ? static_cast(B_sel.amax) : nullptr, + needs_nvfp4_alpha ? ws.nvfp4_computed_alpha : nullptr); NVTE_CHECK_CUDA(cudaGetLastError()); } @@ -1195,9 +1567,14 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT NVTE_API_CALL(nvte_grouped_gemm); using namespace transformer_engine; - // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.3+ + // Grouped GEMM requires Blackwell (SM100) or newer with cuBLAS 13.3+, + // or Hopper (SM90) with cuBLAS 13.4+. check_grouped_gemm_requirements("nvte_grouped_gemm"); + const int current_device = transformer_engine::cuda::current_device(); + const int sm = transformer_engine::cuda::sm_arch(current_device); + const bool use_per_group_alpha_beta = grouped_gemm_supports_per_group_alpha_beta(sm); + // Convert to internal types const GroupedTensor *inputA = convertNVTEGroupedTensorCheck(A); const GroupedTensor *inputB = convertNVTEGroupedTensorCheck(B); @@ -1212,9 +1589,10 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT GroupedMatmulConfig config_ = parse_grouped_gemm_config(config); // Validate inputs and outputs. - const size_t num_tensors = validate_grouped_gemm_inputs(inputA->num_tensors, {inputA, inputB}, - alpha_tensor, beta_tensor); - validate_grouped_gemm_outputs(num_tensors, {inputC_raw, outputD}); + const size_t num_tensors = validate_grouped_gemm_inputs( + inputA->num_tensors, {inputA, inputB}, alpha_tensor, beta_tensor, use_per_group_alpha_beta); + validate_grouped_gemm_outputs(num_tensors, inputA->dtype(), inputB->dtype(), + {inputC_raw, outputD}); // If C is NULL, use D as C (valid when beta=0, cuBLAS won't read C data) const GroupedTensor *inputC = (inputC_raw != nullptr) ? inputC_raw : outputD; @@ -1223,27 +1601,44 @@ void nvte_grouped_gemm(const NVTEGroupedTensor A, int transa, const NVTEGroupedT // mirror the non-grouped GEMM logic for FP8 layout constraints. auto A_sel = select_grouped_operand(inputA, static_cast(transa), /*is_A=*/true); auto B_sel = select_grouped_operand(inputB, static_cast(transb), /*is_A=*/false); + validate_grouped_gemm_scaling_modes(A_sel.scaling_mode, B_sel.scaling_mode, sm, + "nvte_grouped_gemm"); + validate_nvfp4_grouped_gemm_support(A_sel, B_sel, use_per_group_alpha_beta); + validate_fp8_block_grouped_gemm_support(A_sel, B_sel, sm); + + // NVFP4 global-scale alpha requires per-tensor amax for both operands; without it + // the kernel silently drops the (amax_A * amax_B / factor) factor and produces + // numerically wrong output. + if (is_nvfp_scaling(A_sel.scaling_mode)) { + NVTE_CHECK(A_sel.amax != nullptr, "Grouped GEMM: NVFP4 A is missing amax."); + NVTE_CHECK(B_sel.amax != nullptr, "Grouped GEMM: NVFP4 B is missing amax."); + } // Workspaces: setup (pointer arrays) and cuBLAS auto workspace = setup_grouped_gemm_workspace(wspace_setup, wspace_cublas, num_tensors); MultiTensorGroupGemmInputArgs a_multi_tensor_args{}; launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor, - beta_tensor, num_tensors, stream, a_multi_tensor_args, - /*C_list=*/nullptr, /*D_list=*/nullptr, A_sel.dptr, inputC->dtype(), - outputD->dtype()); + beta_tensor, use_per_group_alpha_beta, num_tensors, stream, + a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, A_sel.dptr, + inputC->dtype(), outputD->dtype()); // Compute average dimensions for heuristics // K dimension: if transa, K is A's first dim; if not, K is A's last dim // Use original inputA and transa for heuristics (not modified A_sel.trans) - int64_t avg_m_val = config_.avg_m.value_or(compute_avg_first_dim(outputD)); - int64_t avg_n_val = config_.avg_n.value_or(compute_avg_last_dim(outputD)); - int64_t avg_k_val = + GroupedGemmConfig gemm_config; + gemm_config.use_split_accumulator = config_.use_split_accumulator; + gemm_config.use_fp8 = is_fp8_dtype(A_sel.dtype) || is_fp8_dtype(B_sel.dtype); + gemm_config.use_per_group_alpha_beta = use_per_group_alpha_beta; + gemm_config.alpha_dptr = alpha_tensor->data.dptr; + gemm_config.beta_dptr = beta_tensor->data.dptr; + gemm_config.avg_m = config_.avg_m.value_or(compute_avg_first_dim(outputD)); + gemm_config.avg_n = config_.avg_n.value_or(compute_avg_last_dim(outputD)); + gemm_config.avg_k = config_.avg_k.value_or(transa ? compute_avg_first_dim(inputA) : compute_avg_last_dim(inputA)); - const bool use_fp8 = is_fp8_dtype(A_sel.dtype) || is_fp8_dtype(B_sel.dtype); + gemm_config.sm_count = config_.sm_count; execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, outputD->dtype(), num_tensors, - config_.use_split_accumulator, use_fp8, avg_m_val, avg_n_val, avg_k_val, - workspace.cublas_workspace_ptr, stream, config_.sm_count); + gemm_config, workspace.cublas_workspace_ptr, stream); } void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num_a_tensors, @@ -1255,9 +1650,14 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num NVTE_API_CALL(nvte_grouped_gemm_with_discrete_inputA); using namespace transformer_engine; - // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.3+ + // Grouped GEMM requires Blackwell (SM100) or newer with cuBLAS 13.3+, + // or Hopper (SM90) with cuBLAS 13.4+. check_grouped_gemm_requirements("nvte_grouped_gemm_with_discrete_inputA"); + const int current_device = transformer_engine::cuda::current_device(); + const int sm = transformer_engine::cuda::sm_arch(current_device); + const bool use_per_group_alpha_beta = grouped_gemm_supports_per_group_alpha_beta(sm); + NVTE_CHECK(A_list != nullptr, "Grouped GEMM: A_list is null."); NVTE_CHECK(num_a_tensors > 0, "Grouped GEMM: num_a_tensors must be > 0."); @@ -1273,40 +1673,49 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num GroupedMatmulConfig config_ = parse_grouped_gemm_config(config); // Validate inputs and outputs. - const size_t num_tensors = - validate_grouped_gemm_inputs(num_a_tensors, {inputB}, alpha_tensor, beta_tensor); - - validate_grouped_gemm_outputs(num_tensors, {inputC_raw, outputD}); - - // If C is NULL, use D as C (valid when beta=0, cuBLAS won't read C data) - const GroupedTensor *inputC = (inputC_raw != nullptr) ? inputC_raw : outputD; + const size_t num_tensors = validate_grouped_gemm_inputs(num_a_tensors, {inputB}, alpha_tensor, + beta_tensor, use_per_group_alpha_beta); // Validate A list and selection auto A_list_info = validate_grouped_gemm_multi_inputA_list(A_list, num_a_tensors, num_tensors, "A"); - auto is_fp8_or_16bit = [](transformer_engine::DType dtype) { + auto is_supported_dtype = [](transformer_engine::DType dtype) { return dtype == transformer_engine::DType::kFloat8E4M3 || dtype == transformer_engine::DType::kFloat8E5M2 || + dtype == transformer_engine::DType::kFloat4E2M1 || dtype == transformer_engine::DType::kBFloat16 || dtype == transformer_engine::DType::kFloat16; }; - NVTE_CHECK(is_fp8_or_16bit(A_list_info.all_row ? A_list_info.row_dtype : A_list_info.col_dtype), - "Grouped GEMM: A_list tensors must be FP8, BF16, or FP16."); + NVTE_CHECK( + is_supported_dtype(A_list_info.all_row ? A_list_info.row_dtype : A_list_info.col_dtype), + "Grouped GEMM: A_list tensors must be FP8, NVFP4, BF16, or FP16."); // Cross-operand consistency (mirrors validate_grouped_gemm_inputs). const DType a_rep_dtype = A_list_info.all_row ? A_list_info.row_dtype : A_list_info.col_dtype; - NVTE_CHECK(is_fp8_dtype(a_rep_dtype) == is_fp8_dtype(inputB->dtype()), - "Grouped GEMM: A and B must both be FP8 or both be non-FP8."); - NVTE_CHECK(transformer_engine::is_mxfp_scaling(A_list_info.scaling_mode) == - transformer_engine::is_mxfp_scaling(inputB->scaling_mode), - "Grouped GEMM: A and B must both use MXFP8 scaling or both use tensor scaling."); - if (transformer_engine::is_mxfp_scaling(A_list_info.scaling_mode)) { + const bool a_is_fp8 = is_fp8_dtype(a_rep_dtype); + const bool b_is_fp8 = is_fp8_dtype(inputB->dtype()); + const bool a_is_fp4 = a_rep_dtype == transformer_engine::DType::kFloat4E2M1; + const bool b_is_fp4 = inputB->dtype() == transformer_engine::DType::kFloat4E2M1; + const bool a_is_low_precision = a_is_fp8 || a_is_fp4; + const bool b_is_low_precision = b_is_fp8 || b_is_fp4; + NVTE_CHECK(a_is_low_precision == b_is_low_precision, + "Grouped GEMM: A and B must both be low-precision (FP8/NVFP4) or both not."); + NVTE_CHECK(a_is_fp8 == b_is_fp8, "Grouped GEMM: A and B must both be FP8 or both be non-FP8."); + NVTE_CHECK(a_is_fp4 == b_is_fp4, + "Grouped GEMM: A and B must both be NVFP4 or both be non-NVFP4."); + validate_grouped_gemm_scaling_modes(A_list_info.scaling_mode, inputB->scaling_mode, sm, + "nvte_grouped_gemm_with_discrete_inputA"); + if (transformer_engine::is_mxfp_scaling(A_list_info.scaling_mode) || + transformer_engine::is_nvfp_scaling(A_list_info.scaling_mode)) { NVTE_CHECK(A_list_info.with_gemm_swizzled_scales, - "MXFP8 grouped GEMM: A scales must be swizzled for GEMM."); - NVTE_CHECK(inputB->with_gemm_swizzled_scales, - "MXFP8 grouped GEMM: B scales must be swizzled for GEMM."); + "Grouped GEMM: A scales must be swizzled for GEMM (MXFP8/NVFP4)."); } + validate_grouped_gemm_outputs(num_tensors, a_rep_dtype, inputB->dtype(), {inputC_raw, outputD}); + + // If C is NULL, use D as C (valid when beta=0, cuBLAS won't read C data) + const GroupedTensor *inputC = (inputC_raw != nullptr) ? inputC_raw : outputD; + // Select operand storage for B (row-wise vs column-wise) auto B_sel = select_grouped_operand(inputB, static_cast(transb), /*is_A=*/false); @@ -1314,56 +1723,72 @@ void nvte_grouped_gemm_with_discrete_inputA(const NVTETensor *A_list, size_t num A_sel.scaling_mode = A_list_info.scaling_mode; A_sel.with_gemm_swizzled_scales = A_list_info.with_gemm_swizzled_scales; A_sel.trans = static_cast(transa); + validate_nvfp4_grouped_gemm_support(A_sel, B_sel, use_per_group_alpha_beta); + validate_fp8_block_grouped_gemm_support(A_sel, B_sel, sm); const DType rep_dtype = A_list_info.all_row ? A_list_info.row_dtype : A_list_info.col_dtype; const bool is_fp8 = is_fp8_dtype(rep_dtype); - const bool non_tn_fp8_ok = nvte_is_non_tn_fp8_gemm_supported(); const bool mxfp8 = transformer_engine::is_mxfp_scaling(A_list_info.scaling_mode); + const bool nvfp4 = transformer_engine::is_nvfp_scaling(A_list_info.scaling_mode); + const bool fp8_block = transformer_engine::is_fp8_block_scaling(A_list_info.scaling_mode); + // FP8 block scaling on Hopper requires TN layout (matches select_grouped_operand logic for B). + const bool non_tn_fp8_ok = fp8_block ? false : nvte_is_non_tn_fp8_gemm_supported(); int64_t avg_first_dim = 0; int64_t avg_last_dim = 0; MultiTensorGroupGemmInputArgs a_multi_tensor_args{}; - const auto choice = - choose_grouped_operand_storage(static_cast(transa), /*is_A=*/true, mxfp8, is_fp8, - non_tn_fp8_ok, A_list_info.all_row, A_list_info.all_col, "A"); + const auto choice = choose_grouped_operand_storage(static_cast(transa), /*is_A=*/true, + mxfp8, is_fp8, nvfp4, fp8_block, non_tn_fp8_ok, + A_list_info.all_row, A_list_info.all_col, "A"); A_sel.trans = choice.trans; A_sel.rowwise = choice.use_rowwise; + A_sel.storage_transposed = choice.storage_transposed; if (choice.use_rowwise) { NVTE_CHECK(A_list_info.all_row, "Grouped GEMM: A_list is missing row-wise data"); A_sel.dtype = A_list_info.row_dtype; - a_multi_tensor_args = build_grouped_gemm_multi_inputA_args( - A_list, num_a_tensors, /*use_rowwise=*/true, is_fp8, &avg_first_dim, &avg_last_dim, "A"); } else { NVTE_CHECK(A_list_info.all_col, "Grouped GEMM: A_list is missing column-wise data"); A_sel.dtype = A_list_info.col_dtype; - a_multi_tensor_args = build_grouped_gemm_multi_inputA_args( - A_list, num_a_tensors, /*use_rowwise=*/false, is_fp8, &avg_first_dim, &avg_last_dim, "A"); } + a_multi_tensor_args = build_grouped_gemm_multi_inputA_args( + A_list, num_a_tensors, choice.use_rowwise, choice.storage_transposed, &avg_first_dim, + &avg_last_dim, "A"); - // For discrete A_list, scale pointers are per-tensor; use multi-tensor args. - // Base pointer is unused when providing per-tensor pointers. + // Discrete A_list: per-tensor pointers come from `a_multi_tensor_args` (data/scale/amax). A_sel.scale_inv = nullptr; A_sel.dptr = nullptr; + A_sel.amax = nullptr; + + if (nvfp4) { + for (size_t i = 0; i < num_tensors; ++i) { + NVTE_CHECK(a_multi_tensor_args.amax_ptrs[i] != nullptr, "Grouped GEMM: NVFP4 A_list tensor ", + i, " is missing amax."); + } + NVTE_CHECK(B_sel.amax != nullptr, "Grouped GEMM: NVFP4 B is missing amax."); + } // Workspaces: setup (pointer arrays) and cuBLAS auto workspace = setup_grouped_gemm_workspace(wspace_setup, wspace_cublas, num_tensors); launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, inputC, outputD, alpha_tensor, - beta_tensor, num_tensors, stream, a_multi_tensor_args, - /*C_list=*/nullptr, /*D_list=*/nullptr, nullptr, inputC->dtype(), - outputD->dtype()); - - // Compute average dimensions for heuristics - int64_t avg_m_val = config_.avg_m.value_or(compute_avg_first_dim(outputD)); - int64_t avg_n_val = + beta_tensor, use_per_group_alpha_beta, num_tensors, stream, + a_multi_tensor_args, /*C_list=*/nullptr, /*D_list=*/nullptr, nullptr, + inputC->dtype(), outputD->dtype()); + + GroupedGemmConfig gemm_config; + gemm_config.use_split_accumulator = config_.use_split_accumulator; + gemm_config.use_fp8 = is_fp8_dtype(A_sel.dtype) || is_fp8_dtype(B_sel.dtype); + gemm_config.use_per_group_alpha_beta = use_per_group_alpha_beta; + gemm_config.alpha_dptr = alpha_tensor->data.dptr; + gemm_config.beta_dptr = beta_tensor->data.dptr; + gemm_config.avg_m = config_.avg_m.value_or(compute_avg_first_dim(outputD)); + gemm_config.avg_n = config_.avg_n.value_or(transb ? compute_avg_first_dim(inputB) : compute_avg_last_dim(inputB)); - int64_t avg_k_val = - config_.avg_k.value_or(static_cast(transa) ? avg_last_dim : avg_first_dim); - const bool use_fp8 = is_fp8_dtype(A_sel.dtype) || is_fp8_dtype(B_sel.dtype); + gemm_config.avg_k = config_.avg_k.value_or(transa ? avg_first_dim : avg_last_dim); + gemm_config.sm_count = config_.sm_count; execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, outputD->dtype(), num_tensors, - config_.use_split_accumulator, use_fp8, avg_m_val, avg_n_val, avg_k_val, - workspace.cublas_workspace_ptr, stream, config_.sm_count); + gemm_config, workspace.cublas_workspace_ptr, stream); } void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, @@ -1376,9 +1801,14 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, NVTE_API_CALL(nvte_grouped_gemm_with_discrete_out); using namespace transformer_engine; - // Grouped GEMM requires Blackwell (SM100) or newer and cuBLAS 13.3+ + // Grouped GEMM requires Blackwell (SM100) or newer with cuBLAS 13.3+, + // or Hopper (SM90) with cuBLAS 13.4+. check_grouped_gemm_requirements("nvte_grouped_gemm_with_discrete_out"); + const int current_device = transformer_engine::cuda::current_device(); + const int sm = transformer_engine::cuda::sm_arch(current_device); + const bool use_per_group_alpha_beta = grouped_gemm_supports_per_group_alpha_beta(sm); + NVTE_CHECK(D_list != nullptr, "Grouped GEMM: D_list is null."); NVTE_CHECK(num_d_tensors > 0, "Grouped GEMM: num_d_tensors must be > 0."); if (num_c_tensors > 0) { @@ -1395,20 +1825,15 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, const Tensor *d0 = convertNVTETensorCheck(D_list[0]); const DType d_dtype = d0->dtype(); - const size_t num_tensors = validate_grouped_gemm_inputs(inputA->num_tensors, {inputA, inputB}, - alpha_tensor, beta_tensor); + const size_t num_tensors = validate_grouped_gemm_inputs( + inputA->num_tensors, {inputA, inputB}, alpha_tensor, beta_tensor, use_per_group_alpha_beta); NVTE_CHECK(num_d_tensors == num_tensors, "Grouped GEMM: D_list must have num_tensors (", num_tensors, ") entries, got ", num_d_tensors); if (num_c_tensors > 0) { NVTE_CHECK(num_c_tensors == num_tensors, "Grouped GEMM: C_list must have num_tensors (", num_tensors, ") entries, got ", num_c_tensors); } - auto is_output_dtype = [](transformer_engine::DType dtype) { - return dtype == transformer_engine::DType::kBFloat16 || - dtype == transformer_engine::DType::kFloat16 || - dtype == transformer_engine::DType::kFloat32; - }; - NVTE_CHECK(is_output_dtype(d_dtype), "Grouped GEMM: D must be BF16, FP16, or FP32."); + validate_grouped_gemm_output_dtype(inputA->dtype(), inputB->dtype(), d_dtype, "D"); // Parse config (if provided) GroupedMatmulConfig config_ = parse_grouped_gemm_config(config); @@ -1417,25 +1842,41 @@ void nvte_grouped_gemm_with_discrete_out(const NVTEGroupedTensor A, int transa, // mirror the non-grouped GEMM logic for FP8 layout constraints. auto A_sel = select_grouped_operand(inputA, static_cast(transa), /*is_A=*/true); auto B_sel = select_grouped_operand(inputB, static_cast(transb), /*is_A=*/false); + validate_grouped_gemm_scaling_modes(A_sel.scaling_mode, B_sel.scaling_mode, sm, + "nvte_grouped_gemm_with_discrete_out"); + validate_nvfp4_grouped_gemm_support(A_sel, B_sel, use_per_group_alpha_beta); + validate_fp8_block_grouped_gemm_support(A_sel, B_sel, sm); + + // NVFP4 global-scale alpha requires per-tensor amax for both operands. + if (is_nvfp_scaling(A_sel.scaling_mode)) { + NVTE_CHECK(A_sel.amax != nullptr, "Grouped GEMM: NVFP4 A is missing amax."); + NVTE_CHECK(B_sel.amax != nullptr, "Grouped GEMM: NVFP4 B is missing amax."); + } + // Workspaces: setup (pointer arrays) and cuBLAS auto workspace = setup_grouped_gemm_workspace(wspace_setup, wspace_cublas, num_tensors); MultiTensorGroupGemmInputArgs a_multi_tensor_args{}; launch_grouped_gemm_setup(workspace.setup_workspace, A_sel, B_sel, /*C=*/nullptr, /*D=*/nullptr, - alpha_tensor, beta_tensor, num_tensors, stream, a_multi_tensor_args, - C_list, D_list, A_sel.dptr, d_dtype, d_dtype); - - // Compute average dimensions for heuristics - int64_t avg_m_val = + alpha_tensor, beta_tensor, use_per_group_alpha_beta, num_tensors, + stream, a_multi_tensor_args, C_list, D_list, A_sel.dptr, d_dtype, + d_dtype); + + GroupedGemmConfig gemm_config; + gemm_config.use_split_accumulator = config_.use_split_accumulator; + gemm_config.use_fp8 = is_fp8_dtype(A_sel.dtype) || is_fp8_dtype(B_sel.dtype); + gemm_config.use_per_group_alpha_beta = use_per_group_alpha_beta; + gemm_config.alpha_dptr = alpha_tensor->data.dptr; + gemm_config.beta_dptr = beta_tensor->data.dptr; + gemm_config.avg_m = config_.avg_m.value_or(transa ? compute_avg_last_dim(inputA) : compute_avg_first_dim(inputA)); - int64_t avg_n_val = + gemm_config.avg_n = config_.avg_n.value_or(transb ? compute_avg_first_dim(inputB) : compute_avg_last_dim(inputB)); - int64_t avg_k_val = + gemm_config.avg_k = config_.avg_k.value_or(transa ? compute_avg_first_dim(inputA) : compute_avg_last_dim(inputA)); - const bool use_fp8 = is_fp8_dtype(A_sel.dtype) || is_fp8_dtype(B_sel.dtype); - execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, d_dtype, num_tensors, - config_.use_split_accumulator, use_fp8, avg_m_val, avg_n_val, avg_k_val, - workspace.cublas_workspace_ptr, stream, config_.sm_count); + gemm_config.sm_count = config_.sm_count; + execute_grouped_gemm(workspace.setup_workspace, A_sel, B_sel, d_dtype, num_tensors, gemm_config, + workspace.cublas_workspace_ptr, stream); } namespace { diff --git a/transformer_engine/common/libtransformer_engine.version b/transformer_engine/common/libtransformer_engine.version index 706c237cc..ccd18fc15 100644 --- a/transformer_engine/common/libtransformer_engine.version +++ b/transformer_engine/common/libtransformer_engine.version @@ -7,6 +7,7 @@ transformer_engine::cuda::supports_multicast*; transformer_engine::cuda::stream_priority_range*; transformer_engine::cuda::current_device*; + transformer_engine::cuda::cublas_version*; transformer_engine::cuda_driver::get_symbol*; transformer_engine::cuda_driver::ensure_context_exists*; transformer_engine::ubuf_built_with_mpi*; diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index 561f64d59..ca7eccab0 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -356,6 +356,8 @@ static void CheckGroupedScaleInv(const GroupedTensor &t, const std::string &name // Determine expected dtype based on data type and scaling mode if (is_fp8_dtype(t.dtype()) && is_tensor_scaling(t.scaling_mode)) { check_scales(DType::kFloat32); + } else if (is_fp8_block_scaling(t.scaling_mode)) { + check_scales(DType::kFloat32); } else if (is_mxfp8_scaling(t.scaling_mode)) { check_scales(DType::kFloat8E8M0); } else if (is_nvfp4_scaling(t.scaling_mode)) { diff --git a/transformer_engine/pytorch/cpp_extensions/gemm.py b/transformer_engine/pytorch/cpp_extensions/gemm.py index edf2c1e1c..48751d8c9 100644 --- a/transformer_engine/pytorch/cpp_extensions/gemm.py +++ b/transformer_engine/pytorch/cpp_extensions/gemm.py @@ -5,7 +5,6 @@ """Python interface for GEMM extensions""" from typing import Iterable, Optional, Tuple, Union, List -import ctypes import os import functools import torch @@ -420,20 +419,8 @@ def general_grouped_gemm( @functools.lru_cache(maxsize=None) def get_grouped_gemm_setup_workspace_size(num_tensors: int) -> int: - """Return workspace size for grouped GEMM pointer setup. - Must match GroupedGemmSetupWorkspace::required_setup_size in cublaslt_grouped_gemm.cu. - """ - ptr_bytes = ctypes.sizeof(ctypes.c_void_p) - int_bytes = ctypes.sizeof(ctypes.c_int) - ptr_size = num_tensors * ptr_bytes - int_size = num_tensors * int_bytes - k_ptr_alignment = 16 - # Each pointer array is placed at a 16-byte-aligned offset (matching kPtrAlignment in C++). - # aligned_ptr_size = round_up(num_tensors * ptr_bytes, 16) - aligned_ptr_size = ((ptr_size + k_ptr_alignment - 1) // k_ptr_alignment) * k_ptr_alignment - size = 8 * aligned_ptr_size + 6 * int_size - alignment = 256 - return ((size + alignment - 1) // alignment) * alignment + """Return workspace size for grouped GEMM pointer setup.""" + return tex.get_grouped_gemm_setup_workspace_size(num_tensors) @functools.lru_cache(maxsize=None) @@ -510,13 +497,18 @@ def general_grouped_gemm_for_grouped_tensor( rowwise = B.rowwise_data device = rowwise.device if rowwise is not None else B.columnwise_data.device + # Hopper (SM90) uses a single shared alpha/beta scalar; + # Blackwell+ (SM100) supports per-group alpha/beta arrays. + per_group = torch.cuda.get_device_capability() >= (10, 0) + num_alphabeta = num_tensors if per_group else 1 + if alpha is None: - alpha = _get_fp32_ones_tensor(num_tensors, device) + alpha = _get_fp32_ones_tensor(num_alphabeta, device) if beta is None: if accumulate: - beta = _get_fp32_ones_tensor(num_tensors, device) + beta = _get_fp32_ones_tensor(num_alphabeta, device) else: - beta = _get_fp32_zeros_tensor(num_tensors, device) + beta = _get_fp32_zeros_tensor(num_alphabeta, device) if not alpha.is_cuda or not beta.is_cuda: raise ValueError("alpha and beta must be CUDA tensors.") diff --git a/transformer_engine/pytorch/csrc/extensions/gemm.cpp b/transformer_engine/pytorch/csrc/extensions/gemm.cpp index 9cb1fb7f5..53ea76d83 100644 --- a/transformer_engine/pytorch/csrc/extensions/gemm.cpp +++ b/transformer_engine/pytorch/csrc/extensions/gemm.cpp @@ -86,10 +86,13 @@ GroupedGemmConfig prepare_grouped_gemm_config(at::Tensor alpha, at::Tensor beta, at::Tensor workspace_setup, at::Tensor workspace_cublas, size_t num_tensors, int math_sm_count, bool use_split_accumulator) { - NVTE_CHECK(alpha.numel() == static_cast(num_tensors), - "Grouped GEMM expects alpha to have num_tensors elements."); - NVTE_CHECK(beta.numel() == static_cast(num_tensors), - "Grouped GEMM expects beta to have num_tensors elements."); + const bool per_group = (alpha.numel() == static_cast(num_tensors)); + const bool scalar = (alpha.numel() == 1); + NVTE_CHECK(per_group || scalar, "Grouped GEMM expects alpha to have 1 or num_tensors (", + num_tensors, ") elements, got ", alpha.numel()); + NVTE_CHECK(beta.numel() == alpha.numel(), + "Grouped GEMM expects beta to have the same number of elements as alpha (", + alpha.numel(), "), got ", beta.numel()); GroupedGemmConfig grouped_gemm_config{ makeTransformerEngineTensor(alpha), diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index b5b563882..3acef587f 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -283,6 +283,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("split_quantize", &transformer_engine::pytorch::split_quantize, "Split and multi-tensor quantize", py::arg("tensor"), py::arg("split_sections"), py::arg("quantizer_list"), py::arg("disable_bulk_allocation") = false); + m.def("get_grouped_gemm_setup_workspace_size", &nvte_get_grouped_gemm_setup_workspace_size, + "Required workspace size for grouped GEMM setup"); m.def("te_general_grouped_gemm", &transformer_engine::pytorch::te_general_grouped_gemm, "Grouped GEMM"); m.def("te_general_grouped_gemm_for_grouped_tensor", From 5f1eaffde87336ef3f1bd877044a7159e02e569d Mon Sep 17 00:00:00 2001 From: Xin Yao Date: Thu, 28 May 2026 05:06:01 +0800 Subject: [PATCH 085/180] [PyTorch] Enable head dim 256 for FA4 (#2932) * enable head dim 256 for FA4 Signed-off-by: Xin Yao * update CI, fix lint, resolve comments Signed-off-by: Xin Yao * resolve comments Signed-off-by: Xin Yao * update filter Signed-off-by: Xin Yao --------- Signed-off-by: Xin Yao --- tests/pytorch/attention/test_attention.py | 32 ++++++-- .../dot_product_attention/backends.py | 2 + .../attention/dot_product_attention/utils.py | 75 +++++++++++++------ 3 files changed, 79 insertions(+), 30 deletions(-) diff --git a/tests/pytorch/attention/test_attention.py b/tests/pytorch/attention/test_attention.py index 5c46949f6..401bd6f01 100644 --- a/tests/pytorch/attention/test_attention.py +++ b/tests/pytorch/attention/test_attention.py @@ -344,12 +344,36 @@ def test_dpa_num_splits(dtype, model_configs, model): @pytest.mark.skipif( not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." ) -@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_base]) @pytest.mark.parametrize("model", model_configs_fa4_base.keys()) def test_dpa_fa4_base(dtype, model_configs, model): - """Test DotProductAttention with FA4: base configs, extended head dims, GQA, num_splits""" + """Test DotProductAttention with FA4: base configs, GQA, num_splits""" + test_dot_product_attention(dtype, model_configs, model, False, True, None, False, False) + + +# head_dim=256 is supported only on SM100 via FA4's dedicated kernel +# (flash_attn/cute/sm100_hd256_2cta_fmha_*.py), available in flash-attn-4 > 4.0.0b10. +# On other architectures, _validate_head_dims rejects (256, 256), FA4 is disabled, and +# the test would silently fall back to another backend — defeating the purpose. Gate +# explicitly so the CI signal is unambiguous. +model_configs_fa4_hdim256 = { + "fa4_hdim256": ModelConfig(2, 1024, 8, 256, attn_mask_type="causal"), +} + + +@pytest.mark.skipif( + not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." +) +@pytest.mark.skipif( + device_compute_capability not in ((10, 0), (10, 3)), + reason="FA4 head_dim=256 dedicated kernel is SM100/103-only.", +) +@pytest.mark.parametrize("dtype", param_types_lean) +@pytest.mark.parametrize("model_configs", [model_configs_fa4_hdim256]) +@pytest.mark.parametrize("model", model_configs_fa4_hdim256.keys()) +def test_dpa_fa4_hdim256(dtype, model_configs, model): + """Test DotProductAttention with FA4: head_dim=256 dedicated kernel on SM100""" test_dot_product_attention(dtype, model_configs, model, False, True, None, False, False) @@ -369,7 +393,6 @@ def test_dpa_fa4_base(dtype, model_configs, model): @pytest.mark.skipif( not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." ) -@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_mla]) @pytest.mark.parametrize("model", model_configs_fa4_mla.keys()) @@ -396,7 +419,6 @@ def test_dpa_fa4_mla(dtype, model_configs, model): @pytest.mark.skipif( not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." ) -@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_swa]) @pytest.mark.parametrize("model", model_configs_fa4_swa.keys()) @@ -420,7 +442,6 @@ def test_dpa_fa4_sliding_window(dtype, model_configs, model, qkv_layout): @pytest.mark.skipif( not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." ) -@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_varlen]) @pytest.mark.parametrize("model", model_configs_fa4_varlen.keys()) @@ -446,7 +467,6 @@ def test_dpa_fa4_varlen(dtype, model_configs, model, qkv_layout): @pytest.mark.skipif( not FlashAttentionUtils.v4_is_installed, reason="Flash-attn v4 (flash-attn-4) is required." ) -@pytest.mark.skipif(get_cudnn_version() < (8, 9, 1), reason="cuDNN 8.9.1+ is required.") @pytest.mark.parametrize("dtype", param_types_lean) @pytest.mark.parametrize("model_configs", [model_configs_fa4_mask]) @pytest.mark.parametrize("model", model_configs_fa4_mask.keys()) diff --git a/transformer_engine/pytorch/attention/dot_product_attention/backends.py b/transformer_engine/pytorch/attention/dot_product_attention/backends.py index 6c6adc6e3..9f2dadb68 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/backends.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/backends.py @@ -167,8 +167,10 @@ from flash_attn.cute.interface import ( # pylint: disable=ungrouped-imports,no-name-in-module flash_attn_func as flash_attn_func_v4, flash_attn_varlen_func as flash_attn_varlen_func_v4, + _validate_head_dims as _fa4_validate_head_dims, ) + fa_utils.v4_validate_head_dims = _fa4_validate_head_dims fa_utils.set_flash_attention_4_params() # Float8CurrentScaling: fused_attn_bwd takes O in FP8 by default, this flag allows it in F16 diff --git a/transformer_engine/pytorch/attention/dot_product_attention/utils.py b/transformer_engine/pytorch/attention/dot_product_attention/utils.py index 6565e9f6f..989b65f19 100644 --- a/transformer_engine/pytorch/attention/dot_product_attention/utils.py +++ b/transformer_engine/pytorch/attention/dot_product_attention/utils.py @@ -7,7 +7,7 @@ """ import math import os -from typing import Any, Dict, List, Optional, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, Union import warnings import logging import functools @@ -147,8 +147,11 @@ class FlashAttentionUtils: fa4_version = PkgVersion("0") use_v4 = False v4_installation_steps = """\ -pip install flash-attn-4==4.0.0b8 nvidia-cutlass-dsl[cu13]""" +pip install flash-attn-4==4.0.0b11 nvidia-cutlass-dsl[cu13]""" v4_warning_printed = False + # Set by backends.py if FA4 is installed; calls flash_attn.cute.interface._validate_head_dims + # which raises AssertionError for unsupported (head_dim, head_dim_v) combinations. + v4_validate_head_dims: Callable = None @staticmethod def set_flash_attention_version(): @@ -802,21 +805,25 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt ) use_flash_attention_3 = False - if use_flash_attention_4 and FlashAttentionUtils.v4_is_installed: - # FA4 head dimension support is architecture-dependent - # (matches _validate_head_dims in flash_attn.cute.interface): - # SM90: head_dim <= 256 and head_dim_v <= 256 - # SM100/110: head_dim <= 128 and head_dim_v <= 128, - # OR DeepSeek MLA shape (head_dim=192, head_dim_v=128) - # SM80/120: constrained by shared memory (~256 max in practice) - _fa4_hdim_ok = True - if (10, 0) <= device_compute_capability < (12, 0): - _is_standard = head_dim_qk <= 128 and head_dim_v <= 128 - _is_deepseek = head_dim_qk == 192 and head_dim_v == 128 - _fa4_hdim_ok = _is_standard or _is_deepseek - else: - _fa4_hdim_ok = head_dim_qk <= 256 and head_dim_v <= 256 - if not _fa4_hdim_ok: + if ( + use_flash_attention_4 + and FlashAttentionUtils.v4_is_installed + and FlashAttentionUtils.v4_validate_head_dims is not None + ): + # Defer to FA4's own _validate_head_dims to keep TE in sync with FA4 supported shapes + # (e.g., (256, 256) on SM100, (192, 128) DeepSeek, (64, 512) MLA-absorbed). + # The function asserts on unsupported combinations; SM80/SM120 have no validation branch + # in FA4 so the call passes through silently for those archs. + _fa4_alignment = 16 // torch.empty(0, dtype=qkv_dtype).element_size() + try: + # pylint: disable-next=not-callable + FlashAttentionUtils.v4_validate_head_dims( + head_dim_qk, + head_dim_v, + device_compute_capability[0], + _fa4_alignment, + ) + except AssertionError: logger.debug( "Disabling FlashAttention 4 due to unsupported head dimensions. " "Found: head_dim_qk = %s, head_dim_v = %s, on sm%s.", @@ -825,13 +832,33 @@ def _is_fa3_supported(num_heads, num_gqa_groups, head_dim_qk, head_dim_v, qkv_dt device_compute_capability[0] * 10 + device_compute_capability[1], ) use_flash_attention_4 = False - # Workaround: SM100 backward kernel bug when MLA + 2CTA (head_dim_qk >= 128). - # FlashAttentionBackwardSm100 computes dK_reduce_ncol = gcd(32, tile_hdim // 2) - # based on Q/K head_dim but reuses it for dV TMEM load atoms. When - # (tile_hdimv // 2) % dK_reduce_ncol != 0, dV reads are misaligned. - # See: flash_attn/cute/flash_bwd_sm100.py, line ~262 and ~3890. - elif ( - _fa4_hdim_ok + # flash-attn-4 4.0.0b11 validates (256, 256) on SM100, but its dedicated + # hd256 kernel diverges from the reference for cross-attention/decode-like + # shapes such as sq=1, skv=2048. Keep FA4 enabled for the self-attention + # hd256 path covered by the dedicated test, and fall back for cross-attn. + if ( + use_flash_attention_4 + and (10, 0) <= device_compute_capability < (12, 0) + and head_dim_qk == head_dim_v == 256 + and max_seqlen_q != max_seqlen_kv + ): + logger.debug( + "Disabling FlashAttention 4 for SM100 head_dim=256 cross-attention. " + "Found: max_seqlen_q = %s, max_seqlen_kv = %s.", + max_seqlen_q, + max_seqlen_kv, + ) + use_flash_attention_4 = False + # Workaround: SM100 backward kernel bug when MLA + 2CTA (head_dim_qk >= 128) for + # the standard (non-dedicated) kernel path. FlashAttentionBackwardSm100 computes + # dK_reduce_ncol = gcd(32, tile_hdim // 2) based on Q/K head_dim but reuses it for + # dV TMEM load atoms. When (tile_hdimv // 2) % dK_reduce_ncol != 0, dV reads are + # misaligned (e.g. dqk=128, dv=96 gives 48 % 32 != 0). The dedicated (256, 256) + # kernel uses its own tmem layout and is not affected. + # See: flash_attn/cute/flash_bwd_sm100.py ~L262 and ~L3890. Still present in + # flash-attn-4 4.0.0b11. + if ( + use_flash_attention_4 and is_training and head_dim_qk != head_dim_v and head_dim_qk >= 128 From f3c2e74dd78e82515f0b9084c9a88308cbcf9a0b Mon Sep 17 00:00:00 2001 From: XiaomingFun233 <74387760+XiaomingFun233@users.noreply.github.com> Date: Fri, 29 May 2026 02:00:18 +0800 Subject: [PATCH 086/180] [fused_router][pytorch] Optimize naive topk path and add perf benchmark (#2776) * fused_router: keep low-risk CUDA optimizations - restore forward hot paths to baseline behavior for topk/scores kernels\n- keep warp-level reduction helper for backward normalization\n- handle empty expert_bias safely in fused topk forward Signed-off-by: Xinhao Wei * fused_router: specialize naive_topk_and_mask for topk<=8 Add a lightweight register-based small-k path and keep the generic fallback for compatibility. Signed-off-by: Xinhao Wei * tests: add fused router performance benchmark Add CUDA perf benchmark for fused topk router, aux-loss score, and moe aux-loss kernels. Signed-off-by: Xinhao Wei * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * fused_router: address review feedback Signed-off-by: Xinhao Wei --------- Signed-off-by: Xinhao Wei Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- tests/pytorch/test_fused_router_perf.py | 371 ++++++++++++++++++ .../fused_score_for_moe_aux_loss.cu | 6 +- .../fused_topk_with_score_function.cu | 33 +- .../common/fused_router/utils.h | 95 ++++- 4 files changed, 484 insertions(+), 21 deletions(-) create mode 100644 tests/pytorch/test_fused_router_perf.py diff --git a/tests/pytorch/test_fused_router_perf.py b/tests/pytorch/test_fused_router_perf.py new file mode 100644 index 000000000..122d19dd7 --- /dev/null +++ b/tests/pytorch/test_fused_router_perf.py @@ -0,0 +1,371 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import os +from typing import Callable, Optional, Tuple + +import pytest +import torch + +from transformer_engine.pytorch.router import ( + fused_compute_score_for_moe_aux_loss, + fused_moe_aux_loss, + fused_topk_with_score_function, +) + + +SEED = 42 + + +def _set_seed() -> None: + torch.manual_seed(SEED) + if torch.cuda.is_available(): + torch.cuda.manual_seed(SEED) + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or os.getenv("TE_RUN_PERF_TESTS", "0") != "1", + reason=( + "Benchmark test - run with: TE_RUN_PERF_TESTS=1 pytest" + " tests/pytorch/test_fused_router_perf.py" + ), +) + + +def _benchmark_cuda_kernel(fn: Callable[[], object], warmup: int = 20, iters: int = 100) -> float: + start_event = torch.cuda.Event(enable_timing=True) + end_event = torch.cuda.Event(enable_timing=True) + + for _ in range(warmup): + fn() + torch.cuda.synchronize() + + start_event.record() + for _ in range(iters): + fn() + end_event.record() + torch.cuda.synchronize() + + return start_event.elapsed_time(end_event) / iters + + +def group_limited_topk( + scores: torch.Tensor, + topk: int, + num_tokens: int, + num_experts: int, + num_groups: int, + group_topk: int, +) -> Tuple[torch.Tensor, torch.Tensor]: + group_scores = ( + scores.view(num_tokens, num_groups, -1).topk(topk // group_topk, dim=-1)[0].sum(dim=-1) + ) + group_idx = torch.topk(group_scores, k=group_topk, dim=-1, sorted=False)[1] + group_mask = torch.zeros_like(group_scores) + group_mask.scatter_(1, group_idx, 1) + + score_mask = ( + group_mask.unsqueeze(-1) + .expand(num_tokens, num_groups, num_experts // num_groups) + .reshape(num_tokens, -1) + ) + masked_scores = scores.masked_fill(~score_mask.bool(), float("-inf")) + probs, top_indices = torch.topk(masked_scores, k=topk, dim=-1) + return probs, top_indices + + +def topk_softmax_sigmoid_pytorch( + logits: torch.Tensor, + topk: int, + use_pre_softmax: bool = False, + num_groups: Optional[int] = None, + group_topk: Optional[int] = None, + scaling_factor: Optional[float] = None, + score_function: str = "softmax", + expert_bias: Optional[torch.Tensor] = None, +) -> Tuple[torch.Tensor, torch.Tensor]: + num_tokens, num_experts = logits.shape + + def compute_topk(scores, topk_value, num_groups_value=None, group_topk_value=None): + if group_topk_value: + assert num_groups_value is not None + return group_limited_topk( + scores=scores, + topk=topk_value, + num_tokens=num_tokens, + num_experts=num_experts, + num_groups=num_groups_value, + group_topk=group_topk_value, + ) + return torch.topk(scores, k=topk_value, dim=1) + + if score_function == "softmax": + if use_pre_softmax: + scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits) + probs, top_indices = compute_topk(scores, topk, num_groups, group_topk) + else: + scores, top_indices = compute_topk(logits, topk, num_groups, group_topk) + probs = torch.softmax(scores, dim=-1, dtype=torch.float32).type_as(logits) + elif score_function == "sigmoid": + scores = torch.sigmoid(logits.float()).type_as(logits) + if expert_bias is not None: + scores_for_routing = scores + expert_bias + _, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk) + scores = torch.gather(scores, dim=1, index=top_indices).type_as(logits) + else: + scores, top_indices = compute_topk(scores, topk, num_groups, group_topk) + probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores + else: + raise ValueError(f"Invalid score_function: {score_function}") + + if scaling_factor: + probs = probs * scaling_factor + + topk_masked_gates = torch.zeros_like(logits).scatter(1, top_indices, probs) + topk_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool() + return topk_masked_gates, topk_map + + +def compute_scores_for_aux_loss_pytorch( + logits: torch.Tensor, topk: int, score_function: str +) -> Tuple[torch.Tensor, torch.Tensor]: + if score_function == "softmax": + scores = torch.softmax(logits, dim=-1, dtype=torch.float32) + elif score_function == "sigmoid": + scores = torch.sigmoid(logits) + scores = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores + else: + raise ValueError(f"Invalid score_function: {score_function}") + + _, top_indices = torch.topk(scores, k=topk, dim=1) + routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool() + return routing_map, scores + + +def aux_loss_pytorch( + probs: torch.Tensor, + tokens_per_expert: torch.Tensor, + total_num_tokens: int, + topk: int, + num_experts: int, + moe_aux_loss_coeff: float, +) -> torch.Tensor: + aggregated_probs_per_expert = probs.sum(dim=0) + return torch.sum(aggregated_probs_per_expert * tokens_per_expert) * ( + num_experts * moe_aux_loss_coeff / (topk * total_num_tokens * total_num_tokens) + ) + + +def _make_router_logits( + dtype: torch.dtype, num_tokens: int, num_experts: int, score_function: str +) -> torch.Tensor: + if score_function == "sigmoid": + offset = torch.arange(-num_tokens // 2, num_tokens // 2, dtype=dtype, device="cuda") * 1e-4 + logits = ( + torch.arange(-num_experts // 2, num_experts // 2, device="cuda", dtype=dtype) * 1e-2 + ) + return logits.unsqueeze(0).repeat(num_tokens, 1) + offset.unsqueeze(1) + + logits = ( + torch.arange( + -num_tokens * num_experts // 2, + num_tokens * num_experts // 2, + device="cuda", + dtype=dtype, + ) + * 1e-4 + ) + return logits.view(num_tokens, num_experts) + + +def _make_router_bias(num_experts: int) -> torch.Tensor: + bias = torch.arange(num_experts, device="cuda", dtype=torch.float32) * 0.1 + return torch.flip(bias, dims=[0]) + + +def _print_perf_result(case_name: str, torch_ms: float, fused_ms: float) -> None: + speedup = torch_ms / fused_ms + print(f"{case_name}: torch={torch_ms:.6f} ms, fused={fused_ms:.6f} ms, speedup={speedup:.4f}x") + + +@pytest.mark.parametrize( + "score_function,use_pre_softmax,enable_bias", + [("softmax", False, False), ("sigmoid", False, True)], + ids=["softmax", "sigmoid_with_bias"], +) +def test_fused_topk_router_perf_against_torch( + score_function, use_pre_softmax, enable_bias, record_property +): + _set_seed() + + dtype = torch.float32 + num_tokens = 4096 + num_experts = 192 + topk = 8 + num_groups = 8 + group_topk = 4 + scaling_factor = 1.2 + + logits = _make_router_logits(dtype, num_tokens, num_experts, score_function) + expert_bias = _make_router_bias(num_experts) if enable_bias else None + + torch_probs, torch_map = topk_softmax_sigmoid_pytorch( + logits=logits, + topk=topk, + use_pre_softmax=use_pre_softmax, + num_groups=num_groups, + group_topk=group_topk, + scaling_factor=scaling_factor, + score_function=score_function, + expert_bias=expert_bias, + ) + fused_probs, fused_map = fused_topk_with_score_function( + logits=logits, + topk=topk, + use_pre_softmax=use_pre_softmax, + num_groups=num_groups, + group_topk=group_topk, + scaling_factor=scaling_factor, + score_function=score_function, + expert_bias=expert_bias, + ) + + torch_ms = _benchmark_cuda_kernel( + lambda: topk_softmax_sigmoid_pytorch( + logits=logits, + topk=topk, + use_pre_softmax=use_pre_softmax, + num_groups=num_groups, + group_topk=group_topk, + scaling_factor=scaling_factor, + score_function=score_function, + expert_bias=expert_bias, + ) + ) + fused_ms = _benchmark_cuda_kernel( + lambda: fused_topk_with_score_function( + logits=logits, + topk=topk, + use_pre_softmax=use_pre_softmax, + num_groups=num_groups, + group_topk=group_topk, + scaling_factor=scaling_factor, + score_function=score_function, + expert_bias=expert_bias, + ) + ) + + record_property("torch_ms", round(torch_ms, 6)) + record_property("fused_ms", round(fused_ms, 6)) + record_property("speedup", round(torch_ms / fused_ms, 6)) + _print_perf_result(f"topk_router[{score_function}]", torch_ms, fused_ms) + + torch.testing.assert_close(torch_probs, fused_probs) + torch.testing.assert_close(torch_map, fused_map) + + +@pytest.mark.parametrize("score_function", ["softmax", "sigmoid"]) +def test_fused_scores_for_aux_loss_perf_against_torch(score_function, record_property): + _set_seed() + + dtype = torch.float32 + num_tokens = 8192 + num_experts = 128 + topk = 8 + logits = _make_router_logits(dtype, num_tokens, num_experts, score_function) + + torch_map, torch_scores = compute_scores_for_aux_loss_pytorch( + logits=logits, + topk=topk, + score_function=score_function, + ) + fused_map, fused_scores = fused_compute_score_for_moe_aux_loss( + logits=logits, + topk=topk, + score_function=score_function, + ) + + torch_ms = _benchmark_cuda_kernel( + lambda: compute_scores_for_aux_loss_pytorch( + logits=logits, + topk=topk, + score_function=score_function, + ) + ) + fused_ms = _benchmark_cuda_kernel( + lambda: fused_compute_score_for_moe_aux_loss( + logits=logits, + topk=topk, + score_function=score_function, + ) + ) + + record_property("torch_ms", round(torch_ms, 6)) + record_property("fused_ms", round(fused_ms, 6)) + record_property("speedup", round(torch_ms / fused_ms, 6)) + _print_perf_result(f"scores_for_aux_loss[{score_function}]", torch_ms, fused_ms) + + torch.testing.assert_close(torch_scores, fused_scores) + torch.testing.assert_close(torch_map, fused_map) + + +def test_fused_moe_aux_loss_perf_against_torch(record_property): + _set_seed() + + dtype = torch.float32 + num_tokens = 8192 + num_experts = 128 + topk = 4 + coeff = 0.01 + + offset = torch.arange(-num_tokens // 2, num_tokens // 2, dtype=dtype, device="cuda") * 1e-4 + probs = torch.arange(-num_experts // 2, num_experts // 2, device="cuda", dtype=dtype) * 1e-2 + probs = probs.unsqueeze(0).repeat(num_tokens, 1) + offset.unsqueeze(1) + probs = probs.view(num_tokens, num_experts) + tokens_per_expert = torch.randint(1, 1000, (num_experts,), device="cuda", dtype=torch.int32) + + torch_loss = aux_loss_pytorch( + probs=probs, + tokens_per_expert=tokens_per_expert, + total_num_tokens=num_tokens, + topk=topk, + num_experts=num_experts, + moe_aux_loss_coeff=coeff, + ) + fused_loss = fused_moe_aux_loss( + probs=probs, + tokens_per_expert=tokens_per_expert, + total_num_tokens=num_tokens, + num_experts=num_experts, + topk=topk, + coeff=coeff, + ) + + torch_ms = _benchmark_cuda_kernel( + lambda: aux_loss_pytorch( + probs=probs, + tokens_per_expert=tokens_per_expert, + total_num_tokens=num_tokens, + topk=topk, + num_experts=num_experts, + moe_aux_loss_coeff=coeff, + ) + ) + fused_ms = _benchmark_cuda_kernel( + lambda: fused_moe_aux_loss( + probs=probs, + tokens_per_expert=tokens_per_expert, + total_num_tokens=num_tokens, + num_experts=num_experts, + topk=topk, + coeff=coeff, + ) + ) + + record_property("torch_ms", round(torch_ms, 6)) + record_property("fused_ms", round(fused_ms, 6)) + record_property("speedup", round(torch_ms / fused_ms, 6)) + _print_perf_result("moe_aux_loss", torch_ms, fused_ms) + + torch.testing.assert_close(torch_loss, fused_loss) diff --git a/transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu b/transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu index 4eb4240d7..675f071ab 100644 --- a/transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu +++ b/transformer_engine/common/fused_router/fused_score_for_moe_aux_loss.cu @@ -270,11 +270,7 @@ __global__ void fused_score_for_moe_aux_loss_backward_kernel(const CompType *int for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { local_sum_Output_x_Grad += local_grad[i] * act_output[i]; } - // Warp reduce the sum - for (int s = 16; s > 0; s /= 2) { - local_sum_Output_x_Grad += __shfl_xor_sync(0xffffffff, local_sum_Output_x_Grad, s); - } - CompType sum_Output_x_Grad = local_sum_Output_x_Grad; + CompType sum_Output_x_Grad = warp_reduce_sum_float(local_sum_Output_x_Grad); // In-place update for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { local_grad[i] = local_grad[i] / (sum_fwd_input + epsilon) - diff --git a/transformer_engine/common/fused_router/fused_topk_with_score_function.cu b/transformer_engine/common/fused_router/fused_topk_with_score_function.cu index 9f7a83054..9b8d8b929 100644 --- a/transformer_engine/common/fused_router/fused_topk_with_score_function.cu +++ b/transformer_engine/common/fused_router/fused_topk_with_score_function.cu @@ -284,15 +284,24 @@ void fused_topk_with_score_function_forward(const Tensor logits, int num_tokens, Tensor intermediate_output, cudaStream_t stream) { TE_ROUTER_PROBS_TYPE_SWITCH_ALL( logits.data.dtype, DataType, - TE_ROUTER_PROBS_TYPE_SWITCH_ALL( - expert_bias.data.dtype, BiasType, - fused_topk_with_score_function_forward_kernel_launcher( - reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, - use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, - reinterpret_cast(expert_bias.data.dptr), - reinterpret_cast(probs.data.dptr), - reinterpret_cast(routing_map.data.dptr), - reinterpret_cast(intermediate_output.data.dptr), stream););); + if (expert_bias.has_data()) { + TE_ROUTER_PROBS_TYPE_SWITCH_ALL( + expert_bias.data.dtype, BiasType, + fused_topk_with_score_function_forward_kernel_launcher( + reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, + use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, + reinterpret_cast(expert_bias.data.dptr), + reinterpret_cast(probs.data.dptr), + reinterpret_cast(routing_map.data.dptr), + reinterpret_cast(intermediate_output.data.dptr), stream);); + } else { + fused_topk_with_score_function_forward_kernel_launcher( + reinterpret_cast(logits.data.dptr), num_tokens, num_experts, topk, + use_pre_softmax, num_groups, group_topk, scaling_factor, score_function, nullptr, + reinterpret_cast(probs.data.dptr), + reinterpret_cast(routing_map.data.dptr), + reinterpret_cast(intermediate_output.data.dptr), stream); + }); } template @@ -399,11 +408,7 @@ __global__ void fused_topk_with_score_function_backward_kernel( local_sum_Output_x_Grad += local_grad[i] * act_output[i]; } } - // Warp reduce the sum - for (int s = 16; s > 0; s /= 2) { - local_sum_Output_x_Grad += __shfl_xor_sync(0xffffffff, local_sum_Output_x_Grad, s); - } - CompType sum_Output_x_Grad = local_sum_Output_x_Grad; + CompType sum_Output_x_Grad = warp_reduce_sum_float(local_sum_Output_x_Grad); // In-place update for (int i = lane_id; i < num_experts; i += kThreadsPerWarp) { if (local_routing_map[i]) { diff --git a/transformer_engine/common/fused_router/utils.h b/transformer_engine/common/fused_router/utils.h index 08ad3d16a..09087aff9 100644 --- a/transformer_engine/common/fused_router/utils.h +++ b/transformer_engine/common/fused_router/utils.h @@ -53,6 +53,15 @@ enum ReduceFuncType { MAX, }; +__device__ inline float warp_reduce_sum_float(float val) { + // __shfl_down_sync accumulates the total only in lane 0; + // the broadcast below is required so every lane sees the final sum. + for (int offset = kThreadsPerWarp / 2; offset > 0; offset /= 2) { + val += __shfl_down_sync(0xffffffff, val, offset); + } + return __shfl_sync(0xffffffff, val, 0); +} + template __device__ inline T warp_reduce_on_shmem(T *data_ptr, int data_size, ReduceFuncType type, int lane_id) { @@ -429,8 +438,57 @@ __device__ inline void radix_topk_and_mask(CompType *scores, int data_size, int __syncwarp(); } -__device__ inline void naive_topk_and_mask(CompType *scores, int data_size, int topk, - int *topk_indices, CompType *topk_scores, int lane_id) { +template +__device__ inline void naive_topk_and_mask_smallk(CompType *scores, int data_size, + int *topk_indices, CompType *topk_scores, + int lane_id) { + static_assert(K > 0 && K <= 8, "K must be in [1, 8]"); + int selected[K]; +#pragma unroll + for (int i = 0; i < K; ++i) { + selected[i] = -1; + } + +#pragma unroll + for (int k = 0; k < K; ++k) { + CompType val = -std::numeric_limits::infinity(); + int index = (lane_id < data_size) ? lane_id : -1; + for (int i = lane_id; i < data_size; i += kThreadsPerWarp) { + bool masked = false; +#pragma unroll + for (int j = 0; j < k; ++j) { + masked |= (selected[j] == i); + } + if (masked) continue; + CompType cur_val = scores[i]; + if (cur_val > val) { + val = cur_val; + index = i; + } + } + for (int s = kThreadsPerWarp / 2; s > 0; s /= 2) { + auto shuffled_val = __shfl_xor_sync(0xffffffff, val, s); + auto shuffled_index = __shfl_xor_sync(0xffffffff, index, s); + if (shuffled_val > val) { + val = shuffled_val; + index = shuffled_index; + } + } + + CompType chosen_val = __shfl_sync(0xffffffff, val, 0); + int chosen_index = __shfl_sync(0xffffffff, index, 0); + if (lane_id == 0) { + topk_indices[k] = chosen_index; + topk_scores[k] = chosen_val; + } + selected[k] = chosen_index; + __syncwarp(); + } +} + +__device__ inline void naive_topk_and_mask_generic(CompType *scores, int data_size, int topk, + int *topk_indices, CompType *topk_scores, + int lane_id) { // Check if the index is masked by the later iteration auto is_masked = [&topk_indices](int k, int index) { if (k == 0) return false; @@ -475,6 +533,39 @@ __device__ inline void naive_topk_and_mask(CompType *scores, int data_size, int } } +__device__ inline void naive_topk_and_mask(CompType *scores, int data_size, int topk, + int *topk_indices, CompType *topk_scores, int lane_id) { + switch (topk) { + case 1: + naive_topk_and_mask_smallk<1>(scores, data_size, topk_indices, topk_scores, lane_id); + break; + case 2: + naive_topk_and_mask_smallk<2>(scores, data_size, topk_indices, topk_scores, lane_id); + break; + case 3: + naive_topk_and_mask_smallk<3>(scores, data_size, topk_indices, topk_scores, lane_id); + break; + case 4: + naive_topk_and_mask_smallk<4>(scores, data_size, topk_indices, topk_scores, lane_id); + break; + case 5: + naive_topk_and_mask_smallk<5>(scores, data_size, topk_indices, topk_scores, lane_id); + break; + case 6: + naive_topk_and_mask_smallk<6>(scores, data_size, topk_indices, topk_scores, lane_id); + break; + case 7: + naive_topk_and_mask_smallk<7>(scores, data_size, topk_indices, topk_scores, lane_id); + break; + case 8: + naive_topk_and_mask_smallk<8>(scores, data_size, topk_indices, topk_scores, lane_id); + break; + default: + naive_topk_and_mask_generic(scores, data_size, topk, topk_indices, topk_scores, lane_id); + break; + } +} + template __device__ __forceinline__ void topk_and_mask(CompType *scores, int data_size, int topk, int *topk_indices, CompType *topk_scores, From 439ca21038052271d77123df28afdad1a9272384 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Thu, 28 May 2026 12:50:53 -0700 Subject: [PATCH 087/180] [JAX] Support new JAX triton_kernel_call_ffi for cuda-graph support (#3055) Signed-off-by: Jeremy Berchtold Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> --- .../jax/triton_extensions/utils.py | 24 +++++++++++++------ transformer_engine/jax/version_utils.py | 4 ++++ 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/transformer_engine/jax/triton_extensions/utils.py b/transformer_engine/jax/triton_extensions/utils.py index 332bc6ddb..95ee370c8 100644 --- a/transformer_engine/jax/triton_extensions/utils.py +++ b/transformer_engine/jax/triton_extensions/utils.py @@ -44,14 +44,17 @@ from packaging import version from jax import core +from jaxlib.mlir import ir import jax import jax.numpy as jnp from ..version_utils import ( TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION, + TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION, TRITON_EXTENSION_MIN_JAX_VERSION, is_triton_autotuned_alias_safe, is_triton_extension_supported, + jax_version_meet_requirement, ) @@ -626,12 +629,19 @@ def _normalize_grid(grid_tuple): else: ffi_operand_output_aliases = None - # Use JAX FFI lowering with compressed protobuf - rule = jax.ffi.ffi_lowering( - "triton_kernel_call", # Custom call target registered in gpu_triton.py - api_version=2, - backend_config=zlib.compress(call_proto), - operand_output_aliases=ffi_operand_output_aliases, - ) + compressed_call_proto = zlib.compress(call_proto) + if jax_version_meet_requirement(TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION): + rule = jax.ffi.ffi_lowering( + "triton_kernel_call_ffi", + backend_config={"opaque": ir.StringAttr.get(compressed_call_proto)}, + operand_output_aliases=ffi_operand_output_aliases, + ) + else: + rule = jax.ffi.ffi_lowering( + "triton_kernel_call", # Custom call target registered in gpu_triton.py + api_version=2, + backend_config=compressed_call_proto, + operand_output_aliases=ffi_operand_output_aliases, + ) return rule(ctx, *array_args) diff --git a/transformer_engine/jax/version_utils.py b/transformer_engine/jax/version_utils.py index e6ed9a8ea..e4619d867 100644 --- a/transformer_engine/jax/version_utils.py +++ b/transformer_engine/jax/version_utils.py @@ -25,6 +25,9 @@ def jax_version_meet_requirement(version: str): # Minimum JAX version required for Triton kernel dispatch (jaxlib < 0.8.0 segfaults). TRITON_EXTENSION_MIN_JAX_VERSION = "0.8.0" +# Minimum JAX version for non-legacy Triton kernel FFI (supports CUDA graph capture). +TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION = "0.10.1" + # Nightly and stable floors for safe input_output_aliases in TritonAutotunedKernelCall. # jaxlib/gpu/triton_kernels.cc had a bug in the autotuning save/restore loop: # it iterated over all declared aliases unconditionally, but input_copies only @@ -76,5 +79,6 @@ def is_triton_extension_supported() -> bool: "is_triton_autotuned_alias_safe", "is_triton_extension_supported", "TRITON_EXTENSION_MIN_JAX_VERSION", + "TRITON_EXTENSION_CUDA_GRAPH_MIN_JAX_VERSION", "TRITON_AUTOTUNED_INPUT_OUTPUT_ALIAS_MIN_JAX_VERSION", ] From ace2a9653a2da74a8576a39ef7d1181c7c7923cb Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Thu, 28 May 2026 15:06:01 -0700 Subject: [PATCH 088/180] [PyTorch] Allocate grouped linear wgrads as tensor views (#3049) * Allocate grouped linear wgrads as views Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Tim Moon Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../pytorch/csrc/extensions/allocate.cpp | 10 ++++++++++ transformer_engine/pytorch/module/grouped_linear.py | 12 ++++++------ .../pytorch/ops/basic/grouped_linear.py | 10 +++++----- .../pytorch/ops/fused/backward_grouped_mlp.py | 11 ++++++----- 4 files changed, 27 insertions(+), 16 deletions(-) diff --git a/transformer_engine/pytorch/csrc/extensions/allocate.cpp b/transformer_engine/pytorch/csrc/extensions/allocate.cpp index f972f8a2d..62ed7f373 100644 --- a/transformer_engine/pytorch/csrc/extensions/allocate.cpp +++ b/transformer_engine/pytorch/csrc/extensions/allocate.cpp @@ -12,6 +12,16 @@ namespace transformer_engine { namespace pytorch { +/* Allocate multiple PyTorch tensors backed by the same buffer. + * + * Use with caution and avoid exposing externally. + * + * In order to reduce CPU overhead, we compute pointer offsets + * manually and construct PyTorch tensors with raw pointers. The + * backing buffer is deallocated once the final tensor is destroyed. + * Stream usage is not recorded, so there may be race conditions if + * compute is performed on multiple streams. + */ std::vector bulk_allocate(const std::vector> &shapes, const std::vector &dtypes, std::optional device, diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index 627144345..b2baf1729 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -496,13 +496,13 @@ def backward( if ctx.fuse_wgrad_accumulation: wgrad_list = main_grads else: - weight_shape = list(weights[0].size()) - wgrad_list = tex.bulk_allocate( - [weight_shape] * ctx.num_gemms, - [ctx.activation_dtype] * ctx.num_gemms, - ctx.device, - [256] * ctx.num_gemms, # alignment + wgrad_packed = torch.empty( + ctx.num_gemms, + *weights[0].size(), + dtype=ctx.activation_dtype, + device=ctx.device, ) + wgrad_list = [wgrad_packed[i] for i in range(ctx.num_gemms)] if ctx.save_original_input: inp = inputmats[0] diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index 1f00d9228..dc15bc63b 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -1393,12 +1393,12 @@ def _fuser_backward_split_quantize( ] accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: - grad_weights = tex.bulk_allocate( - [weight_shape] * num_groups, - [ctx.dtype] * num_groups, - device, - [256] * num_groups, # alignment + grad_weights_packed = torch.empty( + grouped_shape, + dtype=ctx.dtype, + device=device, ) + grad_weights = [grad_weights_packed[i] for i in range(num_groups)] final_weight_grads = list(grad_weights) # Perform dgrad GEMMs diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 24aaafc1e..802c1a25d 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -197,12 +197,13 @@ def _compute_grad_params( w_list = [get_main_grad_from_param(w, op_label=op_label) for w in weights] accumulate_into_main_grad = get_accumulate_flag_in_param(weights[0]) else: - w_list = tex.bulk_allocate( - [weight_shape] * num_groups, - [dtype] * num_groups, - device, - [256] * num_groups, # alignment + wgrad_packed = torch.empty( + num_groups, + *weight_shape, + dtype=dtype, + device=device, ) + w_list = [wgrad_packed[i] for i in range(num_groups)] wgrad_output = w_list if ctx.weight_requires_grad: From 9e5a847c90c436c9c9c5f3e756977cf568967ccc Mon Sep 17 00:00:00 2001 From: Tim Moon <4406448+timmoon10@users.noreply.github.com> Date: Thu, 28 May 2026 16:18:33 -0700 Subject: [PATCH 089/180] Optimize function that loads pointers on GPU (#3001) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Remove unnecessary heap allocations Avoid constructing temporary std::vector when converting NVTEBasicTensor to SimpleTensor. Avoid string operations in multi-tensor swizzle. Avoid temporary std::vector when checking scale tensors. Signed-off-by: Tim Moon * Avoid heap allocation in Tensor::flat_first_dim/flat_last_dim Tensor::shape() returns a std::vector by value, allocating on the heap. flat_first_dim and flat_last_dim only need to walk the dims, so the allocation was pure overhead in hot paths. Introduce Tensor::compute_shape() returning an NVTEShape (fixed inline buffer, no heap) as the single source of truth for the format-dependent shape logic. shape() is now a thin std::vector wrapper around it for callers that want a vector; flat_first_dim and flat_last_dim call compute_shape() directly. Signed-off-by: Tim Moon Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Tim Moon * Add Tensor::flat_2d_dims() to compute both matrix dims in one pass flat_first_dim() and flat_last_dim() each called compute_shape() independently. flat_2d_dims() computes both in a single pass; the scalar helpers now delegate to it. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Tim Moon * Use flat_2d_dims() throughout common lib Replace all paired flat_first_dim() + flat_last_dim() calls on the same tensor with a single flat_2d_dims() call. Saves one compute_shape() per tensor in CheckScaleTensorShape, the multi-tensor swizzle loop, and various cast/GEMM dispatch paths. Also adds reserve() to the local vectors in nvte_multi_tensor_swizzle_scaling_factors to avoid reallocation. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Tim Moon * Generalize API for CUDA-Graph-safe copy to GPU. Signed-off-by: Tim Moon * Dedup swizzle logic in get_device_pointer_for_data_and_scales Replace the inline swizzle implementation with a call to multi_tensor_swizzle_scales_for_gemm, which has identical logic (16B-aligned contiguous output buffer, TensorWrapper construction, nvte_multi_tensor_swizzle_scaling_factors kernel). Swizzled pointers are read back from the updated TensorWrappers after the call. Add reserve() to vectors in multi_tensor_swizzle_scales_for_gemm_impl now that this function is on the hot path for get_device_pointer_for_data_and_scales. Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Tim Moon * Make separate functions for load data_ptrs and swizzle + load data_ptrs. Signed-off-by: Tim Moon * Change function name to nvte_load_value_on_device Signed-off-by: Tim Moon * Fix code review issues before opening PR - Use size_t in kernel tail loop (was int64_t) - Zero-initialize Payload before memcpy (Payload{}) - Rename Payload members to kMaxBytes/kVectorSize/kMaxVectors (linter) - Consistent at::empty shape pattern: {static_cast(N)} - Drop intermediate swizzled_scales_bytes variable - Add comment explaining uniform-stride assumption in transform_and_load_data_ptrs_on_device - Rename sfb_buffer -> _sfb_buffer (keepalive, not directly used) Co-Authored-By: Claude Sonnet 4.6 Signed-off-by: Tim Moon * Formatter and review suggestions from @greptile-apps Signed-off-by: Tim Moon * Add Shape class wrapping NVTEShape Provides a std::vector-like interface around NVTEShape without heap allocation, used as the return type of Tensor::shape() in place of the previous std::vector. Disambiguate cute::Shape from transformer_engine::Shape in the hadamard_transform kernels. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Tim Moon * Make SimpleTensor stack-allocatable Store shape in Shape class rather than std::vector. Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Make Shape conversion constructors explicit Signed-off-by: Tim Moon * Make conversion from Shape to std::vector explicit Signed-off-by: Tim Moon * Add batched NVTETensor create/destroy Expose nvte_create_tensors and nvte_destroy_tensors so multi-tensor callers can amortize the TensorAllocator mutex across N tensors instead of locking once per call. nvte_destroy_tensors was already defined internally but not declared in the public header. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Tim Moon * Use batched NVTETensor allocator in transform_and_load_data_ptrs_on_device The uniform swizzle path constructed 2N TensorWrappers and then extracted their raw NVTETensors into separate vectors. Replace with a single 2N nvte_create_tensors call into one contiguous buffer (inputs in the first half, outputs in the second), an RAII guard for nvte_destroy_tensors, and a local set_param lambda for the setters. Drops the separate pack pass and reduces the allocator mutex acquisitions from 4N to 2 per call. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Tim Moon * Expand usage of batched NVTETensor allocator Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Tim Moon * Use string_view in tensor checking functions Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Tweak function names Signed-off-by: Tim Moon * Pass std::string_view by value in Check*Tensor helpers string_view is already a (ptr, len) reference — passing by const-ref adds an indirection without benefit. Matches the C++ Core Guidelines F.16 recommendation. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Tim Moon * Review suggestions from @ptrendx Expand internal usage of Shape class. Zero-initialize in Shape::resize. Make sure dynamic smem querying is per-device. Reuse logic for batched and single tensor alloc/dealloc. Signed-off-by: Tim Moon * Add MultiTensorWrapper for batched NVTETensor allocation Thin RAII wrapper around a batched nvte_create_tensors / nvte_destroy_tensors pair, with operator[], data(), iteration, and implicit conversion to NVTETensor* for multi-tensor C APIs. Replaces the ad-hoc DestroyGuard struct used at each call site in recipe.cpp, swizzle.cpp, and utils.cpp. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Tim Moon Co-authored-by: Claude Opus 4.7 (1M context) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../common/cast/dispatch/quantize.cuh | 9 +- .../common/cast/fp8/quantize_fp8.cuh | 5 +- .../common/cast/mxfp8/dequantize_mxfp8.cuh | 3 +- .../cast/mxfp8/group_quantize_mxfp8.cuh | 2 +- .../common/cast/mxfp8/quantize_mxfp8.cuh | 5 +- .../common/cast/nvfp4/dequantize_nvfp4.cuh | 3 +- .../nvfp4/group_quantize_transpose_nvfp4.cuh | 5 +- .../cast/nvfp4/quantize_transpose_nvfp4.cuh | 8 +- .../quantize_transpose_nvfp4_tuned_1D.cuh | 5 +- .../common/comm_gemm/comm_gemm.cpp | 27 +- transformer_engine/common/common.h | 205 +++++++++++--- .../common/gemm/cublaslt_gemm.cu | 6 +- ...cast_col_hadamard_transform_cast_fusion.cu | 14 +- .../group_hadamard_transform_cast_fusion.cu | 15 +- ...cast_col_hadamard_transform_cast_fusion.cu | 14 +- .../hadamard_transform_cast_fusion.cu | 14 +- ...cast_col_hadamard_transform_cast_fusion.cu | 13 +- .../transformer_engine/transformer_engine.h | 93 +++++++ .../common/include/transformer_engine/utils.h | 23 +- .../common/normalization/common.cpp | 4 +- .../common/normalization/common.h | 6 +- .../common/normalization/layernorm/ln_api.cpp | 5 +- .../normalization/rmsnorm/rmsnorm_api.cpp | 3 +- transformer_engine/common/swizzle/swizzle.cu | 41 +-- .../common/transformer_engine.cpp | 155 +++++------ .../common/transpose/cast_transpose_fusion.cu | 5 +- ...quantize_transpose_vector_blockwise_fp4.cu | 2 +- .../common/transpose/transpose_fusion.cu | 2 +- transformer_engine/common/util/utils.cu | 79 ++++-- transformer_engine/pytorch/csrc/extensions.h | 13 +- .../pytorch/csrc/extensions/pybind.cpp | 15 +- .../pytorch/csrc/extensions/recipe.cpp | 33 +-- .../pytorch/csrc/extensions/swizzle.cpp | 89 +++--- .../pytorch/csrc/extensions/utils.cpp | 256 +++++++++--------- .../pytorch/ops/fused/backward_grouped_mlp.py | 23 +- .../pytorch/ops/fused/forward_grouped_mlp.py | 20 +- 36 files changed, 724 insertions(+), 496 deletions(-) diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh index 316243c97..bad53a03c 100644 --- a/transformer_engine/common/cast/dispatch/quantize.cuh +++ b/transformer_engine/common/cast/dispatch/quantize.cuh @@ -98,8 +98,7 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output, CheckOutputTensor(*output_tensor, "output", false); // Choose kernel - int32_t rows = input_tensor->flat_first_dim(); - int32_t cols = input_tensor->flat_last_dim(); + const auto [rows, cols] = input_tensor->flat_2d_dims(); auto dtype = input_tensor->dtype(); const bool row_scaled_nvfp4 = output_tensor->row_scaled_nvfp4; const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; @@ -260,8 +259,7 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens CheckOutputTensor(*output_tensor, "output", false); // Choose kernel - int32_t rows = grad_tensor->flat_first_dim(); - int32_t cols = grad_tensor->flat_last_dim(); + const auto [rows, cols] = grad_tensor->flat_2d_dims(); auto dtype = grad_tensor->dtype(); const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; NVTE_CHECK(nvfp4_use_4over6 || output_tensor->nvfp4_e4m3_max == 448, @@ -396,8 +394,7 @@ void group_quantize_fwd_host_aware_helper(const NVTETensor input, NVTETensor *ou // output list here is allowed to have empty tensor // Choose kernel - int32_t rows = input_tensor->flat_first_dim(); - int32_t cols = input_tensor->flat_last_dim(); + const auto [rows, cols] = input_tensor->flat_2d_dims(); auto dtype = input_tensor->dtype(); const bool nvfp4_use_4over6 = quant_config_cpp.nvfp4_4over6_mode != kNVTENVFP44Over6Disabled; diff --git a/transformer_engine/common/cast/fp8/quantize_fp8.cuh b/transformer_engine/common/cast/fp8/quantize_fp8.cuh index 96a42b494..bad10c954 100644 --- a/transformer_engine/common/cast/fp8/quantize_fp8.cuh +++ b/transformer_engine/common/cast/fp8/quantize_fp8.cuh @@ -391,8 +391,7 @@ void quantize_2D(const Tensor &input, const Tensor *act_input, Tensor *output, T using namespace quantize_2D_kernel; checkCuDriverContext(stream); - const size_t rows = input.flat_first_dim(); - const size_t cols = input.flat_last_dim(); + const auto [rows, cols] = input.flat_2d_dims(); const size_t chunks_Y = DIVUP(rows, FP8_CHUNK_DIM_Y); const size_t chunks_X = DIVUP(cols, FP8_CHUNK_DIM_X); const size_t blocks_Y = chunks_Y; @@ -406,7 +405,7 @@ void quantize_2D(const Tensor &input, const Tensor *act_input, Tensor *output, T if constexpr (IS_DBIAS) { NVTE_CHECK(dbias->data.dtype == input.data.dtype, "DBias must have the same type as input."); - NVTE_CHECK(dbias->data.shape == std::vector{cols}, "Wrong shape of DBias."); + NVTE_CHECK(dbias->data.shape == Shape{cols}, "Wrong shape of DBias."); NVTE_CHECK(workspace != nullptr, "Workspace must be a tensor."); if (workspace->data.dptr == nullptr) { diff --git a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh index 6441a567a..1face261b 100644 --- a/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/dequantize_mxfp8.cuh @@ -261,8 +261,7 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) const size_t scale_dim_X_rowwise = use_rowwise_scaling ? 32 : 1; const size_t scale_dim_Y_colwise = use_colwise_scaling ? 32 : 1; - const size_t rows = input.flat_first_dim(); - const size_t cols = input.flat_last_dim(); + const auto [rows, cols] = input.flat_2d_dims(); const size_t chunks_Y = DIVUP(rows, CHUNK_DIM_Y); const size_t chunks_X = DIVUP(cols, CHUNK_DIM_X); diff --git a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh index aa697d4bf..14832573d 100644 --- a/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/group_quantize_mxfp8.cuh @@ -867,7 +867,7 @@ void group_quantize(const GroupedTensor *input, const GroupedTensor *activations NVTE_CHECK(dbias->data.dtype == input->dtype(), "DBias must have the same type as input_tensor."); - std::vector expected_shape_dbias_tensor = {num_tensors, last_logical_dim}; + Shape expected_shape_dbias_tensor = {num_tensors, last_logical_dim}; NVTE_CHECK(dbias->data.shape == expected_shape_dbias_tensor, "Wrong shape of DBias."); NVTE_CHECK(workspace != nullptr, "Workspace must be a tensor."); diff --git a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh index 1549a292d..5e71a30e8 100644 --- a/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh +++ b/transformer_engine/common/cast/mxfp8/quantize_mxfp8.cuh @@ -578,8 +578,7 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, constexpr bool CAST_DBIAS_ONLY = IS_DBIAS && (!IS_DACT) && (!IS_ACT); // Tensor dimensions - const size_t rows = input.flat_first_dim(); - const size_t cols = input.flat_last_dim(); + const auto [rows, cols] = input.flat_2d_dims(); // Tensor chunk handled by each CUDA block constexpr size_t CHUNK_DIM_Y = CAST_DBIAS_ONLY ? 128 : 64; @@ -622,7 +621,7 @@ void quantize(const Tensor &input, const Tensor *act_input, const Tensor *noop, if constexpr (IS_DBIAS) { NVTE_CHECK(dbias->data.dtype == input.dtype(), "DBias must have the same type as input."); - NVTE_CHECK(dbias->data.shape == std::vector{cols}, "Wrong shape of DBias."); + NVTE_CHECK(dbias->data.shape == Shape{cols}, "Wrong shape of DBias."); NVTE_CHECK(workspace != nullptr, "Workspace must be a tensor."); if (workspace->data.dptr == nullptr) { diff --git a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh index faf3c58ad..13bb01d50 100644 --- a/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/dequantize_nvfp4.cuh @@ -95,8 +95,7 @@ inline void dequantize(const Tensor &input, Tensor *output, cudaStream_t stream) const int e4m3_max = input.nvfp4_e4m3_max; constexpr int FP4_BLOCK_SIZE = 16; - const size_t N = input.flat_first_dim(); - const size_t M = input.flat_last_dim(); + const auto [N, M] = input.flat_2d_dims(); NVTE_CHECK(M % FP4_BLOCK_SIZE == 0, "Last dimension of FP4 tensors needs to be divisible by ", FP4_BLOCK_SIZE, ", but got ", input.data.shape, "."); diff --git a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh index a2f3dac15..91c6af26b 100644 --- a/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/group_quantize_transpose_nvfp4.cuh @@ -783,8 +783,7 @@ void group_quantize_transpose(const Tensor &input, const Tensor *noop, NVTE_CHECK(input.has_data(), "Cannot quantize tensor without rowwise data."); - const size_t rows = input.flat_first_dim(); - const size_t cols = input.flat_last_dim(); + const auto [rows, cols] = input.flat_2d_dims(); NVTE_CHECK(rows % 32 == 0, "Number of tensor rows must be a multiple of 32"); // 16B alignment for TMA @@ -835,7 +834,7 @@ void group_quantize_transpose(const Tensor &input, const Tensor *noop, Tensor &rng_state_te_tensor = *convertNVTETensor(rng_state_tensor); NVTE_CHECK(rng_state_te_tensor.dtype() == DType::kInt64, "RNG state should contain 2 64-bit values."); - NVTE_CHECK(rng_state_te_tensor.data.shape == std::vector{2}, + NVTE_CHECK(rng_state_te_tensor.data.shape == Shape{2}, "Shape of the RNG state should be [2], but got ", rng_state_te_tensor.data.shape); rng_state = reinterpret_cast(rng_state_te_tensor.data.dptr); } diff --git a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh index 9e4aef5a1..e5100ec86 100644 --- a/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh +++ b/transformer_engine/common/cast/nvfp4/quantize_transpose_nvfp4.cuh @@ -121,8 +121,7 @@ inline void compute_rowwise_amax(const Tensor &input, const Tensor *noop, Tensor #if FP4_TYPE_SUPPORTED using namespace rowwise_amax_kernel; - const size_t rows = input.flat_first_dim(); - const size_t cols = input.flat_last_dim(); + const auto [rows, cols] = input.flat_2d_dims(); NVTE_CHECK(cols % ROWWISE_AMAX_SF_VEC_SIZE == 0, "Row-scaled NVFP4 quantization requires last dim divisible by ", ROWWISE_AMAX_SF_VEC_SIZE, "."); @@ -1359,8 +1358,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, "Transposed scaling tensor must be allocated"); } - const size_t rows = input.flat_first_dim(); - const size_t cols = input.flat_last_dim(); + const auto [rows, cols] = input.flat_2d_dims(); NVTE_CHECK(rows % 32 == 0, "Number of tensor rows must be a multiple of 32"); // 16B alignment for TMA @@ -1391,7 +1389,7 @@ void quantize_transpose(const Tensor &input, const Tensor *noop, Tensor *output, Tensor &rng_state_te_tensor = *convertNVTETensor(rng_state_tensor); NVTE_CHECK(rng_state_te_tensor.dtype() == DType::kInt64, "RNG state should contain 2 64-bit values."); - NVTE_CHECK(rng_state_te_tensor.data.shape == std::vector{2}, + NVTE_CHECK(rng_state_te_tensor.data.shape == Shape{2}, "Shape of the RNG state should be [2], but got ", rng_state_te_tensor.data.shape); rng_state = reinterpret_cast(rng_state_te_tensor.data.dptr); } diff --git a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh index 8adda8213..2cf43b5b6 100644 --- a/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh +++ b/transformer_engine/common/cast/nvfp4/specialized/quantize_transpose_nvfp4_tuned_1D.cuh @@ -718,8 +718,7 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, "Transposed scaling tensor must be allocated"); } - const size_t rows = input.flat_first_dim(); - const size_t cols = input.flat_last_dim(); + const auto [rows, cols] = input.flat_2d_dims(); NVTE_CHECK(rows % 32 == 0, "Number of tensor rows must be a multiple of 32"); // 16B alignment for TMA @@ -750,7 +749,7 @@ inline void quantize_transpose_tuned_1D(const Tensor &input, const Tensor *noop, Tensor &rng_state_te_tensor = *convertNVTETensor(rng_state_tensor); NVTE_CHECK(rng_state_te_tensor.dtype() == DType::kInt64, "RNG state should contain 2 64-bit values."); - NVTE_CHECK(rng_state_te_tensor.data.shape == std::vector{2}, + NVTE_CHECK(rng_state_te_tensor.data.shape == Shape{2}, "Shape of the RNG state should be [2], but got ", rng_state_te_tensor.data.shape); rng_state = reinterpret_cast(rng_state_te_tensor.data.dptr); } diff --git a/transformer_engine/common/comm_gemm/comm_gemm.cpp b/transformer_engine/common/comm_gemm/comm_gemm.cpp index ce389c200..60632b99d 100644 --- a/transformer_engine/common/comm_gemm/comm_gemm.cpp +++ b/transformer_engine/common/comm_gemm/comm_gemm.cpp @@ -130,12 +130,9 @@ int64_t block_size(NVTECommGemmCtx* ctx, int64_t global_size) { void AgGemmInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n, int64_t k, const Tensor* a, const Tensor* b, const Tensor* d, bool transa, bool transb) { - const auto a0 = a->flat_first_dim(); - const auto a1 = a->flat_last_dim(); - const auto b0 = b->flat_first_dim(); - const auto b1 = b->flat_last_dim(); - const auto d0 = d->flat_first_dim(); - const auto d1 = d->flat_last_dim(); + const auto [a0, a1] = a->flat_2d_dims(); + const auto [b0, b1] = b->flat_2d_dims(); + const auto [d0, d1] = d->flat_2d_dims(); if (transa) { NVTE_CHECK(a1 == k, "Unsupported tensor dimension in A: expected ", k, ", got ", a1); @@ -169,12 +166,9 @@ void AgGemmInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n void GemmRsInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n, int64_t k, const Tensor* a, const Tensor* b, const Tensor* d, bool transa, bool transb) { - const auto a0 = a->flat_first_dim(); - const auto a1 = a->flat_last_dim(); - const auto b0 = b->flat_first_dim(); - const auto b1 = b->flat_last_dim(); - const auto d0 = d->flat_first_dim(); - const auto d1 = d->flat_last_dim(); + const auto [a0, a1] = a->flat_2d_dims(); + const auto [b0, b1] = b->flat_2d_dims(); + const auto [d0, d1] = d->flat_2d_dims(); if (transa) { NVTE_CHECK(a0 == m, "Unsupported tensor dimension in A: expected ", m, ", got ", a0); @@ -213,12 +207,9 @@ void GemmRsInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n void GemmArInitMatrices(NVTECommGemmCtx* ctx, int64_t* ldd, int64_t m, int64_t n, int64_t k, const Tensor* a, const Tensor* b, const Tensor* d, bool transa, bool transb) { - const auto a0 = a->flat_first_dim(); - const auto a1 = a->flat_last_dim(); - const auto b0 = b->flat_first_dim(); - const auto b1 = b->flat_last_dim(); - const auto d0 = d->flat_first_dim(); - const auto d1 = d->flat_last_dim(); + const auto [a0, a1] = a->flat_2d_dims(); + const auto [b0, b1] = b->flat_2d_dims(); + const auto [d0, d1] = d->flat_2d_dims(); if (transa) { NVTE_CHECK(a0 == m, "Unsupported tensor dimension in A: expected ", m, ", got ", a0); diff --git a/transformer_engine/common/common.h b/transformer_engine/common/common.h index 6aa8798b7..eb4dcc055 100644 --- a/transformer_engine/common/common.h +++ b/transformer_engine/common/common.h @@ -16,6 +16,7 @@ static_assert(NVTE_BUILD_NUM_PHILOX_ROUNDS > 0, "NVTE_BUILD_NUM_PHILOX_ROUNDS must be a positive integer."); +#include #include #include #include @@ -26,13 +27,16 @@ static_assert(NVTE_BUILD_NUM_PHILOX_ROUNDS > 0, #include #include +#include +#include #include -#include -#include +#include +#include #include +#include #include #include -#include +#include #include #include "./nvtx.h" @@ -105,7 +109,9 @@ inline size_t product(const std::vector &shape, const size_t begin, cons return ret; } -inline size_t product(const std::vector &shape) { +template ::value>> +inline size_t product(const Container &shape) { size_t ret = 1; for (const auto &elem : shape) { ret *= elem; @@ -117,24 +123,138 @@ size_t get_buffer_size_bytes(const size_t N, const DType buffer_dtype); size_t get_buffer_size_bytes(const size_t dim_first, const size_t dim_last, const DType buffer_dtype); +/*! \brief Tensor shape + * + * Wraps NVTEShape with an interface similar to std::vector. + */ +class Shape { + public: + using value_type = size_t; + using size_type = size_t; + using iterator = size_t *; + using const_iterator = const size_t *; + + /*! Maximum number of dimensions this shape can hold. */ + static constexpr size_type max_ndim = std::extent_v; + + constexpr Shape() noexcept = default; + + explicit constexpr Shape(const NVTEShape &shape) noexcept : data_{shape} {} + + Shape(std::initializer_list shape) { + NVTE_CHECK(shape.size() <= max_ndim, "Too many dimensions (requested ", shape.size(), + ", max is ", max_ndim, ")."); + data_.ndim = shape.size(); + std::copy(shape.begin(), shape.end(), data_.data); + } + + // Construct from any container of integers + template ::value>> + explicit Shape(const Container &shape) { + NVTE_CHECK(shape.size() <= max_ndim, "Too many dimensions (requested ", shape.size(), + ", max is ", max_ndim, ")."); + data_.ndim = shape.size(); + std::copy(shape.begin(), shape.end(), data_.data); + } + + constexpr operator NVTEShape() const noexcept { return data_; } + + /*! Cast to std::vector */ + explicit operator std::vector() const { + return std::vector(data_.data, data_.data + data_.ndim); + } + + constexpr size_type size() const noexcept { return data_.ndim; } + constexpr bool empty() const noexcept { return data_.ndim == 0; } + static constexpr size_type capacity() noexcept { return max_ndim; } + + value_type *data() noexcept { return data_.data; } + constexpr const value_type *data() const noexcept { return data_.data; } + + iterator begin() noexcept { return data_.data; } + constexpr const_iterator begin() const noexcept { return data_.data; } + constexpr const_iterator cbegin() const noexcept { return data_.data; } + iterator end() noexcept { return data_.data + data_.ndim; } + constexpr const_iterator end() const noexcept { return data_.data + data_.ndim; } + constexpr const_iterator cend() const noexcept { return data_.data + data_.ndim; } + + const value_type &at(size_type i) const { + NVTE_CHECK(i < data_.ndim, "Attempted to access out-of-bounds entry (requested ", i, + ", size is ", data_.ndim, ")."); + return data_.data[i]; + } + value_type &at(size_type i) { return const_cast(std::as_const(*this).at(i)); } + + value_type &operator[](size_type i) noexcept { return data_.data[i]; } + constexpr const value_type &operator[](size_type i) const noexcept { return data_.data[i]; } + + value_type &front() noexcept { return data_.data[0]; } + constexpr const value_type &front() const noexcept { return data_.data[0]; } + + value_type &back() noexcept { return data_.data[data_.ndim - 1]; } + constexpr const value_type &back() const noexcept { return data_.data[data_.ndim - 1]; } + + void push_back(size_type value) { + NVTE_CHECK(data_.ndim < max_ndim, "Cannot add dimension: shape is at maximum capacity (", + max_ndim, ")."); + data_.data[data_.ndim++] = value; + } + + void resize(size_type count) { + NVTE_CHECK(count <= max_ndim, "Too many dimensions (requested ", count, ", max is ", max_ndim, + ")."); + if (count > data_.ndim) { + std::fill(&data_.data[data_.ndim], &data_.data[count], 0); + } + data_.ndim = count; + } + + void clear() noexcept { data_.ndim = 0; } + + friend bool operator==(const Shape &lhs, const Shape &rhs) noexcept { + return lhs.data_.ndim == rhs.data_.ndim && + std::equal(lhs.data_.data, lhs.data_.data + lhs.data_.ndim, rhs.data_.data); + } + friend bool operator!=(const Shape &lhs, const Shape &rhs) noexcept { return !(lhs == rhs); } + + template ::value>> + friend bool operator==(const Shape &lhs, const Container &rhs) { + return lhs == Shape(rhs); + } + template + friend bool operator==(const T &lhs, const Shape &rhs) { + return rhs == lhs; + } + template + friend bool operator!=(const Shape &lhs, const T &rhs) { + return !(lhs == rhs); + } + template + friend bool operator!=(const T &lhs, const Shape &rhs) { + return !(rhs == lhs); + } + + private: + NVTEShape data_{}; +}; + struct SimpleTensor { void *dptr; - std::vector shape; + Shape shape; DType dtype; SimpleTensor(void *dptr, std::vector shape, DType dtype) - : dptr{dptr}, shape{std::move(shape)}, dtype{dtype} {} + : dptr{dptr}, shape(shape), dtype{dtype} {} - SimpleTensor(const NVTEBasicTensor &tensor) // NOLINT - : dptr(tensor.data_ptr), - shape(tensor.shape.data, tensor.shape.data + tensor.shape.ndim), - dtype(static_cast(tensor.dtype)) {} + SimpleTensor() : SimpleTensor(nullptr, {0}, DType::kFloat32) {} - SimpleTensor() : SimpleTensor(nullptr, std::vector{0}, DType::kFloat32) {} + SimpleTensor(const NVTEBasicTensor &tensor) // NOLINT + : dptr(tensor.data_ptr), shape(tensor.shape), dtype(static_cast(tensor.dtype)) {} operator NVTEBasicTensor() const { - return {dptr, static_cast(dtype), - nvte_make_shape(this->shape.data(), this->shape.size())}; + return {dptr, static_cast(dtype), static_cast(shape)}; } /*! Number of tensor elements. */ @@ -223,12 +343,7 @@ struct Tensor { explicit operator NVTETensor() const noexcept { return nvte_tensor; } /*! Number of tensor elements. */ - size_t numel() const { - if (!has_data() && has_columnwise_data()) { - return product(columnwise_data.shape); - } - return product(data.shape); - } + size_t numel() const { return product(shape()); } /*! Whether the tensor data buffer is not uninitialized. * @@ -268,7 +383,7 @@ struct Tensor { * different shape, e.g. the column-wise data for some tensor * formats are transposed. */ - std::vector shape() const { + Shape shape() const { // Each tensor format interprets its data differently switch (scaling_mode) { case NVTE_DELAYED_TENSOR_SCALING: @@ -278,9 +393,8 @@ struct Tensor { // Row-wise data shape matches tensor logical shape, // column-wise data shape is transpose of logical shape if (!has_data() && has_columnwise_data()) { - std::vector ret; + Shape ret; if (!columnwise_data.shape.empty()) { - ret.reserve(columnwise_data.shape.size()); for (size_t i = 1; i < columnwise_data.shape.size(); i++) { ret.push_back(columnwise_data.shape[i]); } @@ -303,35 +417,36 @@ struct Tensor { } } - /*! Matrix height after tensor is flattened to 2D + /*! Matrix dimensions after flattening tensor to 2D. * * If a tensor has dimensions (D1, D2, ..., Dn), it is reinterpreted * as a (D1*D2*...*D(n-1), Dn) matrix. */ - size_t flat_first_dim() const { - const auto &full_shape = shape(); - size_t ret = 1; - if (!full_shape.empty()) { - for (size_t i = 0; i < full_shape.size() - 1; i++) { - ret *= full_shape[i]; - } + std::array flat_2d_dims() const { + const auto s = shape(); + if (s.empty()) { + return {1, 1}; + } + size_t first_dim = 1; + for (size_t i = 0; i + 1 < s.size(); ++i) { + first_dim *= s[i]; } - return ret; + return {first_dim, s.back()}; } + /*! Matrix height after tensor is flattened to 2D + * + * If a tensor has dimensions (D1, D2, ..., Dn), it is reinterpreted + * as a (D1*D2*...*D(n-1), Dn) matrix. + */ + size_t flat_first_dim() const { return flat_2d_dims()[0]; } + /*! Matrix width after tensor is flattened to 2D * * If a tensor has dimensions (D1, D2, ..., Dn), it is reinterpreted * as a (D1*D2*...*D(n-1), Dn) matrix. */ - size_t flat_last_dim() const { - const auto &full_shape = shape(); - if (full_shape.empty()) { - return 1; - } else { - return full_shape.back(); - } - } + size_t flat_last_dim() const { return flat_2d_dims()[1]; } }; struct GroupedTensor { @@ -1045,9 +1160,9 @@ inline bool is_aligned_tensor_data(const Tensor &t, size_t alignment) { size_t typeToSize(const DType type); size_t typeToNumBits(const DType type); -void CheckNoopTensor(const Tensor &t, const std::string &name); -void CheckInputTensor(const Tensor &t, const std::string &name, bool check_scale_inv_shapes = true); -void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empty = false); +void CheckNoopTensor(const Tensor &t, std::string_view name); +void CheckInputTensor(const Tensor &t, std::string_view name, bool check_scale_inv_shapes = true); +void CheckOutputTensor(const Tensor &t, std::string_view name, bool allow_empty = false); /*! \brief Update a tensor's FP8 scale-inverse * @@ -1082,9 +1197,9 @@ GroupedTensor *convertNVTEGroupedTensor(const NVTEGroupedTensor tensor); GroupedTensor *convertNVTEGroupedTensorCheck(const NVTEGroupedTensor tensor); // Helper functions for GroupedTensor validation -void CheckGroupedTensorShapeArrays(const GroupedTensor &t, const std::string &name); -void CheckInputGroupedTensor(const GroupedTensor &t, const std::string &name); -void CheckOutputGroupedTensor(const GroupedTensor &t, const std::string &name, +void CheckGroupedTensorShapeArrays(const GroupedTensor &t, std::string_view name); +void CheckInputGroupedTensor(const GroupedTensor &t, std::string_view name); +void CheckOutputGroupedTensor(const GroupedTensor &t, std::string_view name, bool allow_empty = false); } // namespace transformer_engine diff --git a/transformer_engine/common/gemm/cublaslt_gemm.cu b/transformer_engine/common/gemm/cublaslt_gemm.cu index 8589d7045..e59e9c00c 100644 --- a/transformer_engine/common/gemm/cublaslt_gemm.cu +++ b/transformer_engine/common/gemm/cublaslt_gemm.cu @@ -322,10 +322,8 @@ void cublas_gemm(const Tensor *inputA, const Tensor *inputB, Tensor *outputD, "cuBLAS GEMM does not support row-scaled NVFP4 inputs."); // Tensor dims in row-major order - const int A0 = inputA->flat_first_dim(); - const int A1 = inputA->flat_last_dim(); - const int B0 = inputB->flat_first_dim(); - const int B1 = inputB->flat_last_dim(); + const auto [A0, A1] = inputA->flat_2d_dims(); + const auto [B0, B1] = inputB->flat_2d_dims(); // GEMM dims in column-major order const int m = transa == CUBLAS_OP_T ? A0 : A1; diff --git a/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu index 0c3a5e929..d5dbd0bd8 100644 --- a/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/graph_safe_group_row_cast_col_hadamard_transform_cast_fusion.cu @@ -46,8 +46,8 @@ namespace { using namespace cute; -// Ensure Tensor refers to cute::Tensor, not transformer_engine::Tensor -using cute::Tensor; +using cute::Shape; // Avoid conflict with transformer_engine::Shape +using cute::Tensor; // Avoid conflict with transformer_engine::Tensor constexpr int kMaxTensorsPerKernel = 64; constexpr int kNVFP4BlockSize = 16; @@ -1343,7 +1343,7 @@ void group_hadamard_transform_cast_fusion_graph_safe(const GroupedTensor *input, const Tensor &rng_state_tensor = *convertNVTETensorCheck(quant_config.rng_state); NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64, "RNG state should contain 2 64-bit values."); - NVTE_CHECK(rng_state_tensor.data.shape == std::vector{2}, + NVTE_CHECK(rng_state_tensor.data.shape == Shape{2}, "Shape of the RNG state should be [2], but got ", rng_state_tensor.data.shape); rng_state = reinterpret_cast(rng_state_tensor.data.dptr); } @@ -1371,11 +1371,9 @@ void group_hadamard_transform_cast_fusion_graph_safe(const GroupedTensor *input, "Hadamard matrix must be BF16 tensor, but dtype is ", to_string(hadamard_matrix_.dtype()), "."); const SimpleTensor &hadamard_matrix = hadamard_matrix_.data; - NVTE_CHECK( - (hadamard_matrix_.shape() == std::vector{kHadamardDimension, kHadamardDimension}), - "Hadamard matrix must have shape=", - std::vector{kHadamardDimension, kHadamardDimension}, - ", but got shape=", hadamard_matrix_.shape(), "."); + NVTE_CHECK((hadamard_matrix_.shape() == Shape{kHadamardDimension, kHadamardDimension}), + "Hadamard matrix must have shape=", Shape{kHadamardDimension, kHadamardDimension}, + ", but got shape=", hadamard_matrix_.shape(), "."); const size_t hadamard_dimension = hadamard_matrix.shape[0]; const size_t num_tensors = input->num_tensors; diff --git a/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu index e6de366f5..e2325dd0f 100644 --- a/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/group_hadamard_transform_cast_fusion.cu @@ -36,8 +36,9 @@ namespace detail { namespace { using namespace cute; -using cute:: - Tensor; // Ensure unqualified Tensor refers to cute::Tensor, not transformer_engine::Tensor + +using cute::Shape; // Avoid conflict with transformer_engine::Shape +using cute::Tensor; // Avoid conflict with transformer_engine::Tensor using Stride2D = cute::Stride>; @@ -891,7 +892,7 @@ void group_hadamard_transform_cast_fusion_columnwise( Tensor &rng_state_tensor = *convertNVTETensor(quant_config.rng_state); NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64, "RNG state should contain 2 64-bit values."); - NVTE_CHECK(rng_state_tensor.data.shape == std::vector{2}, + NVTE_CHECK(rng_state_tensor.data.shape == Shape{2}, "Shape of the RNG state should be [2], but got ", rng_state_tensor.data.shape); rng_state = reinterpret_cast(rng_state_tensor.data.dptr); } @@ -911,11 +912,9 @@ void group_hadamard_transform_cast_fusion_columnwise( "Hadamard matrix must be BF16 tensor, but dtype is ", to_string(hadamard_matrix_.dtype()), "."); const SimpleTensor &hadamard_matrix = hadamard_matrix_.data; - NVTE_CHECK( - (hadamard_matrix_.shape() == std::vector{kHadamardDimension, kHadamardDimension}), - "Hadamard matrix must have shape=", - std::vector{kHadamardDimension, kHadamardDimension}, - ", but got shape=", hadamard_matrix_.shape(), "."); + NVTE_CHECK((hadamard_matrix_.shape() == Shape{kHadamardDimension, kHadamardDimension}), + "Hadamard matrix must have shape=", Shape{kHadamardDimension, kHadamardDimension}, + ", but got shape=", hadamard_matrix_.shape(), "."); const size_t hadamard_dimension = hadamard_matrix.shape[0]; const size_t ndim = input.shape.size(); diff --git a/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu index 1265f2711..359c41f1e 100644 --- a/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/group_row_cast_col_hadamard_transform_cast_fusion.cu @@ -46,8 +46,8 @@ namespace { using namespace cute; -// Ensure Tensor refers to cute::Tensor, not transformer_engine::Tensor -using cute::Tensor; +using cute::Shape; // Avoid conflict with transformer_engine::Shape +using cute::Tensor; // Avoid conflict with transformer_engine::Tensor constexpr int kMaxTensorsPerKernel = 64; @@ -1373,7 +1373,7 @@ void group_hadamard_transform_cast_fusion(const Tensor &input_, std::vector{2}, + NVTE_CHECK(rng_state_tensor.data.shape == Shape{2}, "Shape of the RNG state should be [2], but got ", rng_state_tensor.data.shape); rng_state = reinterpret_cast(rng_state_tensor.data.dptr); } @@ -1401,11 +1401,9 @@ void group_hadamard_transform_cast_fusion(const Tensor &input_, std::vector{kHadamardDimension, kHadamardDimension}), - "Hadamard matrix must have shape=", - std::vector{kHadamardDimension, kHadamardDimension}, - ", but got shape=", hadamard_matrix_.shape(), "."); + NVTE_CHECK((hadamard_matrix_.shape() == Shape{kHadamardDimension, kHadamardDimension}), + "Hadamard matrix must have shape=", Shape{kHadamardDimension, kHadamardDimension}, + ", but got shape=", hadamard_matrix_.shape(), "."); const size_t hadamard_dimension = hadamard_matrix.shape[0]; const size_t ndim = input.shape.size(); diff --git a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu index 957935668..50b9f63bd 100644 --- a/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/hadamard_transform_cast_fusion.cu @@ -38,7 +38,9 @@ namespace detail { namespace { using namespace cute; -using cute::Tensor; // Ensure unqualified Tensor refers to cute::Tensor, not transformer_engine::Tensor + +using cute::Tensor; // Avoid conflict with transformer_engine::Tensor +using cute::Shape; // Avoid conflict with transformer_engine::Shape // calculate the global encode scale factor for a given global amax. __device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax) { @@ -749,7 +751,7 @@ void hadamard_transform_cast_fusion_columnwise(const Tensor &input_, Tensor &out Tensor &rng_state_tensor = *convertNVTETensor(quant_config.rng_state); NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64, "RNG state should contain 2 64-bit values."); - NVTE_CHECK(rng_state_tensor.data.shape == std::vector{2}, + NVTE_CHECK(rng_state_tensor.data.shape == Shape{2}, "Shape of the RNG state should be [2], but got ", rng_state_tensor.data.shape); rng_state = reinterpret_cast(rng_state_tensor.data.dptr); } @@ -771,11 +773,9 @@ void hadamard_transform_cast_fusion_columnwise(const Tensor &input_, Tensor &out "Hadamard matrix must be BF16 tensor, but dtype is ", to_string(hadamard_matrix_.dtype()), "."); const SimpleTensor &hadamard_matrix = hadamard_matrix_.data; - NVTE_CHECK( - (hadamard_matrix_.shape() == std::vector{kHadamardDimension, kHadamardDimension}), - "Hadamard matrix must have shape=", - std::vector{kHadamardDimension, kHadamardDimension}, - ", but got shape=", hadamard_matrix_.shape(), "."); + NVTE_CHECK((hadamard_matrix_.shape() == Shape{kHadamardDimension, kHadamardDimension}), + "Hadamard matrix must have shape=", Shape{kHadamardDimension, kHadamardDimension}, + ", but got shape=", hadamard_matrix_.shape(), "."); const size_t hadamard_dimension = hadamard_matrix.shape[0]; const size_t ndim = input.shape.size(); diff --git a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu index 99060ab62..3f2c08a05 100644 --- a/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu +++ b/transformer_engine/common/hadamard_transform/row_cast_col_hadamard_transform_cast_fusion.cu @@ -48,6 +48,9 @@ namespace { using namespace cute; +using cute::Tensor; // Avoid conflict with transformer_engine::Tensor +using cute::Shape; // Avoid conflict with transformer_engine::Shape + struct CLCResponse { uint32_t data[4] = {0}; }; constexpr int kFp4ConvertChunkElements = 8; @@ -1269,7 +1272,7 @@ void hadamard_transform_cast_fusion(const Tensor &input_, Tensor &output_, Tensor &rng_state_tensor = *convertNVTETensor(quant_config.rng_state); NVTE_CHECK(rng_state_tensor.dtype() == DType::kInt64, "RNG state should contain 2 64-bit values."); - NVTE_CHECK(rng_state_tensor.data.shape == std::vector{2}, + NVTE_CHECK(rng_state_tensor.data.shape == Shape{2}, "Shape of the RNG state should be [2], but got ", rng_state_tensor.data.shape); rng_state = reinterpret_cast(rng_state_tensor.data.dptr); } @@ -1293,11 +1296,9 @@ void hadamard_transform_cast_fusion(const Tensor &input_, Tensor &output_, "Hadamard matrix must be BF16 tensor, but dtype is ", to_string(hadamard_matrix_.dtype()), "."); const SimpleTensor &hadamard_matrix = hadamard_matrix_.data; - NVTE_CHECK( - (hadamard_matrix_.shape() == std::vector{kHadamardDimension, kHadamardDimension}), - "Hadamard matrix must have shape=", - std::vector{kHadamardDimension, kHadamardDimension}, - ", but got shape=", hadamard_matrix_.shape(), "."); + NVTE_CHECK((hadamard_matrix_.shape() == Shape{kHadamardDimension, kHadamardDimension}), + "Hadamard matrix must have shape=", Shape{kHadamardDimension, kHadamardDimension}, + ", but got shape=", hadamard_matrix_.shape(), "."); const size_t hadamard_dimension = hadamard_matrix.shape[0]; const size_t ndim = input.shape.size(); diff --git a/transformer_engine/common/include/transformer_engine/transformer_engine.h b/transformer_engine/common/include/transformer_engine/transformer_engine.h index ffb324315..f675b2f53 100644 --- a/transformer_engine/common/include/transformer_engine/transformer_engine.h +++ b/transformer_engine/common/include/transformer_engine/transformer_engine.h @@ -147,6 +147,20 @@ typedef void *NVTETensor; */ NVTETensor nvte_create_tensor(NVTEScalingMode scaling_mode); +/*! \brief Create a batch of new TE tensors. + * + * Equivalent to calling nvte_create_tensor N times with the same + * scaling mode. Before use, each tensor's parameters need to be set. + * TE tensors are just wrappers on top of raw data and do not own + * memory. + * + * \param[in] scaling_mode Scaling mode shared by all tensors. + * \param[out] tensors Caller-allocated array of length N to + * receive the new tensors. + * \param[in] N Number of tensors to create. + */ +void nvte_create_tensors(NVTEScalingMode scaling_mode, NVTETensor *tensors, size_t N); + /*! \brief Destroy a TE tensor. * * Since the TE tensor does not own memory, the underlying @@ -156,6 +170,17 @@ NVTETensor nvte_create_tensor(NVTEScalingMode scaling_mode); */ void nvte_destroy_tensor(NVTETensor tensor); +/*! \brief Destroy a batch of TE tensors. + * + * Equivalent to calling nvte_destroy_tensor N times. Since TE tensors + * do not own memory, the underlying data is not freed during this + * operation. Null entries are ignored. + * + * \param[in] tensors Array of tensors to be destroyed. + * \param[in] N Number of tensors in the array. + */ +void nvte_destroy_tensors(NVTETensor *tensors, size_t N); + /*! \brief Get a raw pointer to the tensor's rowwise data. * * \param[in] tensor Tensor. @@ -602,6 +627,7 @@ NVTEShape nvte_get_grouped_tensor_logical_shape(const NVTEGroupedTensor tensor); #ifdef __cplusplus } // extern "C" +#include #include /*! \namespace transformer_engine @@ -1043,6 +1069,73 @@ class TensorWrapper { NVTETensor tensor_ = nullptr; }; +/*! \struct MultiTensorWrapper + * \brief C++ wrapper for a batch of NVTETensors allocated together. + */ +class MultiTensorWrapper { + public: + /*! \brief Constructs an empty batch. */ + MultiTensorWrapper() = default; + + /*! \brief Allocates a batch of NVTETensors. + * + * \param[in] num_tensors Number of tensors to allocate. + * \param[in] scaling_mode Scaling mode shared by all tensors. + */ + explicit MultiTensorWrapper(size_t num_tensors, + NVTEScalingMode scaling_mode = NVTE_DELAYED_TENSOR_SCALING) + : tensors_(num_tensors) { + if (!tensors_.empty()) { + nvte_create_tensors(scaling_mode, tensors_.data(), tensors_.size()); + } + } + + ~MultiTensorWrapper() { + if (!tensors_.empty()) { + nvte_destroy_tensors(tensors_.data(), tensors_.size()); + } + } + + MultiTensorWrapper(const MultiTensorWrapper &) = delete; + MultiTensorWrapper &operator=(const MultiTensorWrapper &) = delete; + + MultiTensorWrapper(MultiTensorWrapper &&) noexcept = default; + + MultiTensorWrapper &operator=(MultiTensorWrapper &&other) noexcept { + if (this == &other) return *this; + if (!tensors_.empty()) { + nvte_destroy_tensors(tensors_.data(), tensors_.size()); + } + tensors_ = std::move(other.tensors_); + return *this; + } + + /*! \brief Number of tensors in the batch. */ + size_t size() const noexcept { return tensors_.size(); } + + /*! \brief Whether the batch is empty. */ + bool empty() const noexcept { return tensors_.empty(); } + + /*! \brief Access an NVTETensor by index. */ + NVTETensor operator[](size_t i) const noexcept { return tensors_[i]; } + + /*! \brief Pointer to the underlying NVTETensor array. */ + NVTETensor *data() noexcept { return tensors_.data(); } + const NVTETensor *data() const noexcept { return tensors_.data(); } + + /*! \brief Implicit conversion for multi-tensor C API calls. */ + operator NVTETensor *() noexcept { return tensors_.data(); } + + /*! \brief Iteration over the underlying NVTETensors. */ + auto begin() noexcept { return tensors_.begin(); } + auto end() noexcept { return tensors_.end(); } + auto begin() const noexcept { return tensors_.begin(); } + auto end() const noexcept { return tensors_.end(); } + + private: + std::vector tensors_; +}; + /*! \struct GroupedTensorWrapper * \brief C++ wrapper for the NVTEGroupedTensor class. */ diff --git a/transformer_engine/common/include/transformer_engine/utils.h b/transformer_engine/common/include/transformer_engine/utils.h index eca6f359e..fda49dd54 100644 --- a/transformer_engine/common/include/transformer_engine/utils.h +++ b/transformer_engine/common/include/transformer_engine/utils.h @@ -5,13 +5,14 @@ ************************************************************************/ /*! \file utils.h - * \brief Utility functions (e.g. host-to-device pointer copies). + * \brief Utility functions (e.g. host-to-device value stores). */ #ifndef TRANSFORMER_ENGINE_UTILS_H_ #define TRANSFORMER_ENGINE_UTILS_H_ #include +#include #include #include @@ -19,12 +20,22 @@ extern "C" { #endif -/*! \brief Copy an array of device pointers (held on host) into a device tensor. +/*! \brief Copy a small host buffer into device memory via kernel arguments. * - * \param[in] host_ptrs Host array of device pointer values cast to uint64_t. - * \param[out] output NVTETensor whose rowwise data buffer receives the pointer values. - * \param[in] count Number of pointers. - * \param[in] stream CUDA stream used for the operation. + * The host buffer may be modified or freed after this call returns. + * This is compatible with CUDA Graphs. + * + * \param[in] host_ptr Source in host memory. + * \param[out] device_ptr Destination in device memory. + * \param[in] num_bytes Size of the value in bytes. + * \param[in] stream CUDA stream for the operation. + */ +void nvte_copy_host_to_device_via_kernel(const void *host_ptr, void *device_ptr, size_t num_bytes, + cudaStream_t stream); + +/*! \deprecated Use nvte_copy_host_to_device_via_kernel instead. + * + * \brief Copy an array of device pointers (held on host) into a device tensor. */ void nvte_convert_pointers_to_tensor(const uint64_t *host_ptrs, NVTETensor output, int64_t count, cudaStream_t stream); diff --git a/transformer_engine/common/normalization/common.cpp b/transformer_engine/common/normalization/common.cpp index 7dd942b31..375b109c2 100644 --- a/transformer_engine/common/normalization/common.cpp +++ b/transformer_engine/common/normalization/common.cpp @@ -128,7 +128,7 @@ void TeNormalizationPlan::_build() { } template -std::vector TeNormalizationPlan::getWorkspaceShape() const { +Shape TeNormalizationPlan::getWorkspaceShape() const { size_t workspace_size = _launch_params.getTotalWorkspaceBytes(_is_layernorm); if (workspace_size == 0) { // Workspace size must not be zero since that corresponds to a @@ -431,7 +431,7 @@ void CudnnNormalizationPlan::_build() { _graph.build_plans(_handle, cudnn_frontend::BuildPlanPolicy_t::HEURISTICS_CHOICE).is_good()); } -std::vector CudnnNormalizationPlan::getWorkspaceShape() const { +Shape CudnnNormalizationPlan::getWorkspaceShape() const { size_t workspace_size = _graph.get_workspace_size(); if (workspace_size == 0) { // Workspace size must not be zero since that corresponds to a diff --git a/transformer_engine/common/normalization/common.h b/transformer_engine/common/normalization/common.h index 0cbd5a99f..f5dce6419 100644 --- a/transformer_engine/common/normalization/common.h +++ b/transformer_engine/common/normalization/common.h @@ -219,7 +219,7 @@ class TeNormalizationRegistry { class NormalizationPlanBase { public: virtual ~NormalizationPlanBase() = default; - virtual std::vector getWorkspaceShape() const = 0; + virtual Shape getWorkspaceShape() const = 0; virtual void execute(Tensor* z, void* x_dptr, void* gamma_dptr, void* beta_dptr, void* mean_dptr, void* eps_dptr, void* rsigma_dptr, void* workspace_dptr, @@ -239,7 +239,7 @@ class TeNormalizationPlan : public NormalizationPlanBase { TeNormalizationPlan(NVTE_Norm_Type NormType, NVTE_Norm_Stage NormStage, DType wtype, DType itype, DType otype, DType ctype, const size_t batch_size, const size_t hidden_size, const size_t sm_count, const bool zero_centered_gamma, const bool is_tuned); - std::vector getWorkspaceShape() const override; + Shape getWorkspaceShape() const override; void execute(Tensor* z, void* x_dptr, void* gamma_dptr, void* beta_dptr, void* mean_dptr, void* eps_dptr, void* rsigma_dptr, void* workspace_dptr, @@ -268,7 +268,7 @@ class CudnnNormalizationPlan : public NormalizationPlanBase { const bool zero_centered_gamma, const NVTEScalingMode mode, const bool training); - std::vector getWorkspaceShape() const override; + Shape getWorkspaceShape() const override; void execute(Tensor* z, void* x_dptr, void* gamma_dptr, void* beta_dptr, void* mean_dptr, void* eps_dptr, void* rsigma_dptr, void* workspace_dptr, diff --git a/transformer_engine/common/normalization/layernorm/ln_api.cpp b/transformer_engine/common/normalization/layernorm/ln_api.cpp index 7bd5a1bbd..0f843019c 100644 --- a/transformer_engine/common/normalization/layernorm/ln_api.cpp +++ b/transformer_engine/common/normalization/layernorm/ln_api.cpp @@ -11,7 +11,6 @@ #include #include #include -#include #include "../../common.h" #include "../common.h" @@ -48,11 +47,11 @@ void layernorm_fwd(const Tensor& x, // BxSxhidden_size NVTE_CHECK(z->data.shape == x.data.shape, "Output tensor must have the same shape as x."); - NVTE_CHECK(mu->data.shape == std::vector{x.data.shape[0]}, + NVTE_CHECK(mu->data.shape == Shape{x.data.shape[0]}, "Mu must be 1D tensor with shape (x.shape[0],)."); NVTE_CHECK(mu->data.dtype == DType::kFloat32, "Mu must be a float32 tensor."); - NVTE_CHECK(rsigma->data.shape == std::vector{x.data.shape[0]}, + NVTE_CHECK(rsigma->data.shape == Shape{x.data.shape[0]}, "RSigma must be 1D tensor with shape (x.shape[0],)."); NVTE_CHECK(rsigma->data.dtype == DType::kFloat32, "RSigma must be a float32 tensor."); diff --git a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp index adf2ccee0..07ad3230a 100644 --- a/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp +++ b/transformer_engine/common/normalization/rmsnorm/rmsnorm_api.cpp @@ -8,7 +8,6 @@ #include #include #include -#include #include "../../common.h" #include "../common.h" @@ -40,7 +39,7 @@ void rmsnorm_fwd(const Tensor &x, const Tensor &gamma, const float epsilon, Tens NVTE_CHECK(z->data.shape == x.data.shape, "Output tensor must have the same shape as x."); - NVTE_CHECK(rsigma->data.shape == std::vector{x.data.shape[0]}, + NVTE_CHECK(rsigma->data.shape == Shape{x.data.shape[0]}, "RSigma must be 1D tensor with shape (x.shape[0],)."); NVTE_CHECK(rsigma->data.dtype == DType::kFloat32, "RSigma must be a float32 tensor."); diff --git a/transformer_engine/common/swizzle/swizzle.cu b/transformer_engine/common/swizzle/swizzle.cu index c7ed407a5..51969e10e 100644 --- a/transformer_engine/common/swizzle/swizzle.cu +++ b/transformer_engine/common/swizzle/swizzle.cu @@ -24,15 +24,18 @@ namespace { constexpr int MXFP8_BLOCK_SIZE = 32; constexpr int NVFP4_BLOCK_SIZE = 16; -int get_max_dynamic_smem() { - static int max_smem = -1; - if (max_smem < 0) { - int device; - NVTE_CHECK_CUDA(cudaGetDevice(&device)); - NVTE_CHECK_CUDA( - cudaDeviceGetAttribute(&max_smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, device)); +int get_max_dynamic_smem(int device_id = -1) { + static std::vector cache(cuda::num_devices(), -1); + static std::vector flags(cuda::num_devices()); + if (device_id < 0) { + device_id = cuda::current_device(); } - return max_smem; + auto init = [&]() { + NVTE_CHECK_CUDA(cudaDeviceGetAttribute(&cache[device_id], + cudaDevAttrMaxSharedMemoryPerBlockOptin, device_id)); + }; + std::call_once(flags[device_id], init); + return cache[device_id]; } constexpr __device__ __host__ int TB_DIM = 32; @@ -1456,10 +1459,8 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, // We don't allow empty tensors. They should be filtered out before calling this function. NVTE_CHECK(input[i]->numel() != 0, "Tensor input[", i, "] is empty."); - CheckInputTensor(*input[i], "scaling_factor_input[" + std::to_string(i) + "]", - check_scale_inv_shapes); - CheckInputTensor(*output[i], "scaling_factor_output[" + std::to_string(i) + "]", - check_scale_inv_shapes); + CheckInputTensor(*input[i], "scaling_factor_input", check_scale_inv_shapes); + CheckInputTensor(*output[i], "scaling_factor_output", check_scale_inv_shapes); all_has_data = all_has_data && input[i]->scale_inv.has_data(); all_has_columnwise_data = (all_has_columnwise_data && input[i]->columnwise_scale_inv.has_data()); @@ -1540,17 +1541,18 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, const int pos = kernel_args.num_tensors; kernel_args.m_list[pos] = m; kernel_args.k_list[pos] = k; + const auto [first_dim, last_dim] = input[i]->flat_2d_dims(); if (!all_nvfp4 || all_has_data) { int block_scale_size = all_nvfp4 ? NVFP4_BLOCK_SIZE : MXFP8_BLOCK_SIZE; kernel_args.input_list[pos] = const_cast(input[i]->scale_inv.dptr); kernel_args.output_list[pos] = output[i]->scale_inv.dptr; - kernel_args.original_m_list[pos] = input[i]->flat_first_dim(); - kernel_args.original_k_list[pos] = input[i]->flat_last_dim() / block_scale_size; + kernel_args.original_m_list[pos] = first_dim; + kernel_args.original_k_list[pos] = last_dim / block_scale_size; } else { kernel_args.input_list[pos] = const_cast(input[i]->columnwise_scale_inv.dptr); kernel_args.output_list[pos] = output[i]->columnwise_scale_inv.dptr; - kernel_args.original_m_list[pos] = input[i]->flat_last_dim(); - kernel_args.original_k_list[pos] = input[i]->flat_first_dim() / NVFP4_BLOCK_SIZE; + kernel_args.original_m_list[pos] = last_dim; + kernel_args.original_k_list[pos] = first_dim / NVFP4_BLOCK_SIZE; } kernel_args.num_tensors++; } @@ -1609,8 +1611,9 @@ void multi_tensor_swizzle_scaling_factors(const std::vector& input, kernel_args.output_list[pos] = output[i]->columnwise_scale_inv.dptr; kernel_args.m_list[pos] = m; kernel_args.k_list[pos] = k; - kernel_args.original_m_list[pos] = input[i]->flat_last_dim(); - kernel_args.original_k_list[pos] = input[i]->flat_first_dim() / MXFP8_BLOCK_SIZE; + const auto [first_dim, last_dim] = input[i]->flat_2d_dims(); + kernel_args.original_m_list[pos] = last_dim; + kernel_args.original_k_list[pos] = first_dim / MXFP8_BLOCK_SIZE; kernel_args.num_tensors++; } // Launch the remaining tensors @@ -1958,6 +1961,8 @@ void nvte_multi_tensor_swizzle_scaling_factors(const NVTETensor* inputs, NVTETen using namespace transformer_engine; NVTE_CHECK(num_tensors > 0, "Number of tensors should be greater than 0."); std::vector input_list, output_list; + input_list.reserve(num_tensors); + output_list.reserve(num_tensors); for (size_t i = 0; i < num_tensors; i++) { input_list.push_back(convertNVTETensorCheck(inputs[i])); output_list.push_back(convertNVTETensorCheck(outputs[i])); diff --git a/transformer_engine/common/transformer_engine.cpp b/transformer_engine/common/transformer_engine.cpp index ca7eccab0..b3179d38f 100644 --- a/transformer_engine/common/transformer_engine.cpp +++ b/transformer_engine/common/transformer_engine.cpp @@ -7,6 +7,7 @@ #include #include +#include #include #include #include @@ -51,7 +52,7 @@ std::string to_string(const NVTEScalingMode &mode) { return "Invalid Scaling"; } -void CheckNoopTensor(const Tensor &t, const std::string &name) { +void CheckNoopTensor(const Tensor &t, std::string_view name) { if (t.data.has_data()) { NVTE_CHECK(t.numel() == 1, "Expected 1 element for ", name, " noop, but found ", t.numel(), "."); @@ -60,7 +61,7 @@ void CheckNoopTensor(const Tensor &t, const std::string &name) { } } -void CheckScaleTensorShape(const Tensor &t, const std::string &name) { +void CheckScaleTensorShape(const Tensor &t, std::string_view name) { NVTE_CHECK(t.scaling_mode != NVTE_INVALID_SCALING, "Invalid scaling mode!"); if (is_tensor_scaling(t.scaling_mode)) { if (is_fp8_dtype(t.dtype())) { @@ -91,60 +92,55 @@ void CheckScaleTensorShape(const Tensor &t, const std::string &name) { } else { if (t.scaling_mode == NVTE_MXFP8_1D_SCALING) { // Need (4, 128) alignment even for e8 scaling factor - auto block_alignment = std::vector{128ul, 4ul}; - size_t expected_x, expected_y, alignment; - const size_t block_size_rowwise = 32; - const size_t block_size_colwise = 32; + constexpr std::array block_alignment{128ul, 4ul}; + const auto [first_dim, last_dim] = t.flat_2d_dims(); if (t.has_data()) { - alignment = block_alignment[0]; - expected_x = - DIVUP(DIVUP(t.flat_first_dim(), static_cast(1)), alignment) * alignment; - alignment = block_alignment[1]; - expected_y = - DIVUP(DIVUP(t.flat_last_dim(), static_cast(block_size_rowwise)), alignment) * - alignment; - - const auto &expected = std::vector{expected_x, expected_y}; + constexpr std::array block_shape{1, 32}; + const std::array expected{ + DIVUP_TO_MULTIPLE(DIVUP(first_dim, block_shape[0]), block_alignment[0]), + DIVUP_TO_MULTIPLE(DIVUP(last_dim, block_shape[1]), block_alignment[1])}; NVTE_CHECK(t.scale_inv.shape == expected, "Tensor \"", name, "\" has invalid scale_inv shape (expected ", expected, ", got ", t.scale_inv.shape, ")"); } if (t.has_columnwise_data()) { - alignment = block_alignment[1]; - expected_x = - DIVUP(DIVUP(t.flat_first_dim(), static_cast(block_size_colwise)), alignment) * - alignment; - alignment = block_alignment[0]; - expected_y = DIVUP(DIVUP(t.flat_last_dim(), static_cast(1)), alignment) * alignment; - - const auto &expected = std::vector{expected_x, expected_y}; + constexpr std::array block_shape{32, 1}; + const std::array expected{ + DIVUP_TO_MULTIPLE(DIVUP(first_dim, block_shape[0]), block_alignment[1]), + DIVUP_TO_MULTIPLE(DIVUP(last_dim, block_shape[1]), block_alignment[0])}; NVTE_CHECK(t.columnwise_scale_inv.shape == expected, "Tensor \"", name, "\" has invalid columnwise_scale_inv shape (expected ", expected, ", got ", t.columnwise_scale_inv.shape, ")"); } } else if (t.scaling_mode == NVTE_NVFP4_1D_SCALING) { + const auto [first_dim, last_dim] = t.flat_2d_dims(); + if (t.has_data()) { - const size_t expected_y = DIVUP_TO_MULTIPLE(t.flat_first_dim(), 128); - const size_t expected_x = DIVUP_TO_MULTIPLE(DIVUP(t.flat_last_dim(), 16lu), 4); - const auto &expected = std::vector{expected_y, expected_x}; + constexpr std::array block_shape{1, 16}; + constexpr std::array block_alignment{128, 4}; + const std::array expected{ + DIVUP_TO_MULTIPLE(DIVUP(first_dim, block_shape[0]), block_alignment[0]), + DIVUP_TO_MULTIPLE(DIVUP(last_dim, block_shape[1]), block_alignment[1])}; NVTE_CHECK(t.scale_inv.shape == expected, "Tensor \"", name, "\" has invalid scale_inv shape (expected ", expected, ", got ", t.scale_inv.shape, ")"); } if (t.has_columnwise_data()) { - const size_t expected_y = DIVUP_TO_MULTIPLE(t.flat_last_dim(), 128); - const size_t expected_x = DIVUP_TO_MULTIPLE(DIVUP(t.flat_first_dim(), 16lu), 4); - const auto &expected = std::vector{expected_y, expected_x}; + constexpr std::array block_shape{1, 16}; + constexpr std::array block_alignment{128, 4}; + const std::array expected{ + DIVUP_TO_MULTIPLE(DIVUP(last_dim, block_shape[0]), block_alignment[0]), + DIVUP_TO_MULTIPLE(DIVUP(first_dim, block_shape[1]), block_alignment[1])}; NVTE_CHECK(t.columnwise_scale_inv.shape == expected, "Tensor \"", name, - "\" has invalid columnwise_scale_inv shape (expected ", expected, ", got ", + "\" has invalid columnwise_scale_inv shape (expected ", expected, ", got ", t.columnwise_scale_inv.shape, ")"); } } } } -void CheckInputTensor(const Tensor &t, const std::string &name, bool check_scale_inv_shapes) { +void CheckInputTensor(const Tensor &t, std::string_view name, bool check_scale_inv_shapes) { const DType type = t.dtype(); if (is_fp8_dtype(type)) { // FP8 input needs to have scale_inv @@ -200,7 +196,7 @@ void CheckInputTensor(const Tensor &t, const std::string &name, bool check_scale } } -void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empty) { +void CheckOutputTensor(const Tensor &t, std::string_view name, bool allow_empty) { const DType type = t.dtype(); if (is_fp8_dtype(type)) { // FP8 output needs to have scale, scale_inv and (if delayed scaling) amax @@ -262,7 +258,7 @@ void CheckOutputTensor(const Tensor &t, const std::string &name, bool allow_empt CheckScaleTensorShape(t, name); } -void CheckGroupedTensorShapeArrays(const GroupedTensor &t, const std::string &name) { +void CheckGroupedTensorShapeArrays(const GroupedTensor &t, std::string_view name) { NVTE_CHECK(t.num_tensors > 0, "Grouped tensor ", name, " has no tensors!"); // Helper lambda to validate shape arrays @@ -332,7 +328,7 @@ void CheckGroupedTensorShapeArrays(const GroupedTensor &t, const std::string &na } // Helper function to check scale_inv for both input and output -static void CheckGroupedScaleInv(const GroupedTensor &t, const std::string &name, bool is_output) { +static void CheckGroupedScaleInv(const GroupedTensor &t, std::string_view name, bool is_output) { const char *tensor_type = is_output ? "output" : "input"; // Helper to check scale_inv for both rowwise and columnwise layouts @@ -371,14 +367,14 @@ static void CheckGroupedScaleInv(const GroupedTensor &t, const std::string &name } } -void CheckInputGroupedTensor(const GroupedTensor &t, const std::string &name) { +void CheckInputGroupedTensor(const GroupedTensor &t, std::string_view name) { NVTE_CHECK(t.has_data() || t.has_columnwise_data(), "Input grouped tensor ", name, " not allocated"); CheckGroupedScaleInv(t, name, false); CheckGroupedTensorShapeArrays(t, name); } -void CheckOutputGroupedTensor(const GroupedTensor &t, const std::string &name, bool allow_empty) { +void CheckOutputGroupedTensor(const GroupedTensor &t, std::string_view name, bool allow_empty) { if (!allow_empty) { NVTE_CHECK(t.has_data() || t.has_columnwise_data(), "Output grouped tensor ", name, " not allocated"); @@ -406,51 +402,36 @@ class TensorAllocator { ~TensorAllocator() {} - NVTETensor Allocate(NVTEScalingMode mode) { + void Allocate(NVTEScalingMode mode, NVTETensor *out, size_t N) { std::lock_guard lock(mutex); - if (!free_list.empty()) { - uintptr_t index = free_list.back(); - NVTETensor ret = reinterpret_cast(index); - free_list.pop_back(); - if (debug) { - std::cout << "Allocated " << index - << " from free list. Free list size: " << free_list.size() << " and capacity " - << free_list.capacity() << std::endl; + const size_t available = free_list.size() + (memory.capacity() - memory.size()); + NVTE_CHECK(available >= N, "Cannot allocate ", N, + " new NVTETensors. Maximum number of tensors reached: ", MAX_TENSOR_NUM, + ". There is probably a memory leak in your application."); + for (size_t i = 0; i < N; ++i) { + uintptr_t index; + if (!free_list.empty()) { + index = free_list.back(); + free_list.pop_back(); + } else { + memory.emplace_back(); + index = memory.size(); + size = index; + memory[index - 1].nvte_tensor = reinterpret_cast(index); } - // 1-based indexing memory[index - 1].scaling_mode = mode; - return ret; + out[i] = reinterpret_cast(index); } - if (memory.size() < memory.capacity()) { - memory.emplace_back(); - Tensor &t = memory.back(); - size = memory.size(); - // 1-based indexing - uintptr_t index = memory.size(); - if (debug) { - std::cout << "Allocated " << index << ". Memory size: " << memory.size() << " and capacity " - << memory.capacity() << std::endl; - } - t.scaling_mode = mode; - t.nvte_tensor = reinterpret_cast(index); - return reinterpret_cast(index); + if (debug) { + std::cout << "Allocated range of " << N << " tensors. Free list size: " << free_list.size() + << " and capacity " << free_list.capacity() << std::endl; } - NVTE_ERROR("Cannot allocate a new NVTETensor. Maximum number of tensors reached: ", - MAX_TENSOR_NUM, ". There is probably a memory leak in your application."); } - void Free(NVTETensor t) { - uintptr_t index = reinterpret_cast(t); - if (index == 0) return; - std::lock_guard lock(mutex); - NVTE_CHECK(index <= memory.size(), "Invalid tensor."); - free_list.push_back(index); - // Clean up - memory[index - 1].clear(); - if (debug) { - std::cout << "Freed " << index << ". Free list size: " << free_list.size() << " and capacity " - << free_list.capacity() << std::endl; - } + NVTETensor Allocate(NVTEScalingMode mode) { + NVTETensor t; + Allocate(mode, &t, 1); + return t; } void Free(NVTETensor *t, size_t N) { @@ -464,11 +445,16 @@ class TensorAllocator { memory[index - 1].clear(); } if (debug) { - std::cout << "Freed range of" << N << " tensors. Free list size: " << free_list.size() + std::cout << "Freed range of " << N << " tensors. Free list size: " << free_list.size() << " and capacity " << free_list.capacity() << std::endl; } } + void Free(NVTETensor t) { + if (reinterpret_cast(t) == 0) return; + Free(&t, 1); + } + Tensor *convertNVTETensor(NVTETensor t) { uintptr_t index = reinterpret_cast(t); // 1-based indexing to enable 0-initialization of NVTETensor @@ -601,6 +587,10 @@ NVTETensor nvte_create_tensor(NVTEScalingMode scaling_mode) { return ret; } +void nvte_create_tensors(NVTEScalingMode scaling_mode, NVTETensor *tensors, size_t N) { + transformer_engine::TensorAllocator::instance().Allocate(scaling_mode, tensors, N); +} + void nvte_destroy_tensor(NVTETensor tensor) { transformer_engine::TensorAllocator::instance().Free(tensor); } @@ -638,11 +628,7 @@ NVTEShape nvte_tensor_shape(const NVTETensor tensor) { if (t == nullptr) { NVTE_ERROR("Invalid tensor: received null pointer in nvte_tensor_shape"); } - - // Determine tensor shape depending on tensor format - const std::vector &shape = t->shape(); - - return nvte_make_shape(shape.data(), shape.size()); + return t->shape(); } NVTEShape nvte_tensor_columnwise_shape(const NVTETensor tensor) { @@ -650,8 +636,7 @@ NVTEShape nvte_tensor_columnwise_shape(const NVTETensor tensor) { if (t == nullptr) { NVTE_ERROR("Invalid tensor: received null pointer in nvte_tensor_columnwise_shape"); } - const std::vector &shape = t->columnwise_data.shape; - return nvte_make_shape(shape.data(), shape.size()); + return t->columnwise_data.shape; } size_t nvte_tensor_ndims(const NVTETensor tensor) { return nvte_tensor_shape(tensor).ndim; } @@ -962,10 +947,8 @@ NVTEScalingMode nvte_tensor_scaling_mode(const NVTETensor tensor) { } void nvte_tensor_pack_create(NVTETensorPack *pack) { - for (int i = 0; i < pack->MAX_SIZE; i++) { - pack->tensors[i] = - transformer_engine::TensorAllocator::instance().Allocate(NVTE_DELAYED_TENSOR_SCALING); - } + transformer_engine::TensorAllocator::instance().Allocate(NVTE_DELAYED_TENSOR_SCALING, + pack->tensors, pack->MAX_SIZE); } void nvte_tensor_pack_destroy(NVTETensorPack *pack) { diff --git a/transformer_engine/common/transpose/cast_transpose_fusion.cu b/transformer_engine/common/transpose/cast_transpose_fusion.cu index 77c1322e7..619e60220 100644 --- a/transformer_engine/common/transpose/cast_transpose_fusion.cu +++ b/transformer_engine/common/transpose/cast_transpose_fusion.cu @@ -179,8 +179,7 @@ inline __device__ void cast_and_transpose_regs(const CVec (&in)[nvec_out], void populate_cast_transpose_dbias_workspace_config(const Tensor &cast_output, /*cast*/ Tensor *workspace, const int nvec_out) { - const size_t row_length = cast_output.flat_last_dim(); - const size_t num_rows = cast_output.flat_first_dim(); + const auto [num_rows, row_length] = cast_output.flat_2d_dims(); const size_t tile_size_y = (nvec_out * THREADS_PER_WARP); NVTE_CHECK(num_rows % nvec_out == 0, "Unsupported shape."); @@ -549,7 +548,7 @@ void cast_transpose_fused(const Tensor &input, const Tensor *act_input, Tensor * if constexpr (IS_DBIAS) { NVTE_CHECK(dbias->data.dtype == input.data.dtype, "DBias must have the same type as input."); - NVTE_CHECK(dbias->data.shape == std::vector{row_length}, "Wrong shape of DBias."); + NVTE_CHECK(dbias->data.shape == Shape{row_length}, "Wrong shape of DBias."); } if constexpr (IS_DACT) { NVTE_CHECK(input.dtype() == act_input->dtype(), "Types of both inputs must match."); diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu index cf9821f1a..1596bb3fd 100644 --- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu +++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu @@ -780,7 +780,7 @@ void quantize_transpose_vector_blockwise_fp4( Tensor& rng_state_te_tensor = *convertNVTETensor(rng_state_tensor); NVTE_CHECK(rng_state_te_tensor.dtype() == DType::kInt64, "RNG state should contain 2 64-bit values."); - NVTE_CHECK(rng_state_te_tensor.data.shape == std::vector{2}, + NVTE_CHECK(rng_state_te_tensor.data.shape == Shape{2}, "Shape of the RNG state should be [2], but got ", rng_state_te_tensor.data.shape); rng_state = reinterpret_cast(rng_state_te_tensor.data.dptr); } diff --git a/transformer_engine/common/transpose/transpose_fusion.cu b/transformer_engine/common/transpose/transpose_fusion.cu index 670fe6f92..cfb99d214 100644 --- a/transformer_engine/common/transpose/transpose_fusion.cu +++ b/transformer_engine/common/transpose/transpose_fusion.cu @@ -433,7 +433,7 @@ void fp8_transpose_dbias(const Tensor &input, Tensor *transposed_output, Tensor NVTE_CHECK(transposed_output->data.dtype == input.data.dtype, "T output must have the same type as input."); - NVTE_CHECK(dbias->data.shape == std::vector{row_length}, "Wrong shape of DBias."); + NVTE_CHECK(dbias->data.shape == Shape{row_length}, "Wrong shape of DBias."); TRANSFORMER_ENGINE_TYPE_SWITCH_INPUT( dbias->data.dtype, BiasType, diff --git a/transformer_engine/common/util/utils.cu b/transformer_engine/common/util/utils.cu index a183e6ec5..39d826246 100644 --- a/transformer_engine/common/util/utils.cu +++ b/transformer_engine/common/util/utils.cu @@ -7,45 +7,76 @@ #include #include +#include +#include + #include "../common.h" #include "../util/logging.h" +namespace transformer_engine { +namespace copy_host_to_device_via_kernel { namespace { -constexpr int64_t kMaxKernelAddresses = 256; - -struct HostPointersArgs { - uint64_t ptrs[kMaxKernelAddresses]; +union Payload { + static constexpr size_t kMaxBytes = 2048; + static constexpr size_t kVectorSize = 4; + static constexpr size_t kMaxVectors = kMaxBytes / kVectorSize; + uint8_t bytes[kMaxBytes]; + uint32_t vectors[kMaxVectors]; }; -__global__ void write_pointers_kernel(HostPointersArgs args, uint64_t *out, int64_t count, - int64_t offset) { - const int64_t idx = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (idx < count) { - out[offset + idx] = args.ptrs[idx]; +constexpr size_t block_size = 512; +constexpr size_t num_blocks = DIVUP(Payload::kMaxVectors, block_size); + +__global__ void __launch_bounds__(block_size) kernel(Payload payload, size_t num_bytes, void *out) { + const size_t tid = blockIdx.x * blockDim.x + threadIdx.x; + if (Payload::kVectorSize * (tid + 1) <= num_bytes) { + reinterpret_cast(out)[tid] = payload.vectors[tid]; + } else { + for (size_t i = Payload::kVectorSize * tid; i < num_bytes; ++i) { + reinterpret_cast(out)[i] = payload.bytes[i]; + } } } } // namespace +} // namespace copy_host_to_device_via_kernel +} // namespace transformer_engine + +void nvte_copy_host_to_device_via_kernel(const void *host_ptr, void *device_ptr, size_t num_bytes, + cudaStream_t stream) { + NVTE_API_CALL(nvte_copy_host_to_device_via_kernel); + using namespace transformer_engine::copy_host_to_device_via_kernel; + + // Nothing to be done if size is zero + if (num_bytes == 0) { + return; + } + + // Check pointers + NVTE_CHECK(host_ptr != nullptr, "Attempting to read ", num_bytes, " bytes from a null pointer."); + NVTE_CHECK(device_ptr != nullptr, "Attempting to write ", num_bytes, + " bytes into a null pointer."); + NVTE_CHECK(reinterpret_cast(device_ptr) % Payload::kVectorSize == 0, + "Device pointer is not aligned to ", Payload::kVectorSize, " bytes."); + + // Chunk data to fit in kernel arguments and launch kernels + const uint8_t *src = static_cast(host_ptr); + uint8_t *dst = static_cast(device_ptr); + for (size_t offset = 0; offset < num_bytes; offset += Payload::kMaxBytes) { + const size_t chunk_size = std::min(num_bytes - offset, Payload::kMaxBytes); + Payload payload{}; + std::memcpy(payload.bytes, src + offset, chunk_size); + kernel<<>>(payload, chunk_size, dst + offset); + NVTE_CHECK_CUDA(cudaGetLastError()); + } +} void nvte_convert_pointers_to_tensor(const uint64_t *host_ptrs, NVTETensor output, int64_t count, cudaStream_t stream) { NVTE_API_CALL(nvte_convert_pointers_to_tensor); using namespace transformer_engine; Tensor *out_tensor = convertNVTETensorCheck(output); - uint64_t *out_ptr = static_cast(out_tensor->data.dptr); - NVTE_CHECK(out_ptr != nullptr, "Output tensor data pointer is null."); - - int64_t offset = 0; - while (offset < count) { - const int64_t chunk = std::min(kMaxKernelAddresses, count - offset); - HostPointersArgs args{}; - for (int64_t i = 0; i < chunk; ++i) { - args.ptrs[i] = host_ptrs[offset + i]; - } - constexpr int threads = kMaxKernelAddresses; - write_pointers_kernel<<<1, threads, 0, stream>>>(args, out_ptr, chunk, offset); - NVTE_CHECK_CUDA(cudaGetLastError()); - offset += chunk; - } + nvte_copy_host_to_device_via_kernel(host_ptrs, out_tensor->data.dptr, + static_cast(count) * sizeof(uint64_t), stream); } diff --git a/transformer_engine/pytorch/csrc/extensions.h b/transformer_engine/pytorch/csrc/extensions.h index f8ac778aa..1f9974448 100644 --- a/transformer_engine/pytorch/csrc/extensions.h +++ b/transformer_engine/pytorch/csrc/extensions.h @@ -485,14 +485,15 @@ size_t get_cublasLt_version(); size_t get_cudnn_version(); -std::vector convert_host_pointers_to_tensor( - std::vector> tensor_lists); - -std::tuple get_device_pointer_for_data_and_scales( - std::vector data_tensors, std::vector scale_tensors, bool swizzle, - bool rowwise, transformer_engine::DType data_dtype); at::Tensor splits_to_offsets(const at::Tensor &first_dims, int64_t logical_last_dim); +at::Tensor copy_data_ptrs_to_device(const std::vector &tensors, + const c10::Device &device); + +std::tuple> transform_and_copy_data_ptrs_to_device( + const std::string &transform_type, const std::vector &tensors, + const c10::Device &device); + /*************************************************************************************************** * Support THD format for Context Parallel **************************************************************************************************/ diff --git a/transformer_engine/pytorch/csrc/extensions/pybind.cpp b/transformer_engine/pytorch/csrc/extensions/pybind.cpp index 3acef587f..c1b38a227 100644 --- a/transformer_engine/pytorch/csrc/extensions/pybind.cpp +++ b/transformer_engine/pytorch/csrc/extensions/pybind.cpp @@ -492,15 +492,12 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Get cublasLt version", py::call_guard()); m.def("get_cudnn_version", &transformer_engine::pytorch::get_cudnn_version, "Get cuDNN version", py::call_guard()); - m.def("convert_host_pointers_to_tensor", - &transformer_engine::pytorch::convert_host_pointers_to_tensor, - "Copy host-side device pointers into device tensors", py::arg("tensor_lists"), - py::call_guard()); - m.def("get_device_pointer_for_data_and_scales", - &transformer_engine::pytorch::get_device_pointer_for_data_and_scales, - "Swizzle scales and collect data/scale device pointers into device tensors", - py::arg("data_tensors"), py::arg("scale_tensors"), py::arg("swizzle") = false, - py::arg("rowwise"), py::arg("data_dtype"), py::call_guard()); + m.def("copy_data_ptrs_to_device", &transformer_engine::pytorch::copy_data_ptrs_to_device, + py::arg("tensors"), py::arg("device"), py::call_guard()); + m.def("transform_and_copy_data_ptrs_to_device", + &transformer_engine::pytorch::transform_and_copy_data_ptrs_to_device, + py::arg("transform_type"), py::arg("tensors"), py::arg("device"), + py::call_guard()); m.def("splits_to_offsets", &transformer_engine::pytorch::splits_to_offsets, "Compute grouped tensor offsets from split sizes", py::arg("first_dims"), py::arg("logical_last_dim"), py::call_guard()); diff --git a/transformer_engine/pytorch/csrc/extensions/recipe.cpp b/transformer_engine/pytorch/csrc/extensions/recipe.cpp index c02d2ec61..9be288d2f 100644 --- a/transformer_engine/pytorch/csrc/extensions/recipe.cpp +++ b/transformer_engine/pytorch/csrc/extensions/recipe.cpp @@ -35,35 +35,30 @@ void fused_amax_and_scale_update_after_reduction(const at::Tensor& amax_reductio const std::string& amax_compute_algo, DType fp8_dtype, float margin) { size_t num_tensors = amax_histories.size(); - std::vector te_amax_histories; - std::vector te_scales; - te_amax_histories.reserve(num_tensors); - te_scales.reserve(num_tensors); + + // Allocate amax history and scale NVTETensors as batches + MultiTensorWrapper te_amax_histories(num_tensors, NVTE_DELAYED_TENSOR_SCALING); + MultiTensorWrapper te_scales(num_tensors, NVTE_DELAYED_TENSOR_SCALING); + for (size_t i = 0; i < num_tensors; i++) { - te_amax_histories.push_back(nvte_create_tensor(NVTE_DELAYED_TENSOR_SCALING)); - NVTETensor& amax_history = te_amax_histories.back(); NVTEShape amax_shape = convertTorchShape(amax_histories[i].sizes()); NVTEBasicTensor amax_history_data = {amax_histories[i].data_ptr(), static_cast(DType::kFloat32), amax_shape}; - nvte_set_tensor_param(&amax_history, kNVTERowwiseData, &amax_history_data); + nvte_set_tensor_param_v2(te_amax_histories[i], kNVTERowwiseData, &amax_history_data, + sizeof(amax_history_data)); - te_scales.push_back(nvte_create_tensor(NVTE_DELAYED_TENSOR_SCALING)); - NVTETensor& scale = te_scales.back(); NVTEShape scale_shape = convertTorchShape(scales[i].sizes()); NVTEBasicTensor scale_data = {scales[i].data_ptr(), static_cast(DType::kFloat32), scale_shape}; - nvte_set_tensor_param(&scale, kNVTERowwiseData, &scale_data); + nvte_set_tensor_param_v2(te_scales[i], kNVTERowwiseData, &scale_data, sizeof(scale_data)); } + // The recipe function takes std::vector by value, so + // construct fresh vectors from the batches. nvte_delayed_scaling_recipe_amax_and_scale_update_after_reduction( - makeTransformerEngineTensor(amax_reduction_buffer).data(), te_amax_histories, te_scales, - amax_compute_algo.c_str(), static_cast(fp8_dtype), margin, - at::cuda::getCurrentCUDAStream()); - for (auto& t : te_amax_histories) { - nvte_destroy_tensor(t); - } - for (auto& t : te_scales) { - nvte_destroy_tensor(t); - } + makeTransformerEngineTensor(amax_reduction_buffer).data(), + std::vector(te_amax_histories.begin(), te_amax_histories.end()), + std::vector(te_scales.begin(), te_scales.end()), amax_compute_algo.c_str(), + static_cast(fp8_dtype), margin, at::cuda::getCurrentCUDAStream()); } } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp index 193aed29e..c90a7d6d0 100644 --- a/transformer_engine/pytorch/csrc/extensions/swizzle.cpp +++ b/transformer_engine/pytorch/csrc/extensions/swizzle.cpp @@ -173,6 +173,7 @@ std::optional multi_tensor_swizzle_scales_for_gemm_impl( // Filter out tensors that already have swizzled scales std::vector tensors_needing_swizzle; + tensors_needing_swizzle.reserve(tensors.size()); for (auto &tensor : tensors) { if (!tensor.get_with_gemm_swizzled_scales()) { tensors_needing_swizzle.push_back(&tensor); @@ -184,6 +185,7 @@ std::optional multi_tensor_swizzle_scales_for_gemm_impl( // Determine buffer size needed for swizzled scales std::vector output_scales_offsets; + output_scales_offsets.reserve(tensors_needing_swizzle.size()); size_t output_scales_bytes = 0; for (auto &tensor : tensors_needing_swizzle) { const auto scales_nvte = @@ -202,75 +204,80 @@ std::optional multi_tensor_swizzle_scales_for_gemm_impl( transformer_engine::DType::kByte, false); uint8_t *output_scales_dptr = reinterpret_cast(getDataPtr(output_scales_pyt)); - // Construct TE tensors with only scales - std::vector inputs_nvte, outputs_nvte; - for (size_t i = 0; i < tensors_needing_swizzle.size(); ++i) { + // Allocate input/output NVTETensors as a single batch. The first + // n_swizzle entries are inputs; the next n_swizzle are outputs. + const size_t n_swizzle = tensors_needing_swizzle.size(); + MultiTensorWrapper nvte_tensors(2 * n_swizzle, scaling_mode); + NVTETensor *inputs_nvte = nvte_tensors.data(); + NVTETensor *outputs_nvte = nvte_tensors.data() + n_swizzle; + + auto set_param = [](NVTETensor t, NVTETensorParam param, void *dptr, + transformer_engine::DType dtype, const NVTEShape &shape) { + NVTEBasicTensor data{dptr, static_cast(dtype), shape}; + nvte_set_tensor_param_v2(t, param, &data, sizeof(data)); + }; + + // Cache output scale dtype/shape per tensor so we can update the + // source TensorWrappers without re-reading from the output NVTETensors. + std::vector output_scales_dtypes(n_swizzle); + std::vector output_scales_shapes(n_swizzle); + + for (size_t i = 0; i < n_swizzle; ++i) { auto &tensor = *tensors_needing_swizzle[i]; - inputs_nvte.emplace_back(scaling_mode); - outputs_nvte.emplace_back(scaling_mode); - auto &input_nvte = inputs_nvte.back(); - auto &output_nvte = outputs_nvte.back(); - output_nvte.set_with_gemm_swizzled_scales(true); + const uint8_t swizzled_flag = 1; + nvte_set_tensor_param_v2(outputs_nvte[i], kNVTEWithGEMMSwizzledScales, &swizzled_flag, + sizeof(swizzled_flag)); if (rowwise_usage) { const auto data_nvte = tensor.get_rowwise_data(); const auto scales_nvte = tensor.get_rowwise_scale_inv(); const auto data_dtype = static_cast(data_nvte.dtype); const auto scales_dtype = static_cast(scales_nvte.dtype); - input_nvte.set_rowwise_data(nullptr, data_dtype, data_nvte.shape); - input_nvte.set_rowwise_scale_inv(scales_nvte.data_ptr, scales_dtype, scales_nvte.shape); - output_nvte.set_rowwise_data(nullptr, data_dtype, data_nvte.shape); - output_nvte.set_rowwise_scale_inv(output_scales_dptr + output_scales_offsets[i], scales_dtype, - scales_nvte.shape); + output_scales_dtypes[i] = scales_dtype; + output_scales_shapes[i] = scales_nvte.shape; + set_param(inputs_nvte[i], kNVTERowwiseData, nullptr, data_dtype, data_nvte.shape); + set_param(inputs_nvte[i], kNVTERowwiseScaleInv, scales_nvte.data_ptr, scales_dtype, + scales_nvte.shape); + set_param(outputs_nvte[i], kNVTERowwiseData, nullptr, data_dtype, data_nvte.shape); + set_param(outputs_nvte[i], kNVTERowwiseScaleInv, + output_scales_dptr + output_scales_offsets[i], scales_dtype, scales_nvte.shape); } else { const auto data_nvte = tensor.get_columnwise_data(); const auto scales_nvte = tensor.get_columnwise_scale_inv(); const auto data_dtype = static_cast(data_nvte.dtype); const auto scales_dtype = static_cast(scales_nvte.dtype); - input_nvte.set_columnwise_data(nullptr, data_dtype, data_nvte.shape); - input_nvte.set_columnwise_scale_inv(scales_nvte.data_ptr, scales_dtype, scales_nvte.shape); - output_nvte.set_columnwise_data(nullptr, data_dtype, data_nvte.shape); - output_nvte.set_columnwise_scale_inv(output_scales_dptr + output_scales_offsets[i], - scales_dtype, scales_nvte.shape); + output_scales_dtypes[i] = scales_dtype; + output_scales_shapes[i] = scales_nvte.shape; + set_param(inputs_nvte[i], kNVTEColumnwiseData, nullptr, data_dtype, data_nvte.shape); + set_param(inputs_nvte[i], kNVTEColumnwiseScaleInv, scales_nvte.data_ptr, scales_dtype, + scales_nvte.shape); + set_param(outputs_nvte[i], kNVTEColumnwiseData, nullptr, data_dtype, data_nvte.shape); + set_param(outputs_nvte[i], kNVTEColumnwiseScaleInv, + output_scales_dptr + output_scales_offsets[i], scales_dtype, scales_nvte.shape); } } - // Pack raw NVTETensors into vectors - std::vector inputs_nvte_raw, outputs_nvte_raw; - for (auto &tensor : inputs_nvte) { - inputs_nvte_raw.emplace_back(tensor.data()); - } - for (auto &tensor : outputs_nvte) { - outputs_nvte_raw.emplace_back(tensor.data()); - } - // Launch kernel NVTE_SCOPED_GIL_RELEASE({ if (check_scale_inv_shapes) { - nvte_multi_tensor_swizzle_scaling_factors(inputs_nvte_raw.data(), outputs_nvte_raw.data(), - inputs_nvte_raw.size(), + nvte_multi_tensor_swizzle_scaling_factors(inputs_nvte, outputs_nvte, n_swizzle, at::cuda::getCurrentCUDAStream()); } else { - nvte_multi_tensor_swizzle_scaling_factors_unchecked( - inputs_nvte_raw.data(), outputs_nvte_raw.data(), inputs_nvte_raw.size(), - at::cuda::getCurrentCUDAStream()); + nvte_multi_tensor_swizzle_scaling_factors_unchecked(inputs_nvte, outputs_nvte, n_swizzle, + at::cuda::getCurrentCUDAStream()); } }); // Update tensors with swizzled scales - for (size_t i = 0; i < tensors_needing_swizzle.size(); ++i) { + for (size_t i = 0; i < n_swizzle; ++i) { auto &tensor = *tensors_needing_swizzle[i]; reset_tensor_data(tensor, !rowwise_usage, !columnwise_usage); tensor.set_with_gemm_swizzled_scales(true); if (rowwise_usage) { - auto scales_nvte = outputs_nvte[i].get_rowwise_scale_inv(); - const auto scales_dtype = static_cast(scales_nvte.dtype); - tensor.set_rowwise_scale_inv(output_scales_dptr + output_scales_offsets[i], scales_dtype, - scales_nvte.shape); + tensor.set_rowwise_scale_inv(output_scales_dptr + output_scales_offsets[i], + output_scales_dtypes[i], output_scales_shapes[i]); } else { - auto scales_nvte = outputs_nvte[i].get_columnwise_scale_inv(); - const auto scales_dtype = static_cast(scales_nvte.dtype); - tensor.set_columnwise_scale_inv(output_scales_dptr + output_scales_offsets[i], scales_dtype, - scales_nvte.shape); + tensor.set_columnwise_scale_inv(output_scales_dptr + output_scales_offsets[i], + output_scales_dtypes[i], output_scales_shapes[i]); } } diff --git a/transformer_engine/pytorch/csrc/extensions/utils.cpp b/transformer_engine/pytorch/csrc/extensions/utils.cpp index 9a093608d..453f238c0 100644 --- a/transformer_engine/pytorch/csrc/extensions/utils.cpp +++ b/transformer_engine/pytorch/csrc/extensions/utils.cpp @@ -6,6 +6,10 @@ #include +#include +#include +#include +#include #include #include "common/common.h" @@ -13,153 +17,161 @@ namespace transformer_engine::pytorch { -namespace { - -at::Tensor collect_pointers_in_device_tensor(const std::vector& host_ptrs, - const at::Device& device, cudaStream_t stream) { - const int64_t count = static_cast(host_ptrs.size()); - auto out = at::empty({count}, at::TensorOptions().dtype(at::kLong).device(device)); - auto out_nvte = makeTransformerEngineTensor(out); - nvte_convert_pointers_to_tensor(host_ptrs.data(), out_nvte.data(), count, stream); - return out; -} +at::Tensor copy_data_ptrs_to_device(const std::vector &tensors, + const c10::Device &device) { + // Collect data pointers + std::vector ptrs_host; + ptrs_host.reserve(tensors.size()); + for (const auto &tensor : tensors) { + ptrs_host.push_back(reinterpret_cast(tensor.data_ptr())); + } -} // namespace + // Allocate device buffer + auto ptrs_device = at::empty({static_cast(tensors.size())}, + at::TensorOptions().dtype(at::kLong).device(device)); -std::vector convert_host_pointers_to_tensor( - std::vector> tensor_lists) { - std::vector outputs; - outputs.reserve(tensor_lists.size()); - auto stream = at::cuda::getCurrentCUDAStream(); + // Load pointers on device + nvte_copy_host_to_device_via_kernel(ptrs_host.data(), ptrs_device.data_ptr(), + tensors.size() * sizeof(uint64_t), + at::cuda::getCurrentCUDAStream()); - for (const auto& tensor_list : tensor_lists) { - NVTE_CHECK(!tensor_list.empty(), "Tensor list is empty."); - const auto& first_tensor = tensor_list[0]; - NVTE_CHECK(first_tensor.is_cuda(), "Tensor list must be on CUDA."); - const auto device = first_tensor.device(); - const int64_t count = static_cast(tensor_list.size()); - std::vector host_ptrs(count); - for (int64_t i = 0; i < count; ++i) { - host_ptrs[i] = reinterpret_cast(tensor_list[static_cast(i)].data_ptr()); - } - outputs.push_back(collect_pointers_in_device_tensor(host_ptrs, device, stream)); - } - - return outputs; + return ptrs_device; } -std::tuple get_device_pointer_for_data_and_scales( - std::vector data_tensors, std::vector scale_tensors, bool swizzle, - bool rowwise, transformer_engine::DType data_dtype) { - const size_t num_tensors = data_tensors.size(); - NVTE_CHECK(num_tensors > 0, "data_tensors must not be empty."); - NVTE_CHECK(num_tensors == scale_tensors.size(), - "data_tensors and scale_tensors must have the same size."); - NVTE_CHECK(data_tensors[0].is_cuda(), "data_tensors must be on CUDA."); - const auto device = data_tensors[0].device(); - auto stream = at::cuda::getCurrentCUDAStream(); +std::tuple> transform_and_copy_data_ptrs_to_device( + const std::string &transform_type, const std::vector &tensors, + const c10::Device &device) { + const size_t num_tensors = tensors.size(); - // Infer data shape from the first data tensor (expected 2D: n x k) - NVTE_CHECK(data_tensors[0].dim() == 2, - "data_tensors elements must be 2D, got dim=", data_tensors[0].dim()); - NVTEShape data_shape{}; - data_shape.ndim = 2; - data_shape.data[0] = static_cast(data_tensors[0].size(0)); - data_shape.data[1] = static_cast(data_tensors[0].size(1)); - - // Collect data device pointers - std::vector data_host_ptrs(num_tensors); - for (size_t i = 0; i < num_tensors; ++i) { - data_host_ptrs[i] = reinterpret_cast(data_tensors[i].data_ptr()); + // Trivial cases + if (transform_type.empty()) { + // No transform, just load pointers on device + return {copy_data_ptrs_to_device(tensors, device), std::nullopt}; + } + if (num_tensors == 0) { + // No input tensors, return tensor with no elements + return {at::empty({int64_t{0}}, at::TensorOptions().dtype(at::kLong).device(device)), + std::nullopt}; } - // Swizzle scales and collect scale pointers - at::Tensor swizzled_scales_keepalive; - std::vector scale_host_ptrs(num_tensors); + // CUDA stream + auto stream = at::cuda::getCurrentCUDAStream(); - if (swizzle) { - NVTEScalingMode scaling_mode; - transformer_engine::DType scale_dtype; - if (is_fp8_dtype(data_dtype)) { + // Swizzle scales for GEMM, with uniform tensor sizes + const bool uniform_mxfp8_rowwise_swizzle = transform_type == "uniform_mxfp8_rowwise_swizzle"; + const bool uniform_mxfp8_colwise_swizzle = transform_type == "uniform_mxfp8_columnwise_swizzle"; + const bool uniform_nvfp4_swizzle = transform_type == "uniform_nvfp4_swizzle"; + if (uniform_mxfp8_rowwise_swizzle || uniform_mxfp8_colwise_swizzle || uniform_nvfp4_swizzle) { + // Tensor format + NVTEScalingMode scaling_mode = NVTE_INVALID_SCALING; + if (uniform_mxfp8_rowwise_swizzle || uniform_mxfp8_colwise_swizzle) { scaling_mode = NVTE_MXFP8_1D_SCALING; - scale_dtype = transformer_engine::DType::kFloat8E8M0; - } else if (is_fp4_dtype(data_dtype)) { + } else if (uniform_nvfp4_swizzle) { scaling_mode = NVTE_NVFP4_1D_SCALING; - scale_dtype = transformer_engine::DType::kFloat8E4M3; - } else { - NVTE_ERROR("data_dtype must be an FP8 or FP4 type for swizzling."); } - // Compute output buffer size for swizzled scales (16B aligned per tensor) - std::vector output_offsets; - size_t output_bytes = 0; - for (size_t i = 0; i < num_tensors; ++i) { - const size_t scale_numel = static_cast(scale_tensors[i].numel()); - const size_t dtype_bits = transformer_engine::pytorch::typeToNumBits(scale_dtype); - output_bytes = roundup(output_bytes, 16); - output_offsets.push_back(output_bytes); - output_bytes += ceildiv(scale_numel * dtype_bits, 8); + // Data types + transformer_engine::DType data_dtype, scale_dtype; + switch (scaling_mode) { + case NVTE_MXFP8_1D_SCALING: + data_dtype = transformer_engine::DType::kFloat8E4M3; + scale_dtype = transformer_engine::DType::kFloat8E8M0; + break; + case NVTE_NVFP4_1D_SCALING: + data_dtype = transformer_engine::DType::kFloat4E2M1; + scale_dtype = transformer_engine::DType::kFloat8E4M3; + break; + default: + NVTE_ERROR("Unsupported case."); } - // Allocate single buffer for all swizzled scales - swizzled_scales_keepalive = - allocateSpace(std::vector{output_bytes}, transformer_engine::DType::kByte, false); - uint8_t* output_dptr = reinterpret_cast(getDataPtr(swizzled_scales_keepalive)); + // Scale shape + const NVTEShape scale_shape = convertTorchShape(tensors[0].sizes()); + NVTE_CHECK(scale_shape.ndim == 2, + "Expected 2D scale tensor, but got shape=", getTensorShape(tensors[0]), "."); + const size_t scale_numel = scale_shape.data[0] * scale_shape.data[1]; + const size_t scale_dtype_bits = transformer_engine::pytorch::typeToNumBits(scale_dtype); + const size_t scale_bytes = ceildiv(scale_numel * scale_dtype_bits, 8); + + // Expected data shape + // Note: May not match actual data shape since the scales are padded. + // This is fine since we're not actually touching the data. + NVTEShape data_shape; + data_shape.ndim = 2; + if (uniform_mxfp8_rowwise_swizzle) { + data_shape.data[0] = scale_shape.data[0]; + data_shape.data[1] = scale_shape.data[1] * 32; + } else if (uniform_mxfp8_colwise_swizzle) { + data_shape.data[0] = scale_shape.data[0] * 32; + data_shape.data[1] = scale_shape.data[1]; + } else if (uniform_nvfp4_swizzle) { + data_shape.data[0] = scale_shape.data[0]; + data_shape.data[1] = scale_shape.data[1] * 16; + } else { + NVTE_ERROR("Unsupported case."); + } + + // Allocate single buffer for swizzled scales. + // Uses a uniform stride since all tensors share the same scale shape. + const size_t swizzled_scales_stride = roundup(scale_bytes, 16); // Align to 16 bytes + auto swizzled_scales = at::empty({static_cast(swizzled_scales_stride * num_tensors)}, + at::TensorOptions().dtype(at::kByte).device(device)); + uint8_t *swizzled_scales_dptr = reinterpret_cast(swizzled_scales.data_ptr()); + + // Allocate input/output NVTETensors as a single batch. The first + // num_tensors entries are inputs; the next num_tensors are outputs. + MultiTensorWrapper nvte_tensors(2 * num_tensors, scaling_mode); + NVTETensor *inputs_nvte = nvte_tensors.data(); + NVTETensor *outputs_nvte = nvte_tensors.data() + num_tensors; + + auto set_param = [](NVTETensor t, NVTETensorParam param, void *dptr, + transformer_engine::DType dtype, const NVTEShape &shape) { + NVTEBasicTensor data{dptr, static_cast(dtype), shape}; + nvte_set_tensor_param_v2(t, param, &data, sizeof(data)); + }; - // Build TensorWrapper input/output pairs and get scale shapes - std::vector inputs_nvte, outputs_nvte; - inputs_nvte.reserve(num_tensors); - outputs_nvte.reserve(num_tensors); for (size_t i = 0; i < num_tensors; ++i) { - inputs_nvte.emplace_back(scaling_mode); - outputs_nvte.emplace_back(scaling_mode); - auto& input_nvte = inputs_nvte.back(); - auto& output_nvte = outputs_nvte.back(); - output_nvte.set_with_gemm_swizzled_scales(true); - - NVTEShape scale_shape = convertTorchShape(scale_tensors[i].sizes()); - void* scale_ptr = scale_tensors[i].data_ptr(); - uint8_t* out_scale_ptr = output_dptr + output_offsets[i]; - - if (rowwise) { - input_nvte.set_rowwise_data(nullptr, data_dtype, data_shape); - input_nvte.set_rowwise_scale_inv(scale_ptr, scale_dtype, scale_shape); - output_nvte.set_rowwise_data(nullptr, data_dtype, data_shape); - output_nvte.set_rowwise_scale_inv(out_scale_ptr, scale_dtype, scale_shape); - } else { - input_nvte.set_columnwise_data(nullptr, data_dtype, data_shape); - input_nvte.set_columnwise_scale_inv(scale_ptr, scale_dtype, scale_shape); - output_nvte.set_columnwise_data(nullptr, data_dtype, data_shape); - output_nvte.set_columnwise_scale_inv(out_scale_ptr, scale_dtype, scale_shape); + const uint8_t swizzled_flag = 1; + nvte_set_tensor_param_v2(outputs_nvte[i], kNVTEWithGEMMSwizzledScales, &swizzled_flag, + sizeof(swizzled_flag)); + void *in_scale_ptr = tensors[i].data_ptr(); + void *out_scale_ptr = swizzled_scales_dptr + i * swizzled_scales_stride; + if (uniform_mxfp8_rowwise_swizzle || uniform_nvfp4_swizzle) { + set_param(inputs_nvte[i], kNVTERowwiseData, nullptr, data_dtype, data_shape); + set_param(inputs_nvte[i], kNVTERowwiseScaleInv, in_scale_ptr, scale_dtype, scale_shape); + set_param(outputs_nvte[i], kNVTERowwiseData, nullptr, data_dtype, data_shape); + set_param(outputs_nvte[i], kNVTERowwiseScaleInv, out_scale_ptr, scale_dtype, scale_shape); + } else if (uniform_mxfp8_colwise_swizzle) { + set_param(inputs_nvte[i], kNVTEColumnwiseData, nullptr, data_dtype, data_shape); + set_param(inputs_nvte[i], kNVTEColumnwiseScaleInv, in_scale_ptr, scale_dtype, scale_shape); + set_param(outputs_nvte[i], kNVTEColumnwiseData, nullptr, data_dtype, data_shape); + set_param(outputs_nvte[i], kNVTEColumnwiseScaleInv, out_scale_ptr, scale_dtype, + scale_shape); } } - // Pack raw NVTETensors and launch swizzle kernel - std::vector inputs_raw, outputs_raw; - inputs_raw.reserve(num_tensors); - outputs_raw.reserve(num_tensors); - for (auto& t : inputs_nvte) inputs_raw.push_back(t.data()); - for (auto& t : outputs_nvte) outputs_raw.push_back(t.data()); - - nvte_multi_tensor_swizzle_scaling_factors(inputs_raw.data(), outputs_raw.data(), num_tensors, - stream); + // Launch kernel + nvte_multi_tensor_swizzle_scaling_factors(inputs_nvte, outputs_nvte, num_tensors, stream); - // Collect swizzled scale pointers + // Collect data pointers + std::vector ptrs_host; + ptrs_host.reserve(num_tensors); for (size_t i = 0; i < num_tensors; ++i) { - scale_host_ptrs[i] = reinterpret_cast(output_dptr + output_offsets[i]); + ptrs_host.push_back( + reinterpret_cast(swizzled_scales_dptr + i * swizzled_scales_stride)); } - } else { - swizzled_scales_keepalive = at::empty({0}, at::TensorOptions().dtype(at::kByte).device(device)); - for (size_t i = 0; i < num_tensors; ++i) { - scale_host_ptrs[i] = reinterpret_cast(scale_tensors[i].data_ptr()); - } - } - // Convert pointer arrays to device tensors - auto data_ptrs = collect_pointers_in_device_tensor(data_host_ptrs, device, stream); - auto scale_ptrs = collect_pointers_in_device_tensor(scale_host_ptrs, device, stream); + // Load pointers on device + auto ptrs_device = at::empty({static_cast(num_tensors)}, + at::TensorOptions().dtype(at::kLong).device(device)); + nvte_copy_host_to_device_via_kernel(ptrs_host.data(), ptrs_device.data_ptr(), + num_tensors * sizeof(uint64_t), stream); + + return {std::move(ptrs_device), std::move(swizzled_scales)}; + } - return {std::move(data_ptrs), std::move(scale_ptrs), std::move(swizzled_scales_keepalive)}; + // Unsupported transform + NVTE_ERROR("Unsupported transform type (", transform_type, ")"); } } // namespace transformer_engine::pytorch diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 802c1a25d..25ccad137 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -118,7 +118,7 @@ def _cudnn_compute_wgrad( ) else: # Discrete mode: per-expert wgrad device pointers - (wgrad_ptrs,) = tex.convert_host_pointers_to_tensor([wgrad_output]) + wgrad_ptrs = tex.copy_data_ptrs_to_device(wgrad_output, wgrad_output[0].device) wgrad_kernel_fn( a_tensor=a_tensor, b_tensor=b_tensor, @@ -530,12 +530,14 @@ def fuser_backward( fc2_dactivation_kwargs["b_tensor"] = fc2_w_data fc2_dactivation_kwargs["sfb_tensor"] = fc2_w_scales else: - fc2_b_ptrs, fc2_sfb_ptrs, _fc2_sw = tex.get_device_pointer_for_data_and_scales( + fc2_b_ptrs = tex.copy_data_ptrs_to_device( [w._columnwise_data for w in grouped_fc2_weight], + device, + ) + fc2_sfb_ptrs, _fc2_sfb_buffer = tex.transform_and_copy_data_ptrs_to_device( + "uniform_mxfp8_columnwise_swizzle", [w._columnwise_scale_inv for w in grouped_fc2_weight], - swizzle=True, - rowwise=False, - data_dtype=grouped_fc2_weight[0]._fp8_dtype, + device, ) fc2_dactivation_kwargs["b_ptrs"] = fc2_b_ptrs fc2_dactivation_kwargs["sfb_ptrs"] = fc2_sfb_ptrs @@ -719,14 +721,15 @@ def fuser_backward( fc1_dgrad_kwargs["b_tensor"] = fc1_w_data fc1_dgrad_kwargs["sfb_tensor"] = fc1_w_scales else: - fc1_b_ptrs, fc1_sfb_ptrs, _ = tex.get_device_pointer_for_data_and_scales( + fc1_b_ptrs = tex.copy_data_ptrs_to_device( [w._columnwise_data for w in grouped_fc1_weight], + device, + ) + fc1_sfb_ptrs, _fc1_sfb_buffer = tex.transform_and_copy_data_ptrs_to_device( + "uniform_mxfp8_columnwise_swizzle", [w._columnwise_scale_inv for w in grouped_fc1_weight], - swizzle=True, - rowwise=False, - data_dtype=grouped_fc1_weight[0]._fp8_dtype, + device, ) - fc1_dgrad_kwargs["b_ptrs"] = fc1_b_ptrs fc1_dgrad_kwargs["sfb_ptrs"] = fc1_sfb_ptrs fc1_dgrad_kwargs["n"] = fc1_weight_shape[1] diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index 034d40443..a0c5f766c 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -381,12 +381,14 @@ def fuser_forward( fc1_activation_kwargs["sfb_tensor"] = fc1_w_scales else: # Discrete-weight kernel: per-expert data/scale pointers - fc1_b_ptrs, fc1_sfb_ptrs, _fc1_sw = tex.get_device_pointer_for_data_and_scales( + fc1_b_ptrs = tex.copy_data_ptrs_to_device( [w._rowwise_data for w in grouped_fc1_weight], + device, + ) + fc1_sfb_ptrs, _fc1_sfb_buffer = tex.transform_and_copy_data_ptrs_to_device( + "uniform_mxfp8_rowwise_swizzle", [w._rowwise_scale_inv for w in grouped_fc1_weight], - swizzle=True, - rowwise=True, - data_dtype=grouped_fc1_weight[0]._fp8_dtype, + device, ) fc1_activation_kwargs["b_ptrs"] = fc1_b_ptrs fc1_activation_kwargs["sfb_ptrs"] = fc1_sfb_ptrs @@ -480,12 +482,14 @@ def fuser_forward( fc2_quant_kwargs["b_tensor"] = fc2_w_data fc2_quant_kwargs["sfb_tensor"] = fc2_w_scales else: - fc2_b_ptrs, fc2_sfb_ptrs, _ = tex.get_device_pointer_for_data_and_scales( + fc2_b_ptrs = tex.copy_data_ptrs_to_device( [w._rowwise_data for w in grouped_fc2_weight], + device, + ) + fc2_sfb_ptrs, _fc2_sfb_buffer = tex.transform_and_copy_data_ptrs_to_device( + "uniform_mxfp8_rowwise_swizzle", [w._rowwise_scale_inv for w in grouped_fc2_weight], - swizzle=True, - rowwise=True, - data_dtype=grouped_fc2_weight[0]._fp8_dtype, + device, ) fc2_quant_kwargs["b_ptrs"] = fc2_b_ptrs fc2_quant_kwargs["sfb_ptrs"] = fc2_sfb_ptrs From f8bda5d0ad5f5a8af072c436ed8e070b5bba66cd Mon Sep 17 00:00:00 2001 From: Xin Yao Date: Sat, 30 May 2026 01:45:19 +0800 Subject: [PATCH 090/180] [PyTorch] Make `modules.GroupedLinear` graph-safe (#3038) * make modules.GroupedLinear graph-safe Signed-off-by: Xin Yao * fix tests Signed-off-by: Xin Yao * Review suggestions Handle tensor splits in both legacy and graph-safe impls. Create weight grad tensors as subviews of a larger buffer. Signed-off-by: Tim Moon * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --------- Signed-off-by: Xin Yao Signed-off-by: Tim Moon Co-authored-by: Tim Moon Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- benchmarks/linear/benchmark_grouped_linear.py | 3 + qa/L0_pytorch_unittest/test.sh | 1 + tests/pytorch/test_grouped_linear.py | 1742 +++++++++++++++++ tests/pytorch/test_numerics.py | 1304 +----------- .../pytorch/module/grouped_linear.py | 598 +++++- 5 files changed, 2326 insertions(+), 1322 deletions(-) create mode 100644 tests/pytorch/test_grouped_linear.py diff --git a/benchmarks/linear/benchmark_grouped_linear.py b/benchmarks/linear/benchmark_grouped_linear.py index 815e367f7..cf88faac4 100644 --- a/benchmarks/linear/benchmark_grouped_linear.py +++ b/benchmarks/linear/benchmark_grouped_linear.py @@ -3,6 +3,7 @@ # See LICENSE for license information. import argparse +import os import torch import torch.utils.benchmark as benchmark import pandas as pd @@ -185,6 +186,8 @@ def run_benchmark_linear( x = torch.randn((m, k), dtype=torch.bfloat16, device=device, requires_grad=True) ws = [torch.randn((n, k), dtype=torch.bfloat16, device=device) for _ in range(num_gemms)] m_splits = [m // num_gemms] * num_gemms if m_splits_provided is None else m_splits_provided + if bool(int(os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "0"))): + m_splits = torch.tensor(m_splits, dtype=torch.int64, device=device) # Bias is not supported for GroupedLinear benchmark bias = None diff --git a/qa/L0_pytorch_unittest/test.sh b/qa/L0_pytorch_unittest/test.sh index c35dc4c06..2d3f75f29 100644 --- a/qa/L0_pytorch_unittest/test.sh +++ b/qa/L0_pytorch_unittest/test.sh @@ -29,6 +29,7 @@ python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_recipe.xml $TE_P python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_custom_recipe.xml $TE_PATH/tests/pytorch/test_custom_recipe.py || test_fail "test_custom_recipe.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_deferred_init.xml $TE_PATH/tests/pytorch/test_deferred_init.py || test_fail "test_deferred_init.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_numerics.xml $TE_PATH/tests/pytorch/test_numerics.py || test_fail "test_numerics.py" +PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_grouped_linear.xml $TE_PATH/tests/pytorch/test_grouped_linear.py || test_fail "test_grouped_linear.py" PYTORCH_JIT=0 NVTE_TORCH_COMPILE=0 NVTE_ALLOW_NONDETERMINISTIC_ALGO=0 NVTE_FUSED_ATTN=0 python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_cuda_graphs.xml $TE_PATH/tests/pytorch/test_cuda_graphs.py || test_fail "test_cuda_graphs.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_jit.xml $TE_PATH/tests/pytorch/test_jit.py || test_fail "test_jit.py" python3 -m pytest --tb=auto --junitxml=$XML_LOG_DIR/pytest_test_fused_rope.xml $TE_PATH/tests/pytorch/test_fused_rope.py || test_fail "test_fused_rope.py" diff --git a/tests/pytorch/test_grouped_linear.py b/tests/pytorch/test_grouped_linear.py new file mode 100644 index 000000000..0dc253c18 --- /dev/null +++ b/tests/pytorch/test_grouped_linear.py @@ -0,0 +1,1742 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +import os +import random +from typing import Dict, List, Optional + +import pytest +import torch +import torch.nn as nn +from torch.nn import Parameter + +import transformer_engine.pytorch as te +from transformer_engine.common import recipe +from transformer_engine.pytorch import ( + Float8Quantizer, + Fp8Padding, + Fp8Unpadding, + GroupedLinear, + Linear, + MXFP8Quantizer, + autocast, + is_bf16_available, + quantized_model_init, +) +from transformer_engine.pytorch.cpp_extensions import ( + general_gemm, + general_grouped_gemm, + general_grouped_gemm_for_grouped_tensor, +) +from transformer_engine.pytorch.quantization import ( + FP8GlobalStateManager, + get_align_size_for_quantization, +) +from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor +import transformer_engine_torch as tex +from utils import ModelConfig, recipe_id, reset_rng_states, skip_unsupported_backward_override + +# Only run FP8 tests on supported devices. +fp8_available, reason_for_no_fp8 = te.is_fp8_available(return_reason=True) +fp8_block_scaling_available, _ = te.is_fp8_block_scaling_available(return_reason=True) +mxfp8_available, reason_for_no_mxfp8 = te.is_mxfp8_available(return_reason=True) +nvfp4_available, _ = te.is_nvfp4_available(return_reason=True) + +seed = 1234 +reset_rng_states() + +NVTE_TEST_NVINSPECT_ENABLED = int(os.environ.get("NVTE_TEST_NVINSPECT_ENABLED", "0")) + +if NVTE_TEST_NVINSPECT_ENABLED: + import nvdlfw_inspect.api as debug_api + + debug_api.initialize( + os.environ["NVTE_TEST_NVINSPECT_CONFIG_FILE"], + feature_dirs=os.environ["NVTE_TEST_NVINSPECT_FEATURE_DIRS"], + ) + + +model_configs = { + "126m": ModelConfig(1, 2048, 12, 64, num_layers=12), +} + + +def nvfp4_rht_and_2d_quantization(): + nvfp4_recipe = recipe.NVFP4BlockScaling() + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams( + random_hadamard_transform=True, fp4_2d_quantization=False + ) + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams( + random_hadamard_transform=False, fp4_2d_quantization=True + ) + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams( + random_hadamard_transform=True, fp4_2d_quantization=False + ) + return nvfp4_recipe + + +def nvfp4_row_scaled(): + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + disable_2d_quantization=True, + row_scaled_activation=True, + backward_override="high_precision", + ) + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams() + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + return nvfp4_recipe + + +def nvfp4_4over6(): + nvfp4_recipe = recipe.NVFP4BlockScaling( + disable_rht=True, + disable_stochastic_rounding=True, + nvfp4_4over6="all", + ) + nvfp4_recipe.fp4_quant_fwd_inp = recipe.QParams() + nvfp4_recipe.fp4_quant_fwd_weight = recipe.QParams(fp4_2d_quantization=True) + nvfp4_recipe.fp4_quant_bwd_grad = recipe.QParams() + return nvfp4_recipe + + +def check_rht_usage(recipe: recipe.Recipe) -> bool: + if recipe.nvfp4(): + if ( + recipe.fp4_quant_fwd_inp.random_hadamard_transform + or recipe.fp4_quant_fwd_weight.random_hadamard_transform + or recipe.fp4_quant_bwd_grad.random_hadamard_transform + ): + return True + return False + + +def get_nvfp4_inp_supported_dtypes(recipe: recipe.Recipe, dtype: torch.dtype) -> bool: + supported_input_dtypes = [] + if recipe.nvfp4(): + supported_input_dtypes.append(torch.bfloat16) + if not check_rht_usage(recipe): + supported_input_dtypes.append(torch.float32) + return supported_input_dtypes + + +def dtype_tols(dtype: torch.dtype) -> Dict[str, float]: + if dtype == torch.float32: + return dict(rtol=1.3e-6, atol=1e-5) + if dtype == torch.float16: + return dict(rtol=1e-3, atol=1e-5) + if dtype == torch.bfloat16: + return dict(rtol=1.6e-2, atol=1e-5) + raise ValueError(f"Unsupported dtype ({dtype})") + + +param_types = [torch.float32, torch.float16] +if is_bf16_available(): + param_types.append(torch.bfloat16) + +batch_sizes = [1, 2] +all_boolean = [True, False] + +fp8_recipes = [] +if mxfp8_available: + fp8_recipes.append(recipe.MXFP8BlockScaling()) +if fp8_block_scaling_available: + fp8_recipes.append(recipe.Float8BlockScaling()) +if fp8_available: + fp8_recipes.append(recipe.Float8CurrentScaling()) + fp8_recipes.append(recipe.DelayedScaling()) +if nvfp4_available: + fp8_recipes.append(nvfp4_rht_and_2d_quantization()) + fp8_recipes.append(nvfp4_4over6()) + fp8_recipes.append(nvfp4_row_scaled()) + +use_cutlass_grouped_gemm = [False] +if torch.cuda.get_device_capability() == (9, 0): + use_cutlass_grouped_gemm.append(True) + + +class TorchGroupedLinearWithPadding(nn.Module): + + def __init__( + self, num_gemms, in_features, out_features, bias, params_dtype, parallel_mode, fp8 + ) -> None: + super().__init__() + + self.padding = Fp8Padding(num_gemms) + self.linear_fn = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=bias, + params_dtype=params_dtype, + parallel_mode=parallel_mode, + device="cuda", + ) + self.unpadding = Fp8Unpadding(num_gemms) + + self.fp8 = fp8 + + def forward(self, inp: torch.Tensor, m_splits: List[int]) -> torch.Tensor: + if self.fp8: + orig_m_splits = m_splits + inp, m_splits = self.padding(inp, m_splits) + + out = self.linear_fn(inp, m_splits) + + if self.fp8: + out = self.unpadding(out, orig_m_splits) + + return out + + +def _test_grouped_linear_accuracy( + block, + num_gemms, + bs, + dtype, + config, + recipe, + fp8, + fuse_wgrad_accumulation, + delay_wgrad_compute=False, +): + reset_rng_states() + if fp8: + FP8GlobalStateManager.reset() + + inp_hidden_states = torch.randn( + (config.max_seqlen_q, bs, config.hidden_size), + dtype=dtype, + device="cuda", + requires_grad=True, + ) + inp_hidden_states.retain_grad() + + if num_gemms > 1: + split_size = 1 + if fp8: + split_size = get_align_size_for_quantization(recipe) + m = config.max_seqlen_q // split_size + dist = torch.sort(torch.randint(0, m, (num_gemms - 2,))).values.tolist() + dist.append(dist[-1]) # Manually add a zero + m_splits = torch.tensor(dist + [m]) - torch.tensor([0] + dist) + m_splits = m_splits * split_size + assert m_splits.sum() == config.max_seqlen_q and len(m_splits) == num_gemms + else: + m_splits = torch.tensor([config.max_seqlen_q]) + + with autocast(enabled=fp8, recipe=recipe): + if isinstance(block, GroupedLinear): + m_splits = m_splits * bs + out = block(inp_hidden_states, m_splits.tolist()) + else: + out = torch.cat( + [ + block[i](inp) + for i, inp in enumerate(torch.split(inp_hidden_states, m_splits.tolist())) + ] + ) + loss = out.sum() + loss.backward() + if delay_wgrad_compute: + if isinstance(block, GroupedLinear): + block.backward_dw() + else: + for i in range(num_gemms): + block[i].backward_dw() + + torch.cuda.synchronize() + outputs = [out, inp_hidden_states.grad] + for p in block.parameters(): + if p.requires_grad: + if getattr(p, "main_grad", None) is not None: + outputs.append(p.main_grad) + assert p.grad is None # grad should be None if fuse_wgrad_accumulation is True + else: + outputs.append(p.grad) + return outputs + + +@pytest.mark.parametrize("dtype", param_types, ids=str) +@pytest.mark.parametrize("num_gemms", [3, 6]) +@pytest.mark.parametrize("bs", batch_sizes) +@pytest.mark.parametrize("model", ["126m"]) +@pytest.mark.parametrize("recipe", fp8_recipes + [None], ids=recipe_id) +@pytest.mark.parametrize("fp8_model_params", all_boolean) +@pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean) +@pytest.mark.parametrize("bias", all_boolean) +@pytest.mark.parametrize("delay_wgrad_compute", all_boolean) +def test_grouped_linear_accuracy( + dtype, + num_gemms, + bs, + model, + recipe, + fp8_model_params, + fuse_wgrad_accumulation, + bias, + delay_wgrad_compute, + parallel_mode=None, + use_cutlass=False, +): + fp8 = recipe is not None + if fp8 and fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: + pytest.skip("FP8 parameters are not supported in debug mode.") + if NVTE_TEST_NVINSPECT_ENABLED and delay_wgrad_compute: + pytest.skip("Delayed wgrad compute is not supported in debug mode.") + skip_unsupported_backward_override( + "grouped_linear", recipe, getattr(recipe, "backward_override", None) + ) + + config = model_configs[model] + if config.max_seqlen_q % 16 != 0 and fp8: + pytest.skip("FP8 requires sequence length to be divisible by 16.") + + if recipe is not None and recipe.nvfp4(): + if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype): + pytest.skip( + f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}" + ) + + with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe): + grouped_linear = GroupedLinear( + num_gemms, + config.hidden_size, + 4 * config.hidden_size, + bias=bias, + params_dtype=dtype, + parallel_mode=parallel_mode, + device="cuda", + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + delay_wgrad_compute=delay_wgrad_compute, + save_original_input=False, + ).eval() + sequential_linear = torch.nn.ModuleList( + [ + Linear( + config.hidden_size, + 4 * config.hidden_size, + bias=bias, + params_dtype=dtype, + parallel_mode=parallel_mode, + device="cuda", + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + ).eval() + for _ in range(num_gemms) + ] + ) + + # Share params + with torch.no_grad(): + for i in range(num_gemms): + sequential_linear[i].weight = Parameter(getattr(grouped_linear, f"weight{i}").clone()) + if bias: + sequential_linear[i].bias = Parameter(getattr(grouped_linear, f"bias{i}").clone()) + if fuse_wgrad_accumulation: + weight_i = getattr(grouped_linear, f"weight{i}") + weight_i.main_grad = torch.rand_like(weight_i, dtype=torch.float32) + sequential_linear[i].weight.main_grad = weight_i.main_grad.clone() + + outputs_ref = _test_grouped_linear_accuracy( + sequential_linear, + num_gemms, + bs, + dtype, + config, + recipe, + fp8, + fuse_wgrad_accumulation, + delay_wgrad_compute, + ) + outputs = _test_grouped_linear_accuracy( + grouped_linear, + num_gemms, + bs, + dtype, + config, + recipe, + fp8, + fuse_wgrad_accumulation, + delay_wgrad_compute, + ) + + for o, o_ref in zip(outputs, outputs_ref): + if use_cutlass: + torch.testing.assert_close(o, o_ref, rtol=1e-3, atol=1e-3) + else: + # cuBLAS implementation should be bit-wise match + torch.testing.assert_close(o, o_ref, rtol=0, atol=0) + + +@pytest.mark.skipif( + torch.cuda.get_device_capability() != (9, 0), + reason="Only enable CUTLASS grouped gemm on Hopper", +) +@pytest.mark.parametrize("dtype", param_types, ids=str) +@pytest.mark.parametrize("num_gemms", [3, 6]) +@pytest.mark.parametrize("bs", batch_sizes) +@pytest.mark.parametrize("model", ["126m"]) +@pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean) +@pytest.mark.parametrize("delay_wgrad_compute", all_boolean) +def test_grouped_linear_accuracy_cutlass( + dtype, + num_gemms, + bs, + model, + fuse_wgrad_accumulation, + delay_wgrad_compute, + monkeypatch, +): + monkeypatch.setenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "1") + test_grouped_linear_accuracy( + dtype, + num_gemms, + bs, + model, + None, + False, + fuse_wgrad_accumulation, + False, + delay_wgrad_compute, + None, + use_cutlass=True, + ) + + +@pytest.mark.parametrize("dtype", param_types, ids=str) +@pytest.mark.parametrize("num_gemms", [3]) +@pytest.mark.parametrize("bs", [1]) +@pytest.mark.parametrize("model", ["126m"]) +@pytest.mark.parametrize("recipe", fp8_recipes + [None], ids=recipe_id) +@pytest.mark.parametrize("fp8_model_params", [False]) +@pytest.mark.parametrize("fuse_wgrad_accumulation", [True]) +@pytest.mark.parametrize("bias", [False]) +@pytest.mark.parametrize("delay_wgrad_compute", [True]) +def test_grouped_linear_accuracy_save_original_input( + dtype, + num_gemms, + bs, + model, + recipe, + fp8_model_params, + fuse_wgrad_accumulation, + bias, + delay_wgrad_compute, + parallel_mode=None, +): + fp8 = recipe is not None + if fp8 and fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: + pytest.skip("FP8 parameters are not supported in debug mode.") + if fp8 and recipe.delayed(): + pytest.skip("DelayedScaling recipe is not supported with save_original_input") + if NVTE_TEST_NVINSPECT_ENABLED and delay_wgrad_compute: + pytest.skip("Delayed wgrad compute is not supported in debug mode.") + skip_unsupported_backward_override( + "grouped_linear", recipe, getattr(recipe, "backward_override", None) + ) + + config = model_configs[model] + if config.max_seqlen_q % 16 != 0 and fp8: + pytest.skip("FP8 requires sequence length to be divisible by 16.") + + if recipe is not None and recipe.nvfp4(): + if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype): + pytest.skip( + f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}" + ) + + with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe): + grouped_linear = GroupedLinear( + num_gemms, + config.hidden_size, + 4 * config.hidden_size, + bias=bias, + params_dtype=dtype, + parallel_mode=parallel_mode, + device="cuda", + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + delay_wgrad_compute=delay_wgrad_compute, + save_original_input=True, + ).eval() + sequential_linear = torch.nn.ModuleList( + [ + Linear( + config.hidden_size, + 4 * config.hidden_size, + bias=bias, + params_dtype=dtype, + parallel_mode=parallel_mode, + device="cuda", + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + ).eval() + for _ in range(num_gemms) + ] + ) + + # Share params + with torch.no_grad(): + for i in range(num_gemms): + sequential_linear[i].weight = Parameter(getattr(grouped_linear, f"weight{i}").clone()) + if bias: + sequential_linear[i].bias = Parameter(getattr(grouped_linear, f"bias{i}").clone()) + if fuse_wgrad_accumulation: + weight_i = getattr(grouped_linear, f"weight{i}") + weight_i.main_grad = torch.rand_like(weight_i, dtype=torch.float32) + sequential_linear[i].weight.main_grad = weight_i.main_grad.clone() + + outputs_ref = _test_grouped_linear_accuracy( + sequential_linear, + num_gemms, + bs, + dtype, + config, + recipe, + fp8, + fuse_wgrad_accumulation, + delay_wgrad_compute, + ) + outputs = _test_grouped_linear_accuracy( + grouped_linear, + num_gemms, + bs, + dtype, + config, + recipe, + fp8, + fuse_wgrad_accumulation, + delay_wgrad_compute, + ) + + # Should be bit-wise match + for i, (o, o_ref) in enumerate(zip(outputs, outputs_ref)): + torch.testing.assert_close(o, o_ref, rtol=0, atol=0) + + +@pytest.mark.parametrize("recipe", fp8_recipes + [None], ids=recipe_id) +def test_grouped_linear_accuracy_single_gemm(recipe): + """Split the tests to save CI time""" + test_grouped_linear_accuracy( + dtype=torch.float32, + num_gemms=1, + bs=2, + model="126m", + recipe=recipe, + fp8_model_params=True, + fuse_wgrad_accumulation=True, + bias=True, + delay_wgrad_compute=False, + ) + + +def _test_padding_grouped_linear_accuracy(block, num_gemms, bs, dtype, config, recipe, fp8=False): + + def _pad_tensor_for_fp8(hidden_states, tokens_per_expert): + align_size = get_align_size_for_quantization(recipe) + padded_tokens_per_expert = [ + (num_tokens + align_size - 1) // align_size * align_size + for num_tokens in tokens_per_expert + ] + hidden_states = torch.split(hidden_states, tokens_per_expert) + padded_hidden_states = [] + for hidden_state, actual_num_tokens, padded_num_tokens in zip( + hidden_states, tokens_per_expert, padded_tokens_per_expert + ): + padded_hidden_states.append(hidden_state) + if padded_num_tokens > actual_num_tokens: + pad_tensor = torch.zeros( + padded_num_tokens - actual_num_tokens, + hidden_state.shape[1], + dtype=hidden_state.dtype, + device=hidden_state.device, + ) + padded_hidden_states.append(pad_tensor) + padded_hidden_states = torch.cat(padded_hidden_states, dim=0) + return padded_hidden_states, padded_tokens_per_expert + + def _unpad_tensor_for_fp8(padded_hidden_states, actual_tokens_per_expert, tokens_per_expert): + inputmats = torch.split( + padded_hidden_states.view(-1, padded_hidden_states.shape[-1]), tokens_per_expert + ) + hidden_states = torch.cat( + [ + grad_output_mat[: actual_tokens_per_expert[i]] + for i, grad_output_mat in enumerate(inputmats) + ], + dim=0, + ) + + return hidden_states + + def _generate_random_numbers(n, total_sum): + if n <= 0: + return [] + + # reset seed + random.seed(seed) + + breaks = sorted(random.sample(range(1, total_sum), n - 1)) + random_numbers = ( + [breaks[0]] + + [breaks[i] - breaks[i - 1] for i in range(1, n - 1)] + + [total_sum - breaks[-1]] + ) + + return random_numbers + + reset_rng_states() + if fp8: + FP8GlobalStateManager.reset() + + inp_hidden_states = torch.randn( + (config.max_seqlen_q * bs, config.hidden_size), + dtype=dtype, + device="cuda", + requires_grad=True, + ) + inp_hidden_states.retain_grad() + + m_splits = _generate_random_numbers(num_gemms, config.max_seqlen_q * bs) + + with autocast(enabled=fp8, recipe=recipe): + if isinstance(block, TorchGroupedLinearWithPadding): + out = block(inp_hidden_states, m_splits) + else: + if fp8: + padded_inp_hidden_states, padding_m_splits = _pad_tensor_for_fp8( + inp_hidden_states, m_splits + ) + padded_inp_hidden_states = block(padded_inp_hidden_states, padding_m_splits) + out = _unpad_tensor_for_fp8(padded_inp_hidden_states, m_splits, padding_m_splits) + else: + out = block(inp_hidden_states, m_splits) + + loss = out.sum() + loss.backward() + + torch.cuda.synchronize() + outputs = [out, inp_hidden_states.grad] + for p in block.parameters(): + if p.requires_grad: + outputs.append(p.grad) + return outputs + + +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("num_gemms", [3, 6]) +@pytest.mark.parametrize("bs", batch_sizes) +@pytest.mark.parametrize("model", ["126m"]) +@pytest.mark.parametrize("fp8", [True]) +@pytest.mark.parametrize("recipe", fp8_recipes, ids=recipe_id) +@pytest.mark.parametrize("fp8_model_params", all_boolean) +def test_padding_grouped_linear_accuracy( + dtype, + num_gemms, + bs, + model, + fp8, + recipe, + fp8_model_params, + parallel_mode=None, +): + if fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: + pytest.skip("FP8 parameters are not supported in debug mode.") + skip_unsupported_backward_override( + "grouped_linear", recipe, getattr(recipe, "backward_override", None) + ) + + config = model_configs[model] + if config.max_seqlen_q % 16 != 0 and fp8: + pytest.skip("FP8 requires sequence length to be divisible by 16.") + + if recipe is not None and recipe.nvfp4(): + if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype): + pytest.skip( + f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}" + ) + + with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe): + grouped_linear = TorchGroupedLinearWithPadding( + num_gemms, + config.hidden_size, + 4 * config.hidden_size, + bias=False, + params_dtype=dtype, + parallel_mode=parallel_mode, + fp8=fp8, + ).eval() + + with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe): + ref_grouped_linear = GroupedLinear( + num_gemms, + config.hidden_size, + 4 * config.hidden_size, + bias=False, + params_dtype=dtype, + parallel_mode=parallel_mode, + device="cuda", + save_original_input=False, + ).eval() + + # Share params + with torch.no_grad(): + inner_grouped_linear = grouped_linear.linear_fn + for i in range(num_gemms): + setattr( + ref_grouped_linear, + f"weight{i}", + Parameter(getattr(inner_grouped_linear, f"weight{i}").clone()), + ) + + outputs = _test_padding_grouped_linear_accuracy( + grouped_linear, num_gemms, bs, dtype, config, recipe, fp8 + ) + outputs_ref = _test_padding_grouped_linear_accuracy( + ref_grouped_linear, num_gemms, bs, dtype, config, recipe, fp8 + ) + + # Should be bit-wise match + for i, (o, o_ref) in enumerate(zip(outputs, outputs_ref)): + torch.testing.assert_close(o, o_ref, rtol=0, atol=0) + + +@pytest.mark.parametrize("dtype", param_types) +@pytest.mark.parametrize("num_gemms", [3]) +@pytest.mark.parametrize("bs", [1]) +@pytest.mark.parametrize("model", ["126m"]) +@pytest.mark.parametrize("fp8", [True]) +@pytest.mark.parametrize("recipe", fp8_recipes, ids=recipe_id) +@pytest.mark.parametrize("fp8_model_params", [False]) +def test_padding_grouped_linear_accuracy_save_original_input( + dtype, + num_gemms, + bs, + model, + fp8, + recipe, + fp8_model_params, + parallel_mode=None, +): + if fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: + pytest.skip("FP8 parameters are not supported in debug mode.") + if fp8 and recipe.delayed(): + pytest.skip("DelayedScaling recipe is not supported with save_original_input") + skip_unsupported_backward_override( + "grouped_linear", recipe, getattr(recipe, "backward_override", None) + ) + + config = model_configs[model] + if config.max_seqlen_q % 16 != 0 and fp8: + pytest.skip("FP8 requires sequence length to be divisible by 16.") + + if recipe is not None and recipe.nvfp4(): + if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype): + pytest.skip( + f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}" + ) + + with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe): + grouped_linear = TorchGroupedLinearWithPadding( + num_gemms, + config.hidden_size, + 4 * config.hidden_size, + bias=False, + params_dtype=dtype, + parallel_mode=parallel_mode, + fp8=fp8, + ).eval() + + with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe): + ref_grouped_linear = GroupedLinear( + num_gemms, + config.hidden_size, + 4 * config.hidden_size, + bias=False, + params_dtype=dtype, + parallel_mode=parallel_mode, + device="cuda", + save_original_input=True, + ).eval() + + # Share params + with torch.no_grad(): + inner_grouped_linear = grouped_linear.linear_fn + for i in range(num_gemms): + setattr( + ref_grouped_linear, + f"weight{i}", + Parameter(getattr(inner_grouped_linear, f"weight{i}").clone()), + ) + + outputs = _test_padding_grouped_linear_accuracy( + grouped_linear, num_gemms, bs, dtype, config, recipe, fp8 + ) + outputs_ref = _test_padding_grouped_linear_accuracy( + ref_grouped_linear, num_gemms, bs, dtype, config, recipe, fp8 + ) + + # Should be bit-wise match + for i, (o, o_ref) in enumerate(zip(outputs, outputs_ref)): + torch.testing.assert_close(o, o_ref, rtol=0, atol=0) + + +@pytest.mark.parametrize( + "shape", + [ + (1, 127, 128, 512), + (8, 15, 128, 512), + (8, 1027, 128, 512), + (16, 10027, 128, 512), + ], +) +@pytest.mark.parametrize("dtype", param_types, ids=str) +@pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) +@pytest.mark.parametrize("accumulate", [False, True]) +@pytest.mark.parametrize("use_cutlass", use_cutlass_grouped_gemm) +def test_grouped_gemm(shape, dtype, layout, accumulate, use_cutlass, monkeypatch): + torch.manual_seed(0) + z, m, k, n = shape + + dist = torch.sort(torch.randint(0, m, (z - 1,))).values.tolist() + m_splits = torch.tensor(dist + [m]) - torch.tensor([0] + dist) + assert m_splits.sum() == m and len(m_splits) == z + m_splits = m_splits.tolist() + + if layout == "TN": + A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # weight + B = list(torch.split(torch.randn(m, k, dtype=dtype, device="cuda"), m_splits)) # input + out = [torch.randn(m, n, dtype=dtype, device="cuda")] # output + out_ref = [o.clone() for o in torch.split(out[0], m_splits)] + grad = False + single_output = True + elif layout == "NN": + A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # weight + B = list( + torch.split(torch.randn(m, n, dtype=dtype, device="cuda"), m_splits) + ) # grad_output + out = [torch.randn(m, k, dtype=dtype, device="cuda")] # dgrad + out_ref = [o.clone() for o in torch.split(out[0], m_splits)] + grad = True + single_output = True + else: # layout == "NT" + A = list(torch.split(torch.randn(m, k, dtype=dtype, device="cuda"), m_splits)) # input + B = list( + torch.split(torch.randn(m, n, dtype=dtype, device="cuda"), m_splits) + ) # grad_output + out = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # wgrad + out_ref = [o.clone() for o in out] + grad = True + single_output = False + + if use_cutlass: + monkeypatch.setenv("NVTE_USE_CUTLASS_GROUPED_GEMM", "1") + + for i in range(z): + general_gemm( + A[i], + B[i], + dtype, + grad=grad, + accumulate=accumulate, + layout=layout, + out=out_ref[i], + ) + if single_output: + out_ref = [torch.cat(out_ref)] + + general_grouped_gemm( + A, + B, + out, + [None] * z, + dtype, + m_splits=m_splits, + grad=grad, + accumulate=accumulate, + layout=layout, + single_output=single_output, + ) + + for o, o_ref in zip(out, out_ref): + if not use_cutlass: + # cublas implementation should be bit-wise match + torch.testing.assert_close(o, o_ref, rtol=0, atol=0) + else: + torch.testing.assert_close(o, o_ref, rtol=1.5e-2, atol=1.5e-2) + + +def _pack_grouped_tensor(grouped_tensor: GroupedTensor, tensors: List[torch.Tensor]) -> None: + data = grouped_tensor.rowwise_data + if data is None: + data = grouped_tensor.columnwise_data + if data is None: + raise ValueError("GroupedTensor has no data buffers to pack.") + offset = 0 + for tensor in tensors: + numel = tensor.numel() + data[offset : offset + numel].copy_(tensor.reshape(-1)) + offset += numel + + +def _make_grouped_tensor_from_splits( + m_sizes: List[int], + last_dim: int, + device: torch.device, + dtype: torch.dtype, +) -> GroupedTensor: + first_dims = torch.tensor(m_sizes, device=device, dtype=torch.int64) + return GroupedTensor.make_grouped_tensor( + num_tensors=len(m_sizes), + first_dims=first_dims, + last_dims=None, + logical_first_dim=sum(m_sizes), + logical_last_dim=last_dim, + quantizer=None, + device=device, + dtype=dtype, + ) + + +def _make_grouped_tensor_uniform( + num_tensors: int, + first_dim: int, + last_dim: int, + device: torch.device, + dtype: torch.dtype, +) -> GroupedTensor: + return GroupedTensor.make_grouped_tensor( + num_tensors=num_tensors, + first_dims=None, + last_dims=None, + logical_first_dim=num_tensors * first_dim, + logical_last_dim=last_dim, + quantizer=None, + device=device, + dtype=dtype, + ) + + +def _apply_grouped_bias_ref( + base_outs: List[torch.Tensor], + bias: Optional[List[torch.Tensor]], + bias_scale: Optional[torch.Tensor], + m_sizes: List[int], + dtype: torch.dtype, +) -> List[torch.Tensor]: + """Reference: add (optionally per-row scaled) bias to each group's output, cast to ``dtype``.""" + if bias is None: + return list(base_outs) + if bias_scale is None: + return [(o.float() + b.float()).to(dtype) for o, b in zip(base_outs, bias)] + out = [] + offset = 0 + for i, ms in enumerate(m_sizes): + s = bias_scale[offset : offset + ms].unsqueeze(-1) + out.append((base_outs[i].float() + bias[i].float() * s).to(dtype)) + offset += ms + return out + + +@pytest.mark.parametrize( + "z, m, n, k", + [ + (4, 256, 256, 256), + (4, 512, 256, 512), + (4, 512, 512, 256), + (8, 512, 256, 512), + ], +) +@pytest.mark.parametrize("case", ["no_discrete", "discrete_in", "discrete_out"]) +@pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) +@pytest.mark.parametrize("accumulate", [False, True]) +@pytest.mark.parametrize("use_bias_scale", [False, True]) +def test_grouped_gemm_grouped_tensor(z, m, n, k, case, layout, accumulate, use_bias_scale) -> None: + if torch.cuda.get_device_capability() < (9, 0): + pytest.skip("Grouped GEMM requires Hopper (SM90) or newer.") + if torch.cuda.get_device_capability() < (10, 0): + if tex.get_cublasLt_version() < 130400: + pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") + if tex.get_cublasLt_version() < 130300: + pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") + if not is_bf16_available(): + pytest.skip("bfloat16 is required for grouped GEMM test.") + + torch.manual_seed(0) + + dtype = torch.bfloat16 + + split_points = torch.randperm(m - 1)[: z - 1] + 1 + split_points = torch.sort(split_points).values.tolist() + m_sizes = [split_points[0]] + m_sizes += [b - a for a, b in zip(split_points[:-1], split_points[1:])] + m_sizes.append(m - split_points[-1]) + assert sum(m_sizes) == m and len(m_sizes) == z + + if layout == "NT": + A = [torch.randn(ms, k, dtype=dtype, device="cuda") for ms in m_sizes] # input + B = [torch.randn(ms, n, dtype=dtype, device="cuda") for ms in m_sizes] # grad_output + out = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # wgrad + out_ref = [torch.matmul(B[i].transpose(0, 1).float(), A[i].float()) for i in range(z)] + else: + A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # weight + B = [ + torch.randn(ms, k if layout == "TN" else n, dtype=dtype, device="cuda") + for ms in m_sizes + ] # TN --> input, NN --> grad_output + out = [ + torch.randn(ms, n if layout == "TN" else k, dtype=dtype, device="cuda") + for ms in m_sizes + ] # TN --> output, NN --> dgrad + if layout == "NN": + out_ref = [torch.matmul(B[i].float(), A[i].float()) for i in range(z)] + else: # layout == "TN" + out_ref = [torch.matmul(B[i].float(), A[i].transpose(0, 1).float()) for i in range(z)] + + if accumulate: + out_ref = [out[i].float() + o for i, o in enumerate(out_ref)] + + # Bias is applied after GEMM (broadcasted along rows) + # Match kernel behavior: GEMM output is already in output dtype when bias is added. + out_ref_no_bias = [o.to(dtype) for o in out_ref] + if layout == "TN": + bias_last_dim = n + else: # layout == "NT" or "NN" + bias_last_dim = k + bias = ( + [torch.randn(1, bias_last_dim, dtype=dtype, device="cuda") for _ in range(z)] + if case != "discrete_out" + else None + ) + bias_scale = None + if use_bias_scale and bias is not None and layout != "NT": + bias_scale = torch.randn(m, device="cuda", dtype=torch.float32) + # Bias add in grouped kernel accumulates in FP32 for BF16/FP16. + out_ref = _apply_grouped_bias_ref(out_ref_no_bias, bias, bias_scale, m_sizes, dtype) + # Create grouped tensors based on case + device = A[0].device + grouped_A = A + grouped_out = out + grouped_out_bias = [o.clone() for o in out] + grouped_out_no_bias = [o.clone() for o in out] + grouped_bias = None + if layout == "TN": + grouped_A = ( + _make_grouped_tensor_uniform(z, n, k, device, dtype) if case != "discrete_in" else A + ) # weight + grouped_B = _make_grouped_tensor_from_splits(m_sizes, k, device, dtype) # input + if case != "discrete_out": + grouped_out = _make_grouped_tensor_from_splits(m_sizes, n, device, dtype) # output + grouped_out_bias = _make_grouped_tensor_from_splits(m_sizes, n, device, dtype) + grouped_out_no_bias = _make_grouped_tensor_from_splits(m_sizes, n, device, dtype) + elif layout == "NN": + grouped_A = ( + _make_grouped_tensor_uniform(z, n, k, device, dtype) if case != "discrete_in" else A + ) # weight + grouped_B = _make_grouped_tensor_from_splits(m_sizes, n, device, dtype) # grad_output + if case != "discrete_out": + grouped_out = _make_grouped_tensor_from_splits(m_sizes, k, device, dtype) + grouped_out_bias = _make_grouped_tensor_from_splits(m_sizes, k, device, dtype) + grouped_out_no_bias = _make_grouped_tensor_from_splits(m_sizes, k, device, dtype) + else: # layout == "NT" + grouped_A = ( + _make_grouped_tensor_from_splits(m_sizes, k, device, dtype) + if case != "discrete_in" + else A + ) # input + grouped_B = _make_grouped_tensor_from_splits(m_sizes, n, device, dtype) # grad_output + if case != "discrete_out": + grouped_out = _make_grouped_tensor_uniform(z, n, k, device, dtype) # wgrad + grouped_out_bias = _make_grouped_tensor_uniform(z, n, k, device, dtype) + grouped_out_no_bias = _make_grouped_tensor_uniform(z, n, k, device, dtype) + _pack_grouped_tensor(grouped_B, B) + if case != "discrete_out": + _pack_grouped_tensor(grouped_out, out) + _pack_grouped_tensor(grouped_out_bias, out) + _pack_grouped_tensor(grouped_out_no_bias, out) + if case != "discrete_in": + _pack_grouped_tensor(grouped_A, A) + + if bias is not None: + grouped_bias = _make_grouped_tensor_uniform(z, 1, bias_last_dim, device, dtype) + _pack_grouped_tensor(grouped_bias, bias) + + general_grouped_gemm_for_grouped_tensor( + grouped_A, + grouped_B, + grouped_out_no_bias, + layout=layout, + accumulate=accumulate, + bias=None, + ) + general_grouped_gemm_for_grouped_tensor( + grouped_A, + grouped_B, + grouped_out_bias, + layout=layout, + accumulate=accumulate, + bias=grouped_bias, + bias_scale=bias_scale, + ) + out_grouped_no_bias = ( + grouped_out_no_bias + if isinstance(grouped_out_no_bias, list) + else grouped_out_no_bias.split_into_quantized_tensors() + ) + out_grouped_bias = ( + grouped_out_bias + if isinstance(grouped_out_bias, list) + else grouped_out_bias.split_into_quantized_tensors() + ) + + out_grouped_manual_bias = _apply_grouped_bias_ref( + out_grouped_no_bias, bias, bias_scale, m_sizes, dtype + ) + tols = dtype_tols(dtype) + for o, o_ref in zip(out_grouped_no_bias, out_ref_no_bias): + torch.testing.assert_close(o, o_ref, **tols) + if bias is not None: + for o, o_ref in zip(out_grouped_bias, out_grouped_manual_bias): + torch.testing.assert_close(o, o_ref, **tols) + + +@pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) +@pytest.mark.parametrize("accumulate", [False, True]) +@pytest.mark.parametrize("quant_type", ["bf16", "mxfp8"]) +def test_grouped_gemm_grouped_tensor_zero_work(layout, accumulate, quant_type) -> None: + """Grouped GEMM with all-zero split sizes (zero total work). + + For wgrad (NT layout) the output should be zero when not accumulating, + or unchanged when accumulating with beta=1. + """ + if torch.cuda.get_device_capability() < (10, 0): + pytest.skip("Grouped GEMM requires Blackwell (SM100) or newer.") + if not is_bf16_available(): + pytest.skip("bfloat16 is required for grouped GEMM test.") + if quant_type == "mxfp8" and not mxfp8_available: + pytest.skip(reason_for_no_mxfp8) + + z = 4 + k, n = 256, 256 + dtype = torch.bfloat16 + device = torch.device("cuda") + use_mxfp8 = quant_type == "mxfp8" + + transa = layout[0] == "T" + transb = layout[1] == "T" + zero_first_dims = torch.zeros(z, dtype=torch.int64, device=device) + + def _make_zero_tokens_grouped_tensor(logical_last_dim, is_a): + """Create a GroupedTensor with non-zero logical_shape but zero first_dims.""" + buf = torch.randn(0, logical_last_dim, dtype=dtype, device=device) + if use_mxfp8: + if is_a: + rowwise, columnwise = transa, not transa + else: + rowwise, columnwise = not transb, transb + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + ) + quantizer.optimize_for_gemm = True + return tex.group_quantize(buf, quantizer, z, zero_first_dims) + return GroupedTensor.make_grouped_tensor( + num_tensors=z, + first_dims=zero_first_dims, + last_dims=None, + logical_first_dim=k, + logical_last_dim=logical_last_dim, + quantizer=None, + device=device, + dtype=dtype, + ) + + if layout in ("TN", "NN"): + weight_tensors = [torch.randn(n, k, dtype=dtype, device=device) for _ in range(z)] + if use_mxfp8: + grouped_A = _make_grouped_tensor_quantized_mxfp8( + weight_tensors, + rowwise=transa, + columnwise=not transa, + device=device, + ) + else: + grouped_A = _make_grouped_tensor_uniform(z, n, k, device, dtype) + _pack_grouped_tensor(grouped_A, weight_tensors) + else: # NT + grouped_A = _make_zero_tokens_grouped_tensor(k, is_a=True) + + b_last_dim = k if layout == "TN" else n + grouped_B = _make_zero_tokens_grouped_tensor(b_last_dim, is_a=False) + + if layout == "NT": + out = [torch.randn(n, k, dtype=dtype, device=device) for _ in range(z)] + grouped_out = _make_grouped_tensor_uniform(z, n, k, device, dtype) + _pack_grouped_tensor(grouped_out, out) + else: + out = [torch.zeros(0, dtype=dtype, device=device) for _ in range(z)] + out_last_dim = n if layout == "TN" else k + grouped_out = GroupedTensor.make_grouped_tensor( + num_tensors=z, + first_dims=zero_first_dims, + last_dims=None, + logical_first_dim=k, + logical_last_dim=out_last_dim, + quantizer=None, + device=device, + dtype=dtype, + ) + + out_before = [o.clone() for o in out] + + general_grouped_gemm_for_grouped_tensor( + grouped_A, + grouped_B, + grouped_out, + layout=layout, + accumulate=accumulate, + ) + + out_result = ( + grouped_out if isinstance(grouped_out, list) else grouped_out.split_into_quantized_tensors() + ) + for i in range(z): + if out_result[i].numel() == 0: + continue + if accumulate: + torch.testing.assert_close(out_result[i], out_before[i]) + else: + torch.testing.assert_close(out_result[i], torch.zeros_like(out_result[i])) + + +def _make_grouped_tensor_quantized_mxfp8( + tensors: List[torch.Tensor], + *, + rowwise: bool, + columnwise: bool, + device: torch.device, + is_weight: bool = False, +) -> GroupedTensor: + """Create a quantized MXFP8 GroupedTensor from a list of per-expert tensors. + + For weights (uniform per-expert shape), we generally won't keep it swizzled since we + might need for future dequantize operations. Swizzling is done internally within + general_grouped_gemm_for_grouped_tensor call. + + For non-weight tensors (inputs / grad_outputs), we still pass + ``first_dims`` and keep ``optimize_for_gemm=True``; so the kernel must emit the + already-swizzled layout up front. + """ + if not tensors: + raise ValueError("Expected non-empty tensor list for grouped quantization.") + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + ) + quantizer.optimize_for_gemm = not is_weight + grouped_input = torch.cat(tensors, dim=0) + if is_weight: + first_dims = None + else: + first_dims = torch.tensor([t.shape[0] for t in tensors], dtype=torch.int64, device=device) + return tex.group_quantize(grouped_input, quantizer, len(tensors), first_dims) + + +def _per_tensor_quantize_mxfp8( + tensors: List[torch.Tensor], + *, + rowwise: bool, + columnwise: bool, +) -> List: + """Quantize each tensor individually with MXFP8. + Used to build reference discrete inputs for grouped GEMM. + """ + quantizer = MXFP8Quantizer( + fp8_dtype=tex.DType.kFloat8E4M3, + rowwise=rowwise, + columnwise=columnwise, + ) + return [quantizer(t) for t in tensors] + + +@pytest.mark.parametrize( + "shape", + [ + (1, 128, 128, 512), + (8, 1024, 128, 512), + (16, 4096, 128, 512), + (2, 256, 2880, 2880), + ], +) +@pytest.mark.parametrize("accumulate", [False, True]) +@pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) +@pytest.mark.parametrize("case", ["no_discrete", "discrete_in", "discrete_out"]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) +def test_grouped_gemm_grouped_tensor_mxfp8( + shape, accumulate, layout: str, case: str, dtype: torch.dtype +) -> None: + if tex.get_cublasLt_version() < 130300: + pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") + if torch.cuda.get_device_capability() < (10, 0): + pytest.skip("Grouped GEMM requires Blackwell (SM100) or newer.") + if dtype == torch.bfloat16 and not is_bf16_available(): + pytest.skip("bfloat16 is required for grouped GEMM test.") + + torch.manual_seed(0) + z, m, k, n = shape + m_sizes = [m // z] * z + + if layout == "TN": + A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # weight + B = [torch.randn(ms, k, dtype=dtype, device="cuda") for ms in m_sizes] # input + out = [torch.randn(ms, n, dtype=dtype, device="cuda") for ms in m_sizes] # output + grad = False + elif layout == "NN": + A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # weight + B = [torch.randn(ms, n, dtype=dtype, device="cuda") for ms in m_sizes] # grad_output + out = [torch.randn(ms, k, dtype=dtype, device="cuda") for ms in m_sizes] # dgrad + grad = True + else: # layout == "NT" + A = [torch.randn(ms, k, dtype=dtype, device="cuda") for ms in m_sizes] # input + B = [torch.randn(ms, n, dtype=dtype, device="cuda") for ms in m_sizes] # grad_output + out = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # wgrad + grad = True + + out_ref = [o.clone() for o in out] + + transa = layout[0] == "T" + transb = layout[1] == "T" + a_is_weight = all(t.shape == A[0].shape for t in A) + a_rowwise, a_columnwise = transa, not transa + b_rowwise, b_columnwise = not transb, transb + grouped_A = _make_grouped_tensor_quantized_mxfp8( + A, + rowwise=a_rowwise, + columnwise=a_columnwise, + device="cuda", + is_weight=a_is_weight, + ) + grouped_B = _make_grouped_tensor_quantized_mxfp8( + B, rowwise=b_rowwise, columnwise=b_columnwise, device="cuda" + ) + A_fp8 = _per_tensor_quantize_mxfp8(A, rowwise=a_rowwise, columnwise=a_columnwise) + B_fp8 = _per_tensor_quantize_mxfp8(B, rowwise=b_rowwise, columnwise=b_columnwise) + + general_grouped_gemm( + A_fp8, + B_fp8, + out_ref, + [None] * z, + dtype, + m_splits=m_sizes, + grad=grad, + accumulate=accumulate, + layout=layout, + single_output=False, + ) + + device = A[0].device + + grouped_out = None + if case != "discrete_out": + if layout == "TN": + grouped_out = _make_grouped_tensor_from_splits(m_sizes, n, device, dtype) + elif layout == "NN": + grouped_out = _make_grouped_tensor_from_splits(m_sizes, k, device, dtype) + else: # layout == "NT" + grouped_out = _make_grouped_tensor_uniform(z, n, k, device, dtype) + _pack_grouped_tensor(grouped_out, out) + + grouped_out_input = out if case == "discrete_out" else grouped_out + grouped_A_input = A_fp8 if case == "discrete_in" else grouped_A + general_grouped_gemm_for_grouped_tensor( + grouped_A_input, + grouped_B, + grouped_out_input, + layout=layout, + accumulate=accumulate, + ) + + out_grouped = out if case == "discrete_out" else grouped_out.split_into_quantized_tensors() + tols = dict(rtol=0.125, atol=0.0675) # mxfp8 tolerance + + for o, o_ref in zip(out_grouped, out_ref): + torch.testing.assert_close(o, o_ref, **tols) + + +@pytest.mark.parametrize( + "shape", + [ + (1, 128, 128, 512), + (8, 1024, 128, 512), + (16, 4096, 128, 512), + ], +) +@pytest.mark.parametrize("accumulate", [False, True]) +def test_fp8_grouped_gemm(shape, accumulate): + if not fp8_available: + pytest.skip(reason_for_no_fp8) + + z, m, k, n = shape + m_splits = [m // z] * z + + dtype = torch.bfloat16 + A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # weight + B = torch.split(torch.randn(m, k, dtype=dtype, device="cuda"), m_splits) # input + out = torch.split(torch.randn(m, n, dtype=dtype, device="cuda"), m_splits) # output + out_ref = [o.clone() for o in out] + + # fp8 should be robust enough to this fake scale + scale = 1 + torch.rand(1, dtype=torch.float32, device="cuda").squeeze() + amax = torch.zeros(1, 1, dtype=torch.float32, device="cuda") + + a_quantizers = [ + Float8Quantizer( + scale.clone(), + amax.clone(), + tex.DType.kFloat8E4M3, + ) + for _ in range(z) + ] + b_quantizers = [ + Float8Quantizer( + scale.clone(), + amax.clone(), + tex.DType.kFloat8E4M3, + ) + for _ in range(z) + ] + + A_fp8 = [] + B_fp8 = [] + + for i in range(z): + A_fp8.append(a_quantizers[i](A[i])) + B_fp8.append(b_quantizers[i](B[i])) + + # baseline + for i in range(z): + general_gemm( + A_fp8[i], + B_fp8[i], + dtype, + out=out_ref[i], + accumulate=accumulate, + ) + general_grouped_gemm( + A_fp8, + B_fp8, + out, + [None] * z, + dtype, + m_splits=m_splits, + accumulate=accumulate, + ) + + # should be bit-wise match + for o, o_ref in zip(out, out_ref): + torch.testing.assert_close(o, o_ref, rtol=0, atol=0) + + +_FUSED_GROUPED_GEMM_ENV = "NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM" +_ALL_BOOLEAN = all_boolean +_mxfp8_available, _reason_for_no_mxfp8 = mxfp8_available, reason_for_no_mxfp8 + + +@pytest.fixture(autouse=True) +def _reset_fp8_state(monkeypatch): + monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "0") + yield + FP8GlobalStateManager.reset() + monkeypatch.delenv(_FUSED_GROUPED_GEMM_ENV, raising=False) + + +def _clone_outputs(outputs): + return [None if out is None else out.detach().clone() for out in outputs] + + +def _run_grouped_linear_path( + *, + enable_grouped_tensor_path: bool, + fp8_recipe, + bias: bool, + fp8_model_params: bool, + delay_wgrad_compute: bool, + x_base: torch.Tensor, + dy: torch.Tensor, + weights, + biases, + m_splits, + monkeypatch, +): + FP8GlobalStateManager.reset() + monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1" if enable_grouped_tensor_path else "0") + + dtype = x_base.dtype + num_gemms = len(m_splits) + in_features = weights[0].size(1) + out_features = weights[0].size(0) + use_fp8 = fp8_recipe is not None + + x = x_base.detach().clone().requires_grad_(True) + with quantized_model_init(enabled=fp8_model_params, recipe=fp8_recipe): + grouped_linear = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=bias, + params_dtype=dtype, + device="cuda", + delay_wgrad_compute=delay_wgrad_compute, + ) + with torch.no_grad(): + for i in range(num_gemms): + getattr(grouped_linear, f"weight{i}").copy_(weights[i]) + if bias: + getattr(grouped_linear, f"bias{i}").copy_(biases[i]) + + # The fused path is the graph-safe path and accepts a CUDA tensor for split metadata. + # The legacy path still expects Python split sections in several places. + m_splits_arg = ( + torch.tensor(m_splits, dtype=torch.int64, device="cuda") + if enable_grouped_tensor_path + else m_splits + ) + with autocast(enabled=use_fp8, recipe=fp8_recipe): + y = grouped_linear(x, m_splits_arg) + y.backward(dy) + if delay_wgrad_compute: + grouped_linear.backward_dw() + + outputs = [y, x.grad] + for i in range(num_gemms): + outputs.append(getattr(grouped_linear, f"weight{i}").grad) + if bias: + outputs.append(getattr(grouped_linear, f"bias{i}").grad) + return _clone_outputs(outputs) + + +@pytest.mark.parametrize( + "fp8_recipe", + [ + None, + pytest.param( + recipe.MXFP8BlockScaling(), + marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), + ), + ], + ids=["bf16", "mxfp8"], +) +@pytest.mark.parametrize("bias", _ALL_BOOLEAN) +@pytest.mark.parametrize("fp8_model_params", _ALL_BOOLEAN) +@pytest.mark.parametrize("delay_wgrad_compute", _ALL_BOOLEAN) +def test_grouped_linear_grouped_tensor_path_matches_legacy( + fp8_recipe, bias, fp8_model_params, delay_wgrad_compute, monkeypatch +): + if torch.cuda.get_device_capability() < (10, 0): + pytest.skip("GroupedTensor grouped GEMM path requires SM100+") + + use_fp8 = fp8_recipe is not None + if fp8_model_params and not use_fp8: + pytest.skip("fp8_model_params requires FP8") + + dtype = torch.bfloat16 + num_gemms = 3 + in_features = 64 + out_features = 64 + m_splits = [128, 256, 384] + total_tokens = sum(m_splits) + + torch.manual_seed(1234) + x_base = (0.1 * torch.randn(total_tokens, in_features, device="cuda")).to(dtype) + dy = (0.1 * torch.randn(total_tokens, out_features, device="cuda")).to(dtype) + weights = [ + (0.1 * torch.randn(out_features, in_features, device="cuda")).to(dtype) + for _ in range(num_gemms) + ] + biases = None + if bias: + biases = [ + (0.1 * torch.randn(out_features, device="cuda")).to(dtype) for _ in range(num_gemms) + ] + + outputs_legacy = _run_grouped_linear_path( + enable_grouped_tensor_path=False, + fp8_recipe=fp8_recipe, + bias=bias, + fp8_model_params=fp8_model_params, + delay_wgrad_compute=delay_wgrad_compute, + x_base=x_base, + dy=dy, + weights=weights, + biases=biases, + m_splits=m_splits, + monkeypatch=monkeypatch, + ) + outputs_grouped_tensor = _run_grouped_linear_path( + enable_grouped_tensor_path=True, + fp8_recipe=fp8_recipe, + bias=bias, + fp8_model_params=fp8_model_params, + delay_wgrad_compute=delay_wgrad_compute, + x_base=x_base, + dy=dy, + weights=weights, + biases=biases, + m_splits=m_splits, + monkeypatch=monkeypatch, + ) + + tols = dict(rtol=1e-2, atol=5e-3) + if use_fp8: + tols = dict(rtol=0.05, atol=0.05) + for grouped_tensor_out, legacy_out in zip(outputs_grouped_tensor, outputs_legacy): + assert grouped_tensor_out is not None + assert legacy_out is not None + torch.testing.assert_close(grouped_tensor_out.float(), legacy_out.float(), **tols) + + +def test_grouped_linear_grouped_tensor_path_single_grouped_bias_delay_wgrad(monkeypatch): + if torch.cuda.get_device_capability() < (10, 0): + pytest.skip("GroupedTensor grouped GEMM path requires SM100+") + + monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1") + + dtype = torch.bfloat16 + num_gemms = 3 + in_features = 64 + out_features = 64 + total_tokens = 64 + 96 + 128 + m_splits = torch.tensor([64, 96, 128], dtype=torch.int64, device="cuda") + x = torch.randn(total_tokens, in_features, dtype=dtype, device="cuda").requires_grad_() + dy = torch.randn(x.size(0), out_features, dtype=dtype, device="cuda") + + grouped_linear = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=True, + params_dtype=dtype, + device="cuda", + delay_wgrad_compute=True, + single_grouped_bias=True, + ) + + y = grouped_linear(x, m_splits) + y.backward(dy) + grouped_linear.backward_dw() + + +@pytest.mark.parametrize( + "fp8_recipe", + [ + None, + pytest.param( + recipe.MXFP8BlockScaling(), + marks=pytest.mark.skipif(not _mxfp8_available, reason=_reason_for_no_mxfp8), + ), + ], + ids=["bf16", "mxfp8"], +) +@pytest.mark.parametrize("bias", _ALL_BOOLEAN) +def test_grouped_linear_fused_path_cuda_graph_safe(fp8_recipe, bias, monkeypatch): + """Fused GroupedTensor GEMM path should be CUDA graph capturable.""" + if torch.cuda.get_device_capability() < (10, 0): + pytest.skip("GroupedTensor grouped GEMM path requires SM100+") + + monkeypatch.setenv(_FUSED_GROUPED_GEMM_ENV, "1") + FP8GlobalStateManager.reset() + + use_fp8 = fp8_recipe is not None + dtype = torch.bfloat16 + device = "cuda" + num_gemms = 3 + in_features = 128 + out_features = 128 + split_sizes = [128, 256, 384] + total_tokens = sum(split_sizes) + static_m_splits = torch.tensor(split_sizes, dtype=torch.int64, device=device) + + grouped_linear = GroupedLinear( + num_gemms, + in_features, + out_features, + bias=bias, + params_dtype=dtype, + device=device, + ) + + static_x = torch.randn(total_tokens, in_features, dtype=dtype, device=device) + static_x.requires_grad_(True) + static_dy = torch.randn(total_tokens, out_features, dtype=dtype, device=device) + static_out_buf = torch.empty(total_tokens, out_features, dtype=dtype, device=device) + + def _zero_grads(): + if static_x.grad is not None: + static_x.grad.zero_() + for param in grouped_linear.parameters(): + if param.grad is None: + param.grad = torch.zeros_like(param) + else: + param.grad.zero_() + + def _clone_param_grads(): + return [param.grad.detach().clone() for param in grouped_linear.parameters()] + + def _train_step(x, dy, out_buf, *, use_graphed): + with autocast(enabled=use_fp8, recipe=fp8_recipe): + out = ( + graphed_grouped_linear(x, static_m_splits) + if use_graphed + else grouped_linear(x, static_m_splits) + ) + out.backward(dy) + out_buf.copy_(out) + return out_buf + + graphed_grouped_linear = te.make_graphed_callables( + grouped_linear, + (static_x, static_m_splits), + num_warmup_iters=3, + enabled=use_fp8, + recipe=fp8_recipe, + ) + + fresh_x = torch.randn_like(static_x) + fresh_dy = torch.randn_like(static_dy) + with torch.no_grad(): + static_x.copy_(fresh_x) + static_dy.copy_(fresh_dy) + + _zero_grads() + graph_out = ( + _train_step( + static_x, + static_dy, + static_out_buf, + use_graphed=True, + ) + .detach() + .clone() + ) + torch.cuda.synchronize() + graph_dx = static_x.grad.detach().clone() + graph_param_grads = _clone_param_grads() + + _zero_grads() + expected_x = fresh_x.detach().clone().requires_grad_(True) + expected_dy = fresh_dy.detach().clone() + with autocast(enabled=use_fp8, recipe=fp8_recipe): + expected_out = grouped_linear(expected_x, static_m_splits) + expected_out.backward(expected_dy) + + tols = dict(rtol=1e-2, atol=5e-3) + if use_fp8: + tols = dict(rtol=0.05, atol=0.05) + torch.testing.assert_close(graph_out.float(), expected_out.float(), **tols) + torch.testing.assert_close(graph_dx.float(), expected_x.grad.float(), **tols) + for graph_grad, param in zip(graph_param_grads, grouped_linear.parameters()): + assert param.grad is not None + torch.testing.assert_close(graph_grad.float(), param.grad.float(), **tols) diff --git a/tests/pytorch/test_numerics.py b/tests/pytorch/test_numerics.py index 368c95e27..e087f1e1c 100644 --- a/tests/pytorch/test_numerics.py +++ b/tests/pytorch/test_numerics.py @@ -6,7 +6,6 @@ import os from typing import Dict, List, Tuple, Optional import pytest -import random import torch import torch.nn as nn @@ -14,7 +13,6 @@ from transformer_engine.pytorch.quantization import ( FP8GlobalStateManager, - get_align_size_for_quantization, ) from transformer_engine.pytorch.utils import ( init_method_normal, @@ -28,13 +26,10 @@ LayerNormLinear, LayerNormMLP, Linear, - GroupedLinear, MultiheadAttention, RMSNorm, TransformerLayer, LayerNorm, - Fp8Padding, - Fp8Unpadding, Float8Quantizer, Float8CurrentScalingQuantizer, MXFP8Quantizer, @@ -46,17 +41,11 @@ is_nvfp4_available, ) from transformer_engine.pytorch import checkpoint as te_checkpoint -from transformer_engine.pytorch.cpp_extensions import ( - general_gemm, - general_grouped_gemm, - general_grouped_gemm_for_grouped_tensor, -) -from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor +from transformer_engine.pytorch.cpp_extensions import general_gemm from transformer_engine.common import recipe import transformer_engine_torch as tex from utils import ModelConfig, recipe_id, reset_rng_states, skip_unsupported_backward_override - # Only run FP8 tests on supported devices. fp8_available, reason_for_no_fp8 = is_fp8_available(return_reason=True) mxfp8_available, reason_for_no_mxfp8 = is_mxfp8_available(return_reason=True) @@ -200,11 +189,6 @@ def get_nvfp4_inp_supported_dtypes(recipe: recipe.Recipe, dtype: torch.dtype) -> fp8_recipes.append(nvfp4_4over6()) fp8_recipes.append(nvfp4_row_scaled()) -use_cutlass_grouped_gemm = [False] -# Only enable cutlass grouped gemm on Hopper -if torch.cuda.get_device_capability() == (9, 0): - use_cutlass_grouped_gemm.append(True) - def get_causal_attn_mask(sq: int) -> torch.Tensor: return torch.triu(torch.ones(sq, sq, device="cuda"), diagonal=1).bool() @@ -476,40 +460,6 @@ def forward(self, input: torch.Tensor) -> torch.Tensor: return (input > 0) * input * input -class TorchGroupedLinearWithPadding(nn.Module): - - def __init__( - self, num_gemms, in_features, out_features, bias, params_dtype, parallel_mode, fp8 - ) -> None: - super().__init__() - - self.padding = Fp8Padding(num_gemms) - self.linear_fn = GroupedLinear( - num_gemms, - in_features, - out_features, - bias=bias, - params_dtype=params_dtype, - parallel_mode=parallel_mode, - device="cuda", - ) - self.unpadding = Fp8Unpadding(num_gemms) - - self.fp8 = fp8 - - def forward(self, inp: torch.Tensor, m_splits: List[int]) -> torch.Tensor: - if self.fp8: - orig_m_splits = m_splits - inp, m_splits = self.padding(inp, m_splits) - - out = self.linear_fn(inp, m_splits) - - if self.fp8: - out = self.unpadding(out, orig_m_splits) - - return out - - _supported_act = { "gelu": nn.GELU(approximate="tanh"), "geglu": nn.GELU(approximate="tanh"), @@ -1859,596 +1809,6 @@ def test_layernorm_mlp_accuracy_checkpoint( torch.testing.assert_close(o, o_ref, rtol=0, atol=0) -def _test_grouped_linear_accuracy( - block, - num_gemms, - bs, - dtype, - config, - recipe, - fp8, - fuse_wgrad_accumulation, - delay_wgrad_compute=False, -): - reset_rng_states() - if fp8: - FP8GlobalStateManager.reset() - - inp_hidden_states = torch.randn( - (config.max_seqlen_q, bs, config.hidden_size), - dtype=dtype, - device="cuda", - requires_grad=True, - ) - inp_hidden_states.retain_grad() - - if num_gemms > 1: - split_size = 1 - if fp8: - split_size = get_align_size_for_quantization(recipe) - m = config.max_seqlen_q // split_size - dist = torch.sort(torch.randint(0, m, (num_gemms - 2,))).values.tolist() - dist.append(dist[-1]) # Manually add a zero - m_splits = torch.tensor(dist + [m]) - torch.tensor([0] + dist) - m_splits = m_splits * split_size - assert m_splits.sum() == config.max_seqlen_q and len(m_splits) == num_gemms - else: - m_splits = torch.tensor([config.max_seqlen_q]) - - with autocast(enabled=fp8, recipe=recipe): - if isinstance(block, GroupedLinear): - m_splits = m_splits * bs - out = block(inp_hidden_states, m_splits.tolist()) - else: - out = torch.cat( - [ - block[i](inp) - for i, inp in enumerate(torch.split(inp_hidden_states, m_splits.tolist())) - ] - ) - loss = out.sum() - loss.backward() - if delay_wgrad_compute: - if isinstance(block, GroupedLinear): - block.backward_dw() - else: - for i in range(num_gemms): - block[i].backward_dw() - - torch.cuda.synchronize() - outputs = [out, inp_hidden_states.grad] - for p in block.parameters(): - if p.requires_grad: - if getattr(p, "main_grad", None) is not None: - outputs.append(p.main_grad) - assert p.grad is None # grad should be None if fuse_wgrad_accumulation is True - else: - outputs.append(p.grad) - return outputs - - -@pytest.mark.parametrize("dtype", param_types, ids=str) -@pytest.mark.parametrize("num_gemms", [3, 6]) -@pytest.mark.parametrize("bs", batch_sizes) -@pytest.mark.parametrize("model", ["126m"]) -@pytest.mark.parametrize("recipe", fp8_recipes + [None], ids=recipe_id) -@pytest.mark.parametrize("fp8_model_params", all_boolean) -@pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean) -@pytest.mark.parametrize("bias", all_boolean) -@pytest.mark.parametrize("delay_wgrad_compute", all_boolean) -def test_grouped_linear_accuracy( - dtype, - num_gemms, - bs, - model, - recipe, - fp8_model_params, - fuse_wgrad_accumulation, - bias, - delay_wgrad_compute, - parallel_mode=None, - use_cutlass=False, -): - fp8 = recipe is not None - if fp8 and fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: - pytest.skip("FP8 parameters are not supported in debug mode.") - if NVTE_TEST_NVINSPECT_ENABLED and delay_wgrad_compute: - pytest.skip("Delayed wgrad compute is not supported in debug mode.") - skip_unsupported_backward_override( - "grouped_linear", recipe, getattr(recipe, "backward_override", None) - ) - - config = model_configs[model] - if config.max_seqlen_q % 16 != 0 and fp8: - pytest.skip("FP8 requires sequence length to be divisible by 16.") - - if recipe is not None and recipe.nvfp4(): - if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype): - pytest.skip( - f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}" - ) - - with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe): - grouped_linear = GroupedLinear( - num_gemms, - config.hidden_size, - 4 * config.hidden_size, - bias=bias, - params_dtype=dtype, - parallel_mode=parallel_mode, - device="cuda", - fuse_wgrad_accumulation=fuse_wgrad_accumulation, - delay_wgrad_compute=delay_wgrad_compute, - save_original_input=False, - ).eval() - sequential_linear = torch.nn.ModuleList( - [ - Linear( - config.hidden_size, - 4 * config.hidden_size, - bias=bias, - params_dtype=dtype, - parallel_mode=parallel_mode, - device="cuda", - fuse_wgrad_accumulation=fuse_wgrad_accumulation, - ).eval() - for _ in range(num_gemms) - ] - ) - - # Share params - with torch.no_grad(): - for i in range(num_gemms): - sequential_linear[i].weight = Parameter(getattr(grouped_linear, f"weight{i}").clone()) - if bias: - sequential_linear[i].bias = Parameter(getattr(grouped_linear, f"bias{i}").clone()) - if fuse_wgrad_accumulation: - weight_i = getattr(grouped_linear, f"weight{i}") - weight_i.main_grad = torch.rand_like(weight_i, dtype=torch.float32) - sequential_linear[i].weight.main_grad = weight_i.main_grad.clone() - - outputs_ref = _test_grouped_linear_accuracy( - sequential_linear, - num_gemms, - bs, - dtype, - config, - recipe, - fp8, - fuse_wgrad_accumulation, - delay_wgrad_compute, - ) - outputs = _test_grouped_linear_accuracy( - grouped_linear, - num_gemms, - bs, - dtype, - config, - recipe, - fp8, - fuse_wgrad_accumulation, - delay_wgrad_compute, - ) - - for o, o_ref in zip(outputs, outputs_ref): - if use_cutlass: - torch.testing.assert_close(o, o_ref, rtol=1e-3, atol=1e-3) - else: - # cuBLAS implementation should be bit-wise match - torch.testing.assert_close(o, o_ref, rtol=0, atol=0) - - -@pytest.mark.skipif( - torch.cuda.get_device_capability() != (9, 0), - reason="Only enable CUTLASS grouped gemm on Hopper", -) -@pytest.mark.parametrize("dtype", param_types, ids=str) -@pytest.mark.parametrize("num_gemms", [3, 6]) -@pytest.mark.parametrize("bs", batch_sizes) -@pytest.mark.parametrize("model", ["126m"]) -@pytest.mark.parametrize("fuse_wgrad_accumulation", all_boolean) -@pytest.mark.parametrize("delay_wgrad_compute", all_boolean) -def test_grouped_linear_accuracy_cutlass( - dtype, - num_gemms, - bs, - model, - fuse_wgrad_accumulation, - delay_wgrad_compute, -): - os.environ["NVTE_USE_CUTLASS_GROUPED_GEMM"] = "1" - test_grouped_linear_accuracy( - dtype, - num_gemms, - bs, - model, - None, - False, - fuse_wgrad_accumulation, - False, - delay_wgrad_compute, - None, - use_cutlass=True, - ) - os.environ.pop("NVTE_USE_CUTLASS_GROUPED_GEMM", None) - - -@pytest.mark.parametrize("dtype", param_types, ids=str) -@pytest.mark.parametrize("num_gemms", [3]) -@pytest.mark.parametrize("bs", [1]) -@pytest.mark.parametrize("model", ["126m"]) -@pytest.mark.parametrize("recipe", fp8_recipes + [None], ids=recipe_id) -@pytest.mark.parametrize("fp8_model_params", [False]) -@pytest.mark.parametrize("fuse_wgrad_accumulation", [True]) -@pytest.mark.parametrize("bias", [False]) -@pytest.mark.parametrize("delay_wgrad_compute", [True]) -def test_grouped_linear_accuracy_save_original_input( - dtype, - num_gemms, - bs, - model, - recipe, - fp8_model_params, - fuse_wgrad_accumulation, - bias, - delay_wgrad_compute, - parallel_mode=None, -): - fp8 = recipe is not None - if fp8 and fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: - pytest.skip("FP8 parameters are not supported in debug mode.") - if fp8 and recipe.delayed(): - pytest.skip("DelayedScaling recipe is not supported with save_original_input") - if NVTE_TEST_NVINSPECT_ENABLED and delay_wgrad_compute: - pytest.skip("Delayed wgrad compute is not supported in debug mode.") - skip_unsupported_backward_override( - "grouped_linear", recipe, getattr(recipe, "backward_override", None) - ) - - config = model_configs[model] - if config.max_seqlen_q % 16 != 0 and fp8: - pytest.skip("FP8 requires sequence length to be divisible by 16.") - - if recipe is not None and recipe.nvfp4(): - if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype): - pytest.skip( - f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}" - ) - - with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe): - grouped_linear = GroupedLinear( - num_gemms, - config.hidden_size, - 4 * config.hidden_size, - bias=bias, - params_dtype=dtype, - parallel_mode=parallel_mode, - device="cuda", - fuse_wgrad_accumulation=fuse_wgrad_accumulation, - delay_wgrad_compute=delay_wgrad_compute, - save_original_input=True, - ).eval() - sequential_linear = torch.nn.ModuleList( - [ - Linear( - config.hidden_size, - 4 * config.hidden_size, - bias=bias, - params_dtype=dtype, - parallel_mode=parallel_mode, - device="cuda", - fuse_wgrad_accumulation=fuse_wgrad_accumulation, - ).eval() - for _ in range(num_gemms) - ] - ) - - # Share params - with torch.no_grad(): - for i in range(num_gemms): - sequential_linear[i].weight = Parameter(getattr(grouped_linear, f"weight{i}").clone()) - if bias: - sequential_linear[i].bias = Parameter(getattr(grouped_linear, f"bias{i}").clone()) - if fuse_wgrad_accumulation: - weight_i = getattr(grouped_linear, f"weight{i}") - weight_i.main_grad = torch.rand_like(weight_i, dtype=torch.float32) - sequential_linear[i].weight.main_grad = weight_i.main_grad.clone() - - outputs_ref = _test_grouped_linear_accuracy( - sequential_linear, - num_gemms, - bs, - dtype, - config, - recipe, - fp8, - fuse_wgrad_accumulation, - delay_wgrad_compute, - ) - outputs = _test_grouped_linear_accuracy( - grouped_linear, - num_gemms, - bs, - dtype, - config, - recipe, - fp8, - fuse_wgrad_accumulation, - delay_wgrad_compute, - ) - - # Shoule be bit-wise match - for i, (o, o_ref) in enumerate(zip(outputs, outputs_ref)): - torch.testing.assert_close(o, o_ref, rtol=0, atol=0) - - -@pytest.mark.parametrize("recipe", fp8_recipes + [None], ids=recipe_id) -def test_grouped_linear_accuracy_single_gemm(recipe): - """Split the tests to save CI time""" - test_grouped_linear_accuracy( - dtype=torch.float32, - num_gemms=1, - bs=2, - model="126m", - recipe=recipe, - fp8_model_params=True, - fuse_wgrad_accumulation=True, - bias=True, - delay_wgrad_compute=False, - ) - - -def _test_padding_grouped_linear_accuracy(block, num_gemms, bs, dtype, config, recipe, fp8=False): - - def _pad_tensor_for_fp8(hidden_states, tokens_per_expert): - align_size = get_align_size_for_quantization(recipe) - padded_tokens_per_expert = [ - (num_tokens + align_size - 1) // align_size * align_size - for num_tokens in tokens_per_expert - ] - hidden_states = torch.split(hidden_states, tokens_per_expert) - padded_hidden_states = [] - for hidden_state, actual_num_tokens, padded_num_tokens in zip( - hidden_states, tokens_per_expert, padded_tokens_per_expert - ): - padded_hidden_states.append(hidden_state) - if padded_num_tokens > actual_num_tokens: - pad_tensor = torch.zeros( - padded_num_tokens - actual_num_tokens, - hidden_state.shape[1], - dtype=hidden_state.dtype, - device=hidden_state.device, - ) - padded_hidden_states.append(pad_tensor) - padded_hidden_states = torch.cat(padded_hidden_states, dim=0) - return padded_hidden_states, padded_tokens_per_expert - - def _unpad_tensor_for_fp8(padded_hidden_states, actual_tokens_per_expert, tokens_per_expert): - inputmats = torch.split( - padded_hidden_states.view(-1, padded_hidden_states.shape[-1]), tokens_per_expert - ) - hidden_states = torch.cat( - [ - grad_output_mat[: actual_tokens_per_expert[i]] - for i, grad_output_mat in enumerate(inputmats) - ], - dim=0, - ) - - return hidden_states - - def _generate_random_numbers(n, total_sum): - if n <= 0: - return [] - - # reset seed - random.seed(seed) - - breaks = sorted(random.sample(range(1, total_sum), n - 1)) - random_numbers = ( - [breaks[0]] - + [breaks[i] - breaks[i - 1] for i in range(1, n - 1)] - + [total_sum - breaks[-1]] - ) - - return random_numbers - - reset_rng_states() - if fp8: - FP8GlobalStateManager.reset() - - inp_hidden_states = torch.randn( - (config.max_seqlen_q * bs, config.hidden_size), - dtype=dtype, - device="cuda", - requires_grad=True, - ) - inp_hidden_states.retain_grad() - - m_splits = _generate_random_numbers(num_gemms, config.max_seqlen_q * bs) - - with autocast(enabled=fp8, recipe=recipe): - if isinstance(block, TorchGroupedLinearWithPadding): - out = block(inp_hidden_states, m_splits) - else: - if fp8: - padded_inp_hidden_states, padding_m_splits = _pad_tensor_for_fp8( - inp_hidden_states, m_splits - ) - padded_inp_hidden_states = block(padded_inp_hidden_states, padding_m_splits) - out = _unpad_tensor_for_fp8(padded_inp_hidden_states, m_splits, padding_m_splits) - else: - out = block(inp_hidden_states, m_splits) - - loss = out.sum() - loss.backward() - - torch.cuda.synchronize() - outputs = [out, inp_hidden_states.grad] - for p in block.parameters(): - if p.requires_grad: - outputs.append(p.grad) - return outputs - - -@pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("num_gemms", [3, 6]) -@pytest.mark.parametrize("bs", batch_sizes) -@pytest.mark.parametrize("model", ["126m"]) -@pytest.mark.parametrize("fp8", [True]) -@pytest.mark.parametrize("recipe", fp8_recipes, ids=recipe_id) -@pytest.mark.parametrize("fp8_model_params", all_boolean) -def test_padding_grouped_linear_accuracy( - dtype, - num_gemms, - bs, - model, - fp8, - recipe, - fp8_model_params, - parallel_mode=None, -): - if fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: - pytest.skip("FP8 parameters are not supported in debug mode.") - skip_unsupported_backward_override( - "grouped_linear", recipe, getattr(recipe, "backward_override", None) - ) - - config = model_configs[model] - if config.max_seqlen_q % 16 != 0 and fp8: - pytest.skip("FP8 requires sequence length to be divisible by 16.") - - if recipe is not None and recipe.nvfp4(): - if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype): - pytest.skip( - f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}" - ) - - with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe): - grouped_linear = TorchGroupedLinearWithPadding( - num_gemms, - config.hidden_size, - 4 * config.hidden_size, - bias=False, - params_dtype=dtype, - parallel_mode=parallel_mode, - fp8=fp8, - ).eval() - - with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe): - ref_grouped_linear = GroupedLinear( - num_gemms, - config.hidden_size, - 4 * config.hidden_size, - bias=False, - params_dtype=dtype, - parallel_mode=parallel_mode, - device="cuda", - save_original_input=False, - ).eval() - - # Share params - with torch.no_grad(): - inner_grouped_linear = grouped_linear.linear_fn - for i in range(num_gemms): - setattr( - ref_grouped_linear, - f"weight{i}", - Parameter(getattr(inner_grouped_linear, f"weight{i}").clone()), - ) - - outputs = _test_padding_grouped_linear_accuracy( - grouped_linear, num_gemms, bs, dtype, config, recipe, fp8 - ) - outputs_ref = _test_padding_grouped_linear_accuracy( - ref_grouped_linear, num_gemms, bs, dtype, config, recipe, fp8 - ) - - # Shoule be bit-wise match - for i, (o, o_ref) in enumerate(zip(outputs, outputs_ref)): - torch.testing.assert_close(o, o_ref, rtol=0, atol=0) - - -@pytest.mark.parametrize("dtype", param_types) -@pytest.mark.parametrize("num_gemms", [3]) -@pytest.mark.parametrize("bs", [1]) -@pytest.mark.parametrize("model", ["126m"]) -@pytest.mark.parametrize("fp8", [True]) -@pytest.mark.parametrize("recipe", fp8_recipes, ids=recipe_id) -@pytest.mark.parametrize("fp8_model_params", [False]) -def test_padding_grouped_linear_accuracy_save_original_input( - dtype, - num_gemms, - bs, - model, - fp8, - recipe, - fp8_model_params, - parallel_mode=None, -): - if fp8_model_params and NVTE_TEST_NVINSPECT_ENABLED: - pytest.skip("FP8 parameters are not supported in debug mode.") - if fp8 and recipe.delayed(): - pytest.skip("DelayedScaling recipe is not supported with save_original_input") - skip_unsupported_backward_override( - "grouped_linear", recipe, getattr(recipe, "backward_override", None) - ) - - config = model_configs[model] - if config.max_seqlen_q % 16 != 0 and fp8: - pytest.skip("FP8 requires sequence length to be divisible by 16.") - - if recipe is not None and recipe.nvfp4(): - if dtype not in get_nvfp4_inp_supported_dtypes(recipe, dtype): - pytest.skip( - f"Input dtype {dtype} not supported for NVFP4 Recipe {recipe.__class__.__name__}" - ) - - with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe): - grouped_linear = TorchGroupedLinearWithPadding( - num_gemms, - config.hidden_size, - 4 * config.hidden_size, - bias=False, - params_dtype=dtype, - parallel_mode=parallel_mode, - fp8=fp8, - ).eval() - - with quantized_model_init(enabled=fp8 and fp8_model_params, recipe=recipe): - ref_grouped_linear = GroupedLinear( - num_gemms, - config.hidden_size, - 4 * config.hidden_size, - bias=False, - params_dtype=dtype, - parallel_mode=parallel_mode, - device="cuda", - save_original_input=True, - ).eval() - - # Share params - with torch.no_grad(): - inner_grouped_linear = grouped_linear.linear_fn - for i in range(num_gemms): - setattr( - ref_grouped_linear, - f"weight{i}", - Parameter(getattr(inner_grouped_linear, f"weight{i}").clone()), - ) - - outputs = _test_padding_grouped_linear_accuracy( - grouped_linear, num_gemms, bs, dtype, config, recipe, fp8 - ) - outputs_ref = _test_padding_grouped_linear_accuracy( - ref_grouped_linear, num_gemms, bs, dtype, config, recipe, fp8 - ) - - # Shoule be bit-wise match - for i, (o, o_ref) in enumerate(zip(outputs, outputs_ref)): - torch.testing.assert_close(o, o_ref, rtol=0, atol=0) - - def _test_gpt_e2e_cuda_graph(block, bs, dtype, config, graph): reset_rng_states() @@ -2761,594 +2121,6 @@ def test_transformer_layer_hidden_states_format(dtype, bs, model): ) -@pytest.mark.parametrize( - "shape", - [ - (1, 127, 128, 512), - (8, 15, 128, 512), - (8, 1027, 128, 512), - (16, 10027, 128, 512), - ], -) -@pytest.mark.parametrize("dtype", param_types, ids=str) -@pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) -@pytest.mark.parametrize("accumulate", [False, True]) -@pytest.mark.parametrize("use_cutlass", use_cutlass_grouped_gemm) -def test_grouped_gemm(shape, dtype, layout, accumulate, use_cutlass): - torch.manual_seed(0) - z, m, k, n = shape - - dist = torch.sort(torch.randint(0, m, (z - 1,))).values.tolist() - m_splits = torch.tensor(dist + [m]) - torch.tensor([0] + dist) - assert m_splits.sum() == m and len(m_splits) == z - m_splits = m_splits.tolist() - - if layout == "TN": - A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # weight - B = list(torch.split(torch.randn(m, k, dtype=dtype, device="cuda"), m_splits)) # input - out = [torch.randn(m, n, dtype=dtype, device="cuda")] # output - out_ref = [o.clone() for o in torch.split(out[0], m_splits)] - grad = False - single_output = True - elif layout == "NN": - A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # weight - B = list( - torch.split(torch.randn(m, n, dtype=dtype, device="cuda"), m_splits) - ) # grad_output - out = [torch.randn(m, k, dtype=dtype, device="cuda")] # dgrad - out_ref = [o.clone() for o in torch.split(out[0], m_splits)] - grad = True - single_output = True - else: # layout == "NT" - A = list(torch.split(torch.randn(m, k, dtype=dtype, device="cuda"), m_splits)) # input - B = list( - torch.split(torch.randn(m, n, dtype=dtype, device="cuda"), m_splits) - ) # grad_output - out = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # wgrad - out_ref = [o.clone() for o in out] - grad = True - single_output = False - - if use_cutlass: - os.environ["NVTE_USE_CUTLASS_GROUPED_GEMM"] = "1" - - for i in range(z): - general_gemm( - A[i], - B[i], - dtype, - grad=grad, - accumulate=accumulate, - layout=layout, - out=out_ref[i], - ) - if single_output: - out_ref = [torch.cat(out_ref)] - - general_grouped_gemm( - A, - B, - out, - [None] * z, - dtype, - m_splits=m_splits, - grad=grad, - accumulate=accumulate, - layout=layout, - single_output=single_output, - ) - - for o, o_ref in zip(out, out_ref): - if not use_cutlass: - # cublas implementation should be bit-wise match - torch.testing.assert_close(o, o_ref, rtol=0, atol=0) - else: - torch.testing.assert_close(o, o_ref, rtol=1.5e-2, atol=1.5e-2) - - if use_cutlass: - os.environ.pop("NVTE_USE_CUTLASS_GROUPED_GEMM", None) - - -def _pack_grouped_tensor(grouped_tensor: GroupedTensor, tensors: List[torch.Tensor]) -> None: - data = grouped_tensor.rowwise_data - if data is None: - data = grouped_tensor.columnwise_data - if data is None: - raise ValueError("GroupedTensor has no data buffers to pack.") - offset = 0 - for tensor in tensors: - numel = tensor.numel() - data[offset : offset + numel].copy_(tensor.reshape(-1)) - offset += numel - - -def _make_grouped_tensor_from_splits( - m_sizes: List[int], - last_dim: int, - device: torch.device, - dtype: torch.dtype, -) -> GroupedTensor: - first_dims = torch.tensor(m_sizes, device=device, dtype=torch.int64) - return GroupedTensor.make_grouped_tensor( - num_tensors=len(m_sizes), - first_dims=first_dims, - last_dims=None, - logical_first_dim=sum(m_sizes), - logical_last_dim=last_dim, - quantizer=None, - device=device, - dtype=dtype, - ) - - -def _make_grouped_tensor_uniform( - num_tensors: int, - first_dim: int, - last_dim: int, - device: torch.device, - dtype: torch.dtype, -) -> GroupedTensor: - return GroupedTensor.make_grouped_tensor( - num_tensors=num_tensors, - first_dims=None, - last_dims=None, - logical_first_dim=num_tensors * first_dim, - logical_last_dim=last_dim, - quantizer=None, - device=device, - dtype=dtype, - ) - - -def _apply_grouped_bias_ref( - base_outs: List[torch.Tensor], - bias: Optional[List[torch.Tensor]], - bias_scale: Optional[torch.Tensor], - m_sizes: List[int], - dtype: torch.dtype, -) -> List[torch.Tensor]: - """Reference: add (optionally per-row scaled) bias to each group's output, cast to ``dtype``.""" - if bias is None: - return list(base_outs) - if bias_scale is None: - return [(o.float() + b.float()).to(dtype) for o, b in zip(base_outs, bias)] - out = [] - offset = 0 - for i, ms in enumerate(m_sizes): - s = bias_scale[offset : offset + ms].unsqueeze(-1) - out.append((base_outs[i].float() + bias[i].float() * s).to(dtype)) - offset += ms - return out - - -@pytest.mark.parametrize( - "z, m, n, k", - [ - (4, 256, 256, 256), - (4, 512, 256, 512), - (4, 512, 512, 256), - (8, 512, 256, 512), - ], -) -@pytest.mark.parametrize("case", ["no_discrete", "discrete_in", "discrete_out"]) -@pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) -@pytest.mark.parametrize("accumulate", [False, True]) -@pytest.mark.parametrize("use_bias_scale", [False, True]) -def test_grouped_gemm_grouped_tensor(z, m, n, k, case, layout, accumulate, use_bias_scale) -> None: - if torch.cuda.get_device_capability() < (9, 0): - pytest.skip("Grouped GEMM requires Hopper (SM90) or newer.") - if torch.cuda.get_device_capability() < (10, 0): - if tex.get_cublasLt_version() < 130400: - pytest.skip("Grouped GEMM on Hopper requires cuBLAS 13.4+.") - if tex.get_cublasLt_version() < 130300: - pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") - if not is_bf16_available(): - pytest.skip("bfloat16 is required for grouped GEMM test.") - - torch.manual_seed(0) - - dtype = torch.bfloat16 - - split_points = torch.randperm(m - 1)[: z - 1] + 1 - split_points = torch.sort(split_points).values.tolist() - m_sizes = [split_points[0]] - m_sizes += [b - a for a, b in zip(split_points[:-1], split_points[1:])] - m_sizes.append(m - split_points[-1]) - assert sum(m_sizes) == m and len(m_sizes) == z - - if layout == "NT": - A = [torch.randn(ms, k, dtype=dtype, device="cuda") for ms in m_sizes] # input - B = [torch.randn(ms, n, dtype=dtype, device="cuda") for ms in m_sizes] # grad_output - out = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # wgrad - out_ref = [torch.matmul(B[i].transpose(0, 1).float(), A[i].float()) for i in range(z)] - else: - A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # weight - B = [ - torch.randn(ms, k if layout == "TN" else n, dtype=dtype, device="cuda") - for ms in m_sizes - ] # TN --> input, NN --> grad_output - out = [ - torch.randn(ms, n if layout == "TN" else k, dtype=dtype, device="cuda") - for ms in m_sizes - ] # TN --> output, NN --> dgrad - if layout == "NN": - out_ref = [torch.matmul(B[i].float(), A[i].float()) for i in range(z)] - else: # layout == "TN" - out_ref = [torch.matmul(B[i].float(), A[i].transpose(0, 1).float()) for i in range(z)] - - if accumulate: - out_ref = [out[i].float() + o for i, o in enumerate(out_ref)] - - # Bias is applied after GEMM (broadcasted along rows) - # Match kernel behavior: GEMM output is already in output dtype when bias is added. - out_ref_no_bias = [o.to(dtype) for o in out_ref] - if layout == "TN": - bias_last_dim = n - else: # layout == "NT" or "NN" - bias_last_dim = k - bias = ( - [torch.randn(1, bias_last_dim, dtype=dtype, device="cuda") for _ in range(z)] - if case != "discrete_out" - else None - ) - bias_scale = None - if use_bias_scale and bias is not None and layout != "NT": - bias_scale = torch.randn(m, device="cuda", dtype=torch.float32) - # Bias add in grouped kernel accumulates in FP32 for BF16/FP16. - out_ref = _apply_grouped_bias_ref(out_ref_no_bias, bias, bias_scale, m_sizes, dtype) - # Create grouped tensors based on case - device = A[0].device - grouped_A = A - grouped_out = out - grouped_out_bias = [o.clone() for o in out] - grouped_out_no_bias = [o.clone() for o in out] - grouped_bias = None - if layout == "TN": - grouped_A = ( - _make_grouped_tensor_uniform(z, n, k, device, dtype) if case != "discrete_in" else A - ) # weight - grouped_B = _make_grouped_tensor_from_splits(m_sizes, k, device, dtype) # input - if case != "discrete_out": - grouped_out = _make_grouped_tensor_from_splits(m_sizes, n, device, dtype) # output - grouped_out_bias = _make_grouped_tensor_from_splits(m_sizes, n, device, dtype) - grouped_out_no_bias = _make_grouped_tensor_from_splits(m_sizes, n, device, dtype) - elif layout == "NN": - grouped_A = ( - _make_grouped_tensor_uniform(z, n, k, device, dtype) if case != "discrete_in" else A - ) # weight - grouped_B = _make_grouped_tensor_from_splits(m_sizes, n, device, dtype) # grad_output - if case != "discrete_out": - grouped_out = _make_grouped_tensor_from_splits(m_sizes, k, device, dtype) - grouped_out_bias = _make_grouped_tensor_from_splits(m_sizes, k, device, dtype) - grouped_out_no_bias = _make_grouped_tensor_from_splits(m_sizes, k, device, dtype) - else: # layout == "NT" - grouped_A = ( - _make_grouped_tensor_from_splits(m_sizes, k, device, dtype) - if case != "discrete_in" - else A - ) # input - grouped_B = _make_grouped_tensor_from_splits(m_sizes, n, device, dtype) # grad_output - if case != "discrete_out": - grouped_out = _make_grouped_tensor_uniform(z, n, k, device, dtype) # wgrad - grouped_out_bias = _make_grouped_tensor_uniform(z, n, k, device, dtype) - grouped_out_no_bias = _make_grouped_tensor_uniform(z, n, k, device, dtype) - _pack_grouped_tensor(grouped_B, B) - if case != "discrete_out": - _pack_grouped_tensor(grouped_out, out) - _pack_grouped_tensor(grouped_out_bias, out) - _pack_grouped_tensor(grouped_out_no_bias, out) - if case != "discrete_in": - _pack_grouped_tensor(grouped_A, A) - - if bias is not None: - grouped_bias = _make_grouped_tensor_uniform(z, 1, bias_last_dim, device, dtype) - _pack_grouped_tensor(grouped_bias, bias) - - general_grouped_gemm_for_grouped_tensor( - grouped_A, - grouped_B, - grouped_out_no_bias, - layout=layout, - accumulate=accumulate, - bias=None, - ) - general_grouped_gemm_for_grouped_tensor( - grouped_A, - grouped_B, - grouped_out_bias, - layout=layout, - accumulate=accumulate, - bias=grouped_bias, - bias_scale=bias_scale, - ) - out_grouped_no_bias = ( - grouped_out_no_bias - if isinstance(grouped_out_no_bias, list) - else grouped_out_no_bias.split_into_quantized_tensors() - ) - out_grouped_bias = ( - grouped_out_bias - if isinstance(grouped_out_bias, list) - else grouped_out_bias.split_into_quantized_tensors() - ) - - out_grouped_manual_bias = _apply_grouped_bias_ref( - out_grouped_no_bias, bias, bias_scale, m_sizes, dtype - ) - tols = dtype_tols(dtype) - for o, o_ref in zip(out_grouped_no_bias, out_ref_no_bias): - torch.testing.assert_close(o, o_ref, **tols) - if bias is not None: - for o, o_ref in zip(out_grouped_bias, out_grouped_manual_bias): - torch.testing.assert_close(o, o_ref, **tols) - - -@pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) -@pytest.mark.parametrize("accumulate", [False, True]) -@pytest.mark.parametrize("quant_type", ["bf16", "mxfp8"]) -def test_grouped_gemm_grouped_tensor_zero_work(layout, accumulate, quant_type) -> None: - """Grouped GEMM with all-zero split sizes (zero total work). - - For wgrad (NT layout) the output should be zero when not accumulating, - or unchanged when accumulating with beta=1. - """ - if torch.cuda.get_device_capability() < (10, 0): - pytest.skip("Grouped GEMM requires Blackwell (SM100) or newer.") - if not is_bf16_available(): - pytest.skip("bfloat16 is required for grouped GEMM test.") - if quant_type == "mxfp8" and not mxfp8_available: - pytest.skip(reason_for_no_mxfp8) - - z = 4 - k, n = 256, 256 - dtype = torch.bfloat16 - device = torch.device("cuda") - use_mxfp8 = quant_type == "mxfp8" - - transa = layout[0] == "T" - transb = layout[1] == "T" - zero_first_dims = torch.zeros(z, dtype=torch.int64, device=device) - - def _make_zero_tokens_grouped_tensor(logical_last_dim, is_a): - """Create a GroupedTensor with non-zero logical_shape but zero first_dims.""" - buf = torch.randn(0, logical_last_dim, dtype=dtype, device=device) - if use_mxfp8: - if is_a: - rowwise, columnwise = transa, not transa - else: - rowwise, columnwise = not transb, transb - quantizer = MXFP8Quantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=rowwise, - columnwise=columnwise, - ) - quantizer.optimize_for_gemm = True - return tex.group_quantize(buf, quantizer, z, zero_first_dims) - return GroupedTensor.make_grouped_tensor( - num_tensors=z, - first_dims=zero_first_dims, - last_dims=None, - logical_first_dim=k, - logical_last_dim=logical_last_dim, - quantizer=None, - device=device, - dtype=dtype, - ) - - if layout in ("TN", "NN"): - weight_tensors = [torch.randn(n, k, dtype=dtype, device=device) for _ in range(z)] - if use_mxfp8: - grouped_A = _make_grouped_tensor_quantized_mxfp8( - weight_tensors, - rowwise=transa, - columnwise=not transa, - device=device, - ) - else: - grouped_A = _make_grouped_tensor_uniform(z, n, k, device, dtype) - _pack_grouped_tensor(grouped_A, weight_tensors) - else: # NT - grouped_A = _make_zero_tokens_grouped_tensor(k, is_a=True) - - b_last_dim = k if layout == "TN" else n - grouped_B = _make_zero_tokens_grouped_tensor(b_last_dim, is_a=False) - - if layout == "NT": - out = [torch.randn(n, k, dtype=dtype, device=device) for _ in range(z)] - grouped_out = _make_grouped_tensor_uniform(z, n, k, device, dtype) - _pack_grouped_tensor(grouped_out, out) - else: - out = [torch.zeros(0, dtype=dtype, device=device) for _ in range(z)] - out_last_dim = n if layout == "TN" else k - grouped_out = GroupedTensor.make_grouped_tensor( - num_tensors=z, - first_dims=zero_first_dims, - last_dims=None, - logical_first_dim=k, - logical_last_dim=out_last_dim, - quantizer=None, - device=device, - dtype=dtype, - ) - - out_before = [o.clone() for o in out] - - general_grouped_gemm_for_grouped_tensor( - grouped_A, - grouped_B, - grouped_out, - layout=layout, - accumulate=accumulate, - ) - - out_result = ( - grouped_out if isinstance(grouped_out, list) else grouped_out.split_into_quantized_tensors() - ) - for i in range(z): - if out_result[i].numel() == 0: - continue - if accumulate: - torch.testing.assert_close(out_result[i], out_before[i]) - else: - torch.testing.assert_close(out_result[i], torch.zeros_like(out_result[i])) - - -def _make_grouped_tensor_quantized_mxfp8( - tensors: List[torch.Tensor], - *, - rowwise: bool, - columnwise: bool, - device: torch.device, - is_weight: bool = False, -) -> GroupedTensor: - """Create a quantized MXFP8 GroupedTensor from a list of per-expert tensors. - - For weights (uniform per-expert shape), we generally won't keep it swizzled since we - might need for future dequantize operations. Swizzling is done internally within - general_grouped_gemm_for_grouped_tensor call. - - For non-weight tensors (inputs / grad_outputs), we still pass - ``first_dims`` and keep ``optimize_for_gemm=True``; so the kernel must emit the - already-swizzled layout up front. - """ - if not tensors: - raise ValueError("Expected non-empty tensor list for grouped quantization.") - quantizer = MXFP8Quantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=rowwise, - columnwise=columnwise, - ) - quantizer.optimize_for_gemm = not is_weight - grouped_input = torch.cat(tensors, dim=0) - if is_weight: - first_dims = None - else: - first_dims = torch.tensor([t.shape[0] for t in tensors], dtype=torch.int64, device=device) - return tex.group_quantize(grouped_input, quantizer, len(tensors), first_dims) - - -def _per_tensor_quantize_mxfp8( - tensors: List[torch.Tensor], - *, - rowwise: bool, - columnwise: bool, -) -> List: - """Quantize each tensor individually with MXFP8. - Used to build reference discrete inputs for grouped GEMM. - """ - quantizer = MXFP8Quantizer( - fp8_dtype=tex.DType.kFloat8E4M3, - rowwise=rowwise, - columnwise=columnwise, - ) - return [quantizer(t) for t in tensors] - - -@pytest.mark.parametrize( - "shape", - [ - (1, 128, 128, 512), - (8, 1024, 128, 512), - (16, 4096, 128, 512), - (2, 256, 2880, 2880), - ], -) -@pytest.mark.parametrize("accumulate", [False, True]) -@pytest.mark.parametrize("layout", ["TN", "NN", "NT"]) -@pytest.mark.parametrize("case", ["no_discrete", "discrete_in", "discrete_out"]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16]) -def test_grouped_gemm_grouped_tensor_mxfp8( - shape, accumulate, layout: str, case: str, dtype: torch.dtype -) -> None: - if tex.get_cublasLt_version() < 130300: - pytest.skip("Grouped GEMM requires cuBLAS 13.3+.") - if torch.cuda.get_device_capability() < (10, 0): - pytest.skip("Grouped GEMM requires Blackwell (SM100) or newer.") - if dtype == torch.bfloat16 and not is_bf16_available(): - pytest.skip("bfloat16 is required for grouped GEMM test.") - - torch.manual_seed(0) - z, m, k, n = shape - m_sizes = [m // z] * z - - if layout == "TN": - A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # weight - B = [torch.randn(ms, k, dtype=dtype, device="cuda") for ms in m_sizes] # input - out = [torch.randn(ms, n, dtype=dtype, device="cuda") for ms in m_sizes] # output - grad = False - elif layout == "NN": - A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # weight - B = [torch.randn(ms, n, dtype=dtype, device="cuda") for ms in m_sizes] # grad_output - out = [torch.randn(ms, k, dtype=dtype, device="cuda") for ms in m_sizes] # dgrad - grad = True - else: # layout == "NT" - A = [torch.randn(ms, k, dtype=dtype, device="cuda") for ms in m_sizes] # input - B = [torch.randn(ms, n, dtype=dtype, device="cuda") for ms in m_sizes] # grad_output - out = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # wgrad - grad = True - - out_ref = [o.clone() for o in out] - - transa = layout[0] == "T" - transb = layout[1] == "T" - a_is_weight = all(t.shape == A[0].shape for t in A) - a_rowwise, a_columnwise = transa, not transa - b_rowwise, b_columnwise = not transb, transb - grouped_A = _make_grouped_tensor_quantized_mxfp8( - A, - rowwise=a_rowwise, - columnwise=a_columnwise, - device="cuda", - is_weight=a_is_weight, - ) - grouped_B = _make_grouped_tensor_quantized_mxfp8( - B, rowwise=b_rowwise, columnwise=b_columnwise, device="cuda" - ) - A_fp8 = _per_tensor_quantize_mxfp8(A, rowwise=a_rowwise, columnwise=a_columnwise) - B_fp8 = _per_tensor_quantize_mxfp8(B, rowwise=b_rowwise, columnwise=b_columnwise) - - general_grouped_gemm( - A_fp8, - B_fp8, - out_ref, - [None] * z, - dtype, - m_splits=m_sizes, - grad=grad, - accumulate=accumulate, - layout=layout, - single_output=False, - ) - - device = A[0].device - - grouped_out = None - if case != "discrete_out": - if layout == "TN": - grouped_out = _make_grouped_tensor_from_splits(m_sizes, n, device, dtype) - elif layout == "NN": - grouped_out = _make_grouped_tensor_from_splits(m_sizes, k, device, dtype) - else: # layout == "NT" - grouped_out = _make_grouped_tensor_uniform(z, n, k, device, dtype) - _pack_grouped_tensor(grouped_out, out) - - grouped_out_input = out if case == "discrete_out" else grouped_out - grouped_A_input = A_fp8 if case == "discrete_in" else grouped_A - general_grouped_gemm_for_grouped_tensor( - grouped_A_input, - grouped_B, - grouped_out_input, - layout=layout, - accumulate=accumulate, - ) - - out_grouped = out if case == "discrete_out" else grouped_out.split_into_quantized_tensors() - tols = dict(rtol=0.125, atol=0.0675) # mxfp8 tolerance - - for o, o_ref in zip(out_grouped, out_ref): - torch.testing.assert_close(o, o_ref, **tols) - - @pytest.mark.parametrize("N", [32]) @pytest.mark.parametrize("datatype", [torch.float16, torch.bfloat16]) @pytest.mark.parametrize( @@ -3414,80 +2186,6 @@ def test_fp8gemm_with_unfused_quantization(N, datatype, input_quantizer, out_qua torch.testing.assert_close(expected_quantized_out.dequantize(), quantized_out.dequantize()) -@pytest.mark.parametrize( - "shape", - [ - (1, 128, 128, 512), - (8, 1024, 128, 512), - (16, 4096, 128, 512), - ], -) -@pytest.mark.parametrize("accumulate", [False, True]) -def test_fp8_grouped_gemm(shape, accumulate): - if not fp8_available: - pytest.skip(reason_for_no_fp8) - - z, m, k, n = shape - m_splits = [m // z] * z - - dtype = torch.bfloat16 - A = [torch.randn(n, k, dtype=dtype, device="cuda") for _ in range(z)] # weight - B = torch.split(torch.randn(m, k, dtype=dtype, device="cuda"), m_splits) # input - out = torch.split(torch.randn(m, n, dtype=dtype, device="cuda"), m_splits) # output - out_ref = [o.clone() for o in out] - - # fp8 should be robust enough to this fake scale - scale = 1 + torch.rand(1, dtype=torch.float32, device="cuda").squeeze() - amax = torch.zeros(1, 1, dtype=torch.float32, device="cuda") - - a_quantizers = [ - Float8Quantizer( - scale.clone(), - amax.clone(), - tex.DType.kFloat8E4M3, - ) - for _ in range(z) - ] - b_quantizers = [ - Float8Quantizer( - scale.clone(), - amax.clone(), - tex.DType.kFloat8E4M3, - ) - for _ in range(z) - ] - - A_fp8 = [] - B_fp8 = [] - - for i in range(z): - A_fp8.append(a_quantizers[i](A[i])) - B_fp8.append(b_quantizers[i](B[i])) - - # baseline - for i in range(z): - general_gemm( - A_fp8[i], - B_fp8[i], - dtype, - out=out_ref[i], - accumulate=accumulate, - ) - general_grouped_gemm( - A_fp8, - B_fp8, - out, - [None] * z, - dtype, - m_splits=m_splits, - accumulate=accumulate, - ) - - # should be bit-wise match - for o, o_ref in zip(out, out_ref): - torch.testing.assert_close(o, o_ref, rtol=0, atol=0) - - def test_noncontiguous(): def _create2modules(m, params): mod1 = m(*params) diff --git a/transformer_engine/pytorch/module/grouped_linear.py b/transformer_engine/pytorch/module/grouped_linear.py index b2baf1729..15ec3fe32 100644 --- a/transformer_engine/pytorch/module/grouped_linear.py +++ b/transformer_engine/pytorch/module/grouped_linear.py @@ -3,8 +3,10 @@ # See LICENSE for license information. """GroupedLinear API""" + from typing import Union, Optional, Callable, Tuple, List from itertools import chain +import os import warnings import weakref @@ -29,6 +31,7 @@ divide, cast_if_needed, clear_tensor_data, + get_device_compute_capability, init_method_constant, requires_grad, resolve_grouped_linear_single_param_flags, @@ -42,12 +45,15 @@ ) from ..cpp_extensions import ( general_grouped_gemm, + general_grouped_gemm_for_grouped_tensor, ) from ..constants import GemmParallelModes, dist_group_type from ..jit import no_torch_dynamo from ..cpu_offload import is_cpu_offload_enabled, mark_not_offload, start_offload +from ..triton.grouped_dbias_dscales import compute_grouped_dbias from ..tensor.float8_tensor import Float8CurrentScalingQuantizer, Float8Quantizer +from ..tensor.mxfp8_tensor import MXFP8Quantizer from ..quantized_tensor import ( QuantizedTensorStorage, Quantizer, @@ -65,11 +71,316 @@ class _GroupedLinear(torch.autograd.Function): Calls custom cuda extensions. """ + @staticmethod + def _maybe_dequantize( + tensor: Union[torch.Tensor, QuantizedTensorStorage], + dtype: torch.dtype, + ) -> torch.Tensor: + """Dequantize quantized tensors or cast regular tensors to ``dtype``.""" + if isinstance(tensor, QuantizedTensorStorage): + return tensor.dequantize(dtype=dtype) + return cast_if_needed(tensor, dtype) + + @staticmethod + def _is_grouped_tensor_path_supported( + *, + fp8: bool, + fp8_calibration: bool, + debug: bool, + cpu_offloading: bool, + backward_override: Optional[str], + save_original_input: bool, + activation_dtype: torch.dtype, + input_quantizers: List[Optional[Quantizer]], + weight_quantizers: List[Optional[Quantizer]], + output_quantizers: List[Optional[Quantizer]], + grad_output_quantizers: List[Optional[Quantizer]], + ) -> bool: + """Whether to use cublasLt grouped GEMM through GroupedTensor metadata. + + There are no checks whether split sizes are supported. Splits + may be in a CUDA tensor, so checking would hurt performance + and be incompatible with CUDA Graphs. + + """ + if not bool(int(os.getenv("NVTE_GROUPED_LINEAR_USE_FUSED_GROUPED_GEMM", "0"))): + return False + if ( + debug + or cpu_offloading + or fp8_calibration + or backward_override is not None + or save_original_input + ): + return False + if get_device_compute_capability() < (10, 0): + return False + if any(q is not None for q in output_quantizers): + return False + if fp8: + return ( + activation_dtype in (torch.bfloat16, torch.float16) + and all(isinstance(q, MXFP8Quantizer) for q in input_quantizers) + and all(isinstance(q, MXFP8Quantizer) for q in weight_quantizers) + and all(q is None or isinstance(q, MXFP8Quantizer) for q in grad_output_quantizers) + ) + return activation_dtype in (torch.bfloat16, torch.float16) + + @staticmethod + def _make_grouped_tensor( + data: torch.Tensor, + *, + num_gemms: int, + split_sizes: torch.Tensor, + base_split_offsets: torch.Tensor, + last_dim: int, + dtype: torch.dtype, + ) -> GroupedTensor: + """Wrap a packed 2D buffer as a varying-first-dimension GroupedTensor.""" + return GroupedTensor( + shape=(data.size(0), last_dim), + dtype=dtype, + num_tensors=num_gemms, + quantizer=None, + data=data.reshape(-1), + first_dims=split_sizes, + tensor_offsets=base_split_offsets * last_dim, + ) + + @staticmethod + def _make_grouped_bias( + biases: Tuple[torch.Tensor, ...], + *, + num_gemms: int, + out_features: int, + dtype: torch.dtype, + ) -> GroupedTensor: + """Pack per-GEMM biases into the grouped GEMM bias format.""" + bias_data = torch.stack( + [_GroupedLinear._maybe_dequantize(bias, dtype) for bias in biases], + dim=0, + ).contiguous() + return GroupedTensor( + shape=(num_gemms, out_features), + dtype=dtype, + num_tensors=num_gemms, + shapes=[(1, out_features)] * num_gemms, + quantizer=None, + data=bias_data.reshape(-1), + ) + + @staticmethod + def _prepare_weights_for_grouped_tensor_gemm( + weights: Tuple[torch.Tensor, ...], + weight_quantizers: List[Optional[Quantizer]], + weight_workspaces: List[Optional[QuantizedTensorStorage]], + *, + with_quantized_compute: bool, + columnwise_usage: bool, + activation_dtype: torch.dtype, + is_first_microbatch: Optional[bool], + skip_fp8_weight_update: Optional[torch.Tensor], + cache_weight: bool, + ) -> Tuple[List[torch.Tensor], List[Optional[QuantizedTensorStorage]]]: + """Prepare discrete weight tensors for GroupedTensor GEMM.""" + weights_for_gemm: List[torch.Tensor] = [] + new_workspaces: List[Optional[QuantizedTensorStorage]] = [None] * len(weights) + if not with_quantized_compute: + return ( + [_GroupedLinear._maybe_dequantize(weight, activation_dtype) for weight in weights], + new_workspaces, + ) + + update_ws = is_first_microbatch is None or is_first_microbatch + for idx, weight in enumerate(weights): + weight_quantizer = weight_quantizers[idx] + weight_quantizer.set_usage(rowwise=True, columnwise=columnwise_usage) + weight_fp8, new_workspaces[idx] = quantize_weight( + tensor=weight, + quantizer=weight_quantizer, + workspace=weight_workspaces[idx] if weight_workspaces else None, + update_workspace=update_ws, + skip_update_flag=skip_fp8_weight_update, + workspace_dtype=activation_dtype, + cache=cache_weight, + ) + weights_for_gemm.append(weight_fp8) + return weights_for_gemm, new_workspaces + + @staticmethod + def _forward_grouped_tensor( + ctx, + *, + inp: torch.Tensor, + m_splits: torch.Tensor, + use_bias: bool, + is_first_microbatch: Optional[bool], + fp8: bool, + wgrad_store: Optional[WeightGradStore], + input_quantizers: List[Optional[Quantizer]], + weight_quantizers: List[Optional[Quantizer]], + grad_input_quantizers: List[Optional[Quantizer]], + grad_weight_quantizers: List[Optional[Quantizer]], + grad_output_quantizers: List[Optional[Quantizer]], + fuse_wgrad_accumulation: bool, + activation_dtype: torch.dtype, + is_grad_enabled: bool, + weight_workspaces: List[Optional[QuantizedTensorStorage]], + cache_weight: bool, + skip_fp8_weight_update: Optional[torch.Tensor], + weights: Tuple[torch.Tensor, ...], + biases: Tuple[torch.Tensor, ...], + ) -> Tuple[torch.Tensor, list]: + """Forward path backed by GroupedTensor + cublasLt grouped GEMM.""" + num_gemms = len(m_splits) + device = inp.device + in_features = weights[0].size(-1) + out_features = weights[0].size(0) + weight_requires_grad = weights[0].requires_grad + + split_sizes = m_splits.to(device=device) + base_split_offsets = tex.splits_to_offsets(split_sizes, 1) + + inp_view = inp.reshape(-1, in_features) + x = cast_if_needed(inp_view, activation_dtype) + if fp8: + input_quantizer = input_quantizers[0] + input_quantizer.set_usage( + rowwise=True, + columnwise=is_grad_enabled and weight_requires_grad, + ) + input_quantizer.optimize_for_gemm = True + grouped_x = tex.group_quantize(x, input_quantizer, num_gemms, split_sizes) + else: + grouped_x = _GroupedLinear._make_grouped_tensor( + x, + num_gemms=num_gemms, + split_sizes=split_sizes, + base_split_offsets=base_split_offsets, + last_dim=in_features, + dtype=activation_dtype, + ) + + columnwise_usage = is_grad_enabled and inp.requires_grad + weights_for_gemm, new_workspaces = _GroupedLinear._prepare_weights_for_grouped_tensor_gemm( + weights, + weight_quantizers, + weight_workspaces, + with_quantized_compute=fp8, + columnwise_usage=columnwise_usage, + activation_dtype=activation_dtype, + is_first_microbatch=is_first_microbatch, + skip_fp8_weight_update=skip_fp8_weight_update, + cache_weight=cache_weight, + ) + + out = torch.empty( + [x.size(0), out_features], + dtype=activation_dtype, + device=device, + ) + grouped_out = _GroupedLinear._make_grouped_tensor( + out, + num_gemms=num_gemms, + split_sizes=split_sizes, + base_split_offsets=base_split_offsets, + last_dim=out_features, + dtype=activation_dtype, + ) + + grouped_bias = None + if use_bias: + grouped_bias = _GroupedLinear._make_grouped_bias( + biases, + num_gemms=num_gemms, + out_features=out_features, + dtype=activation_dtype, + ) + + use_split_accumulator = _2X_ACC_FPROP + if fp8: + recipe = FP8GlobalStateManager.get_fp8_recipe() + if hasattr(recipe, "fp8_gemm_fprop"): + use_split_accumulator = recipe.fp8_gemm_fprop.use_split_accumulator + + general_grouped_gemm_for_grouped_tensor( + weights_for_gemm, + grouped_x, + grouped_out, + layout="TN", + bias=grouped_bias, + use_split_accumulator=use_split_accumulator, + ) + + if is_grad_enabled: + if weight_requires_grad: + if fp8: + grouped_x.rowwise_data = None + grouped_x.scale_inv = None + else: + grouped_x = None + + weights_to_save = weights_for_gemm if inp.requires_grad else [None] * num_gemms + tensors_to_save, tensor_objects = prepare_for_saving( + grouped_x, + *weights_to_save, + split_sizes, + base_split_offsets, + ) + ctx.save_for_backward(*tensors_to_save) + ctx.tensor_objects = tensor_objects + + ctx.use_grouped_tensor_path = True + ctx.weight_quantizers = weight_quantizers + ctx.weights_shape_0 = out_features + ctx.weights_shape_1 = in_features + ctx.grad_input_quantizers = grad_input_quantizers + ctx.grad_output_quantizers = grad_output_quantizers + ctx.grad_weight_quantizers = grad_weight_quantizers + ctx.weights_requires_grad = weight_requires_grad + if fuse_wgrad_accumulation and ctx.weights_requires_grad: + ctx.origin_weight_refs = [weakref.ref(w) for w in weights] + ctx.origin_weights_overwrite_main_grad = getattr( + weights[0], "overwrite_main_grad", False + ) + if hasattr(weights[0], "__fsdp_param__"): + ctx.main_grad_funcs = [weights[i].get_main_grad for i in range(num_gemms)] + else: + ctx.main_grad_funcs = [ + lambda j=i: weights[j].main_grad for i in range(num_gemms) + ] + ctx.device = device + ctx.m_splits = None + ctx.num_gemms = num_gemms + ctx.activation_dtype = activation_dtype + ctx.fp8 = fp8 + ctx.fp8_recipe = FP8GlobalStateManager.get_fp8_recipe() if fp8 else None + ctx.backward_override = None + ctx.fuse_wgrad_accumulation = fuse_wgrad_accumulation + ctx.cpu_offloading = False + ctx.is_first_microbatch = is_first_microbatch + ctx.use_bias = use_bias + ctx.inp_shape = inp.shape + ctx.requires_dgrad = inp.requires_grad + ctx.reduce_and_update_bwd_fp8_tensors = False + if ctx.fp8 and requires_grad(inp, weights[0], biases[0]): + ctx.reduce_and_update_bwd_fp8_tensors = ( + ctx.reduce_and_update_bwd_fp8_tensors + or FP8GlobalStateManager.is_first_fp8_module() + ) + ctx.wgrad_store = wgrad_store + ctx.debug = False + ctx.save_original_input = False + ctx.input_quantizers = input_quantizers + + return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces + # pylint: disable=keyword-arg-before-vararg @staticmethod def forward( ctx, inp: torch.Tensor, + m_splits: torch.Tensor, non_tensor_args: Tuple, *weights_and_biases, ) -> Tuple[torch.Tensor, list]: @@ -78,7 +389,6 @@ def forward( # Reduce number of arguments to autograd function in order # to reduce CPU overhead due to pytorch arg checking. ( - m_splits, use_bias, is_first_microbatch, fp8, @@ -168,6 +478,46 @@ def forward( f"Input tensor (shape={tuple(inp.size())}) is not compatible with " f"weight tensor (shape={tuple(weights[0].size())})" ) + + if _GroupedLinear._is_grouped_tensor_path_supported( + fp8=fp8, + fp8_calibration=fp8_calibration, + debug=debug, + cpu_offloading=cpu_offloading, + backward_override=backward_override, + save_original_input=save_original_input, + activation_dtype=activation_dtype, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + output_quantizers=output_quantizers, + grad_output_quantizers=grad_output_quantizers, + ): + return _GroupedLinear._forward_grouped_tensor( + ctx, + inp=inp, + m_splits=m_splits, + use_bias=use_bias, + is_first_microbatch=is_first_microbatch, + fp8=fp8, + wgrad_store=wgrad_store, + input_quantizers=input_quantizers, + weight_quantizers=weight_quantizers, + grad_input_quantizers=grad_input_quantizers, + grad_weight_quantizers=grad_weight_quantizers, + grad_output_quantizers=grad_output_quantizers, + fuse_wgrad_accumulation=fuse_wgrad_accumulation, + activation_dtype=activation_dtype, + is_grad_enabled=is_grad_enabled, + weight_workspaces=weight_workspaces, + cache_weight=cache_weight, + skip_fp8_weight_update=skip_fp8_weight_update, + weights=weights, + biases=biases, + ) + + # Convert splits to list of ints for compatibility with split functions + m_splits = m_splits.tolist() + inp_view = inp.reshape(-1, in_features) inputmats: list if fp8 and not debug: @@ -256,6 +606,7 @@ def forward( mark_not_offload(*weights_fp8, *weights) if is_grad_enabled: + ctx.use_grouped_tensor_path = False ctx.weight_quantizers = weight_quantizers ctx.weights_shape_1 = weights[0].shape[1] @@ -276,10 +627,18 @@ def forward( else: inputmats = [None] * num_gemms + # Original weights are only needed by high_precision dgrad. The weakrefs + # used for fused wgrad accumulation serve a different purpose: restoring + # Python parameter attributes without keeping the parameter alive here. + saved_weights = ( + weights + if backward_override == "high_precision" and inp.requires_grad + else [None] * num_gemms + ) tensors_to_save, tensor_objects = prepare_for_saving( *inputmats, *weights_fp8, - *weights, + *saved_weights, *biases, ) ctx.save_for_backward(*tensors_to_save) @@ -349,12 +708,201 @@ def forward( # [*, in_features] -> [*, out_features] except first dimension changes for SP return out.view(-1, *inp.shape[1:-1], out.shape[-1]), new_workspaces + @staticmethod + def _backward_grouped_tensor( + ctx, + grad_output: torch.Tensor, + ) -> Tuple[Union[torch.Tensor, None], ...]: + """Backward path paired with ``_forward_grouped_tensor``.""" + saved_tensors = restore_from_func_ctx(ctx) + N = ctx.num_gemms + grouped_x = saved_tensors[0] + weights = saved_tensors[1 : 1 + N] + split_sizes = saved_tensors[1 + N] + base_split_offsets = saved_tensors[2 + N] + + origin_weights = [None] * N + main_grads = [None] * N + if ctx.fuse_wgrad_accumulation and ctx.weights_requires_grad: + origin_weight_refs = ctx.origin_weight_refs + ctx.origin_weight_refs = None + origin_weights = [ref() if ref is not None else None for ref in origin_weight_refs] + assert all( + w is not None for w in origin_weights + ), "weight was removed while fuse_wgrad_accumulation=True" + main_grads = [main_grad_func() for main_grad_func in ctx.main_grad_funcs] + for origin_weight, main_grad in zip(origin_weights, main_grads): + if main_grad is not None: + origin_weight.main_grad = main_grad + + grad_output_view = grad_output.contiguous().view(-1, grad_output.shape[-1]) + dy_2d = cast_if_needed(grad_output_view, ctx.activation_dtype) + dbias_packed = None + if ctx.fp8: + grad_output_quantizer = ctx.grad_output_quantizers[0] + grad_output_quantizer.set_usage( + rowwise=ctx.requires_dgrad, + columnwise=ctx.weights_requires_grad, + ) + grad_output_quantizer.optimize_for_gemm = True + if ctx.use_bias: + grouped_dy, dbias_packed = tex.bgrad_group_quantize( + dy_2d, + grad_output_quantizer, + N, + split_sizes, + ) + else: + grouped_dy = tex.group_quantize( + dy_2d, + grad_output_quantizer, + N, + split_sizes, + ) + else: + grouped_dy = _GroupedLinear._make_grouped_tensor( + dy_2d, + num_gemms=N, + split_sizes=split_sizes, + base_split_offsets=base_split_offsets, + last_dim=ctx.weights_shape_0, + dtype=ctx.activation_dtype, + ) + + grad_biases = [None] * N + if ctx.use_bias: + if dbias_packed is None: + dbias_packed = compute_grouped_dbias(dy_2d, base_split_offsets, N) + grad_biases = [dbias_packed[i].to(dtype=ctx.activation_dtype) for i in range(N)] + + dgrad = None + if ctx.requires_dgrad: + dgrad_gemm_use_split_accumulator = _2X_ACC_DGRAD + if ctx.fp8: + recipe = ctx.fp8_recipe + if hasattr(recipe, "fp8_gemm_dgrad"): + dgrad_gemm_use_split_accumulator = recipe.fp8_gemm_dgrad.use_split_accumulator + for weight in weights: + if isinstance(weight, QuantizedTensorStorage): + weight.update_usage(columnwise_usage=True) + dgrad = torch.empty( + (dy_2d.size(0), ctx.weights_shape_1), + dtype=ctx.activation_dtype, + device=ctx.device, + ) + grouped_dgrad = _GroupedLinear._make_grouped_tensor( + dgrad, + num_gemms=N, + split_sizes=split_sizes, + base_split_offsets=base_split_offsets, + last_dim=ctx.weights_shape_1, + dtype=ctx.activation_dtype, + ) + general_grouped_gemm_for_grouped_tensor( + weights, + grouped_dy, + grouped_dgrad, + layout="NN", + use_split_accumulator=dgrad_gemm_use_split_accumulator, + ) + + if ctx.is_first_microbatch is not None: + accumulate_wgrad_into_param_main_grad = ( + ctx.fuse_wgrad_accumulation and not ctx.is_first_microbatch + ) + else: + accumulate_wgrad_into_param_main_grad = ctx.fuse_wgrad_accumulation + + if ctx.weights_requires_grad: + wgrad_gemm_use_split_accumulator = _2X_ACC_WGRAD + if ctx.fp8: + recipe = ctx.fp8_recipe + if hasattr(recipe, "fp8_gemm_wgrad"): + wgrad_gemm_use_split_accumulator = recipe.fp8_gemm_wgrad.use_split_accumulator + if ctx.fuse_wgrad_accumulation: + wgrad_list = main_grads + else: + wgrad_packed = torch.empty( + N, + ctx.weights_shape_0, + ctx.weights_shape_1, + dtype=ctx.activation_dtype, + device=ctx.device, + ) + wgrad_list = [wgrad_packed[i] for i in range(N)] + + accumulate = ( + accumulate_wgrad_into_param_main_grad + if not getattr(ctx, "origin_weights_overwrite_main_grad", False) + else False + ) + + def grouped_gemm_wgrad(inputmats, grad_output_mats, grad_weights): + general_grouped_gemm_for_grouped_tensor( + inputmats, + grad_output_mats, + grad_weights, + layout="NT", + use_split_accumulator=wgrad_gemm_use_split_accumulator, + accumulate=accumulate, + ) + return None, [None] * N, None + + if ctx.wgrad_store is not None and ctx.wgrad_store.delay_wgrad_compute(): + ctx.wgrad_store.put([grouped_x, grouped_dy, wgrad_list], grouped_gemm_wgrad) + else: + grouped_gemm_wgrad(grouped_x, grouped_dy, wgrad_list) + + def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): + if ctx.weights_requires_grad: + if ctx.fuse_wgrad_accumulation and hasattr(weight, "grad_added_to_main_grad"): + weight.grad_added_to_main_grad = True + if getattr(weight, "zero_out_wgrad", False): + wgrad = get_dummy_wgrad( + list(main_grad.shape), + weight.dtype, + zero=True, + ) + else: + wgrad = get_dummy_wgrad( + list(main_grad.shape), + weight.dtype, + ) + elif ctx.fuse_wgrad_accumulation: + wgrad = None + else: + wgrad = None + return wgrad + + wgrad_list = [ + handle_custom_ddp_from_mcore(weight, main_grad, wgrad) + for weight, main_grad, wgrad in zip(origin_weights, main_grads, wgrad_list) + ] + else: + wgrad_list = [None] * N + + if not ctx.use_bias: + grad_biases = [None] * N + + if ctx.reduce_and_update_bwd_fp8_tensors: + FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) + return ( + dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, + None, # m_splits + None, # non_tensor_args + *wgrad_list, + *grad_biases, + ) + @staticmethod def backward( ctx, grad_output: torch.Tensor, _grad_workspaces ) -> Tuple[Union[torch.Tensor, None], ...]: # pylint: disable=missing-function-docstring with get_nvtx_range_context("_GroupedLinear_backward"): + if ctx.use_grouped_tensor_path: + return _GroupedLinear._backward_grouped_tensor(ctx, grad_output) + saved_tensors = restore_from_func_ctx(ctx) N = ctx.num_gemms inputmats = saved_tensors[:N] @@ -613,7 +1161,8 @@ def handle_custom_ddp_from_mcore(weight, main_grad, wgrad): FP8GlobalStateManager.reduce_and_update_fp8_tensors(forward=False) return ( dgrad.view(ctx.inp_shape) if ctx.requires_dgrad else None, - None, + None, # m_splits + None, # non_tensor_args *wgrad_list, *grad_biases, ) @@ -1105,7 +1654,7 @@ def _load_from_state_dict( def forward( self, inp: torch.Tensor, - m_splits: List[int], + m_splits: torch.Tensor, is_first_microbatch: Optional[bool] = None, ) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: """ @@ -1115,8 +1664,8 @@ def forward( ---------- inp : torch.Tensor Input tensor. - m_splits : List[int] - List of integers representing the split of the input tensor. + m_splits : torch.Tensor + Split sizes for the input tensor. is_first_microbatch : {True, False, None}, default = None During training using either gradient accumulation or pipeline parallelism a minibatch of data is further split @@ -1132,18 +1681,26 @@ def forward( produced) """ debug = self.is_debug_iter() - - if isinstance(inp, QuantizedTensorStorage): - raise TypeError("GroupedLinear doesn't support input tensor in FP8.") - if len(m_splits) != self.num_gemms: + is_grad_enabled = torch.is_grad_enabled() + num_gemms = self.num_gemms + + # Make sure splits are in expected format + if not isinstance(m_splits, torch.Tensor): + # Convert list of ints to tensor for backward compatibility + m_splits = torch.tensor(m_splits, dtype=torch.int64, device="cpu") + elif m_splits.dtype != torch.int64: + m_splits = m_splits.to(dtype=torch.int64) + if m_splits.size() != (num_gemms,): raise ValueError( - f"Number of splits ({len(m_splits)}) should match number of" - f" GEMMs ({self.num_gemms})." + f"Shape of splits tensor ({tuple(m_splits.size())}) " + f"does not match number of GEMMs ({num_gemms})." ) - is_grad_enabled = torch.is_grad_enabled() - + # Preprocess input tensor + if isinstance(inp, QuantizedTensorStorage): + raise TypeError("GroupedLinear doesn't support input tensor in FP8.") inp = self.prepare_forward(inp, num_gemms=self.num_gemms) + try: weight_tensors = self._get_weight_tensors() bias_tensors = self._get_bias_tensors() @@ -1171,7 +1728,6 @@ def forward( linear_fn = _GroupedLinear.forward autograd_ctx = [None] - num_gemms = len(m_splits) cache_weight = is_first_microbatch is not None weight_workspaces = ( [self._fp8_workspaces.get(f"weight{i}") for i in range(num_gemms)] @@ -1180,7 +1736,6 @@ def forward( ) non_tensor_args = ( - m_splits, self.apply_bias, is_first_microbatch, self.fp8, @@ -1204,7 +1759,7 @@ def forward( debug, ) out, new_workspaces = linear_fn( - *autograd_ctx, inp, non_tensor_args, *weight_tensors, *bias_tensors + *autograd_ctx, inp, m_splits, non_tensor_args, *weight_tensors, *bias_tensors ) if cache_weight: @@ -1237,9 +1792,14 @@ def backward_dw(self): if not self.fuse_wgrad_accumulation: for i in range(self.num_gemms): weight_params[i].grad = wgrad_list[i].to(weight_params[i].dtype) - if self.use_bias: + has_grad_biases = [ + grad_bias is not None and grad_bias.numel() != 0 for grad_bias in grad_biases_ + ] + if self.use_bias and any(has_grad_biases): grouped_bias = getattr(self, "bias", None) if grouped_bias is not None: + if not all(has_grad_biases): + raise RuntimeError("Expected all grouped bias gradients to be present.") gstack = torch.stack(grad_biases_, dim=0).to(grouped_bias.dtype) if grouped_bias.grad is None: grouped_bias.grad = gstack @@ -1248,7 +1808,7 @@ def backward_dw(self): else: bias_params = [getattr(self, f"bias{i}") for i in range(self.num_gemms)] for i in range(self.num_gemms): - if bias_params[i].grad is None: + if has_grad_biases[i] and bias_params[i].grad is None: bias_params[i].grad = grad_biases_[i].to(bias_params[i].dtype) del grad_biases_ del wgrad_list From af5d1e0d43aa626f56240cbeb750a07a7da6b1f8 Mon Sep 17 00:00:00 2001 From: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Date: Fri, 29 May 2026 10:55:54 -0700 Subject: [PATCH 091/180] [JAX] Fix L0_jax_unittest docs example test to enforce single-GPU (#3059) Update test.sh Signed-off-by: jberchtold-nvidia <158520091+jberchtold-nvidia@users.noreply.github.com> Co-authored-by: Teddy Do --- qa/L0_jax_unittest/test.sh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/qa/L0_jax_unittest/test.sh b/qa/L0_jax_unittest/test.sh index e4bcdc4e5..12bb027b9 100644 --- a/qa/L0_jax_unittest/test.sh +++ b/qa/L0_jax_unittest/test.sh @@ -45,7 +45,7 @@ NVTE_JAX_CUSTOM_CALLS="false" python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini # Exercise the docs/examples/jax tutorials. The multi-GPU tests are # skipped at runtime when fewer than 4 devices are visible, so this is safe on # single-GPU runners. -python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_docs_examples_jax.xml $TE_PATH/docs/examples/jax/ || test_fail "docs/examples/jax" +CUDA_VISIBLE_DEVICES=0 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_docs_examples_jax.xml $TE_PATH/docs/examples/jax/ || test_fail "docs/examples/jax" if [ $RET -ne 0 ]; then echo "Error: some sub-tests failed: $FAILED_CASES" From d1920cf524f224fa9e70642588659eaba87949b9 Mon Sep 17 00:00:00 2001 From: Teddy Do Date: Fri, 29 May 2026 11:07:21 -0700 Subject: [PATCH 092/180] [JAX] Add an MoE Block (Layer) that compound router, permutation, groupedGEMM and communication (#2912) Refactor MoEBlock into a unified MoE custom_vjp, add tests Replace the per-primitive custom_vjp boundaries in MoEBlock with a single jax.custom_vjp covering routing, dispatch, expert FFN, and combine. Helper functions group permute -> ragged_all_to_all -> local-permute into a single dispatch / combine pair, with a hand- derived bwd that mirrors the forward and runs entirely inside the EP shard_map body. Add a multi-process (one-GPU-per-process) test suite for the new unified VJP under a 2x2 (ep, fsdp) mesh: * tests/jax/test_multiprocess_moe_vjp.py -- fwd/bwd + aux_loss + PURE_JAX vs TRITON parity at Mixtral-ish shapes (batch=16, seq=2048, hidden=1024, intermediate=4096, num_experts=8, topk=2). * tests/jax/run_multiprocess_moe_vjp.sh -- launcher; forks one pytest process per visible GPU (mirrors examples/jax/encoder/run_test_multiprocessing_encoder.sh). * tests/jax/conftest.py -- pytest --num-process / --process-id options for the launcher. * qa/L0_jax_distributed_unittest/test.sh -- CI hook for the multiprocess smoke. Signed-off-by: tdophung * [JAX] Fix EP+TRITON combine bwd: save post-A2A expert_outputs Under EP, _combine reassigns expert_outputs locally to the post- ragged_all_to_all tensor before Step 3 (the global combine). Saving the input (pre-A2A) tensor in ctx.expert_outputs meant _combine_bwd's Step-3 inverse (unpermute_bwd_with_merging_probs) consumed a tensor with the wrong shape and contents, silently corrupting d_expert_outputs. _combine now returns (output, expert_outputs_post_ep). The caller stashes the second value as the bwd residual so the Step-3 inverse sees the same tensor the forward Step 3 saw. Signed-off-by: Teddy Do Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- qa/L0_jax_distributed_unittest/test.sh | 6 + tests/jax/conftest.py | 14 + tests/jax/run_multiprocess_moe_vjp.sh | 132 + tests/jax/test_moe_vjp.py | 443 ++++ tests/jax/test_multiprocess_moe_vjp.py | 406 ++++ .../common/triton/permutation.py | 50 +- transformer_engine/jax/flax/__init__.py | 2 + transformer_engine/jax/flax/moe.py | 284 +++ transformer_engine/jax/moe.py | 2165 +++++++++++++++++ transformer_engine/jax/permutation.py | 720 +++++- transformer_engine/jax/sharding.py | 34 + 11 files changed, 4208 insertions(+), 48 deletions(-) create mode 100755 tests/jax/run_multiprocess_moe_vjp.sh create mode 100644 tests/jax/test_moe_vjp.py create mode 100644 tests/jax/test_multiprocess_moe_vjp.py create mode 100644 transformer_engine/jax/flax/moe.py create mode 100644 transformer_engine/jax/moe.py diff --git a/qa/L0_jax_distributed_unittest/test.sh b/qa/L0_jax_distributed_unittest/test.sh index c62d7a4ba..bf4652c31 100644 --- a/qa/L0_jax_distributed_unittest/test.sh +++ b/qa/L0_jax_distributed_unittest/test.sh @@ -37,6 +37,12 @@ wait TE_PATH=$TE_PATH bash $TE_PATH/examples/jax/collective_gemm/run_test_cgemm.sh || test_fail "run_test_cgemm.sh" wait +# MoE custom_vjp distributed suite. Runs one Python process per GPU +# via tests/jax/run_multiprocess_moe_vjp.sh (mirrors the pattern in +# examples/jax/encoder/run_test_multiprocessing_encoder.sh). Requires +# >=4 visible GPUs. +TE_PATH=$TE_PATH bash $TE_PATH/tests/jax/run_multiprocess_moe_vjp.sh \ + || test_fail "test_multiprocess_moe_vjp.py" # Exercise the multi-GPU tutorial in docs/examples/jax (needs >= 4 GPUs; # auto-skips otherwise). CUDA_VISIBLE_DEVICES=0,1,2,3 python3 -m pytest -c $TE_PATH/tests/jax/pytest.ini -v --junitxml=$XML_LOG_DIR/pytest_docs_examples_jax_distributed.xml -k multi_gpu $TE_PATH/docs/examples/jax/ || test_fail "docs/examples/jax (multi-GPU)" diff --git a/tests/jax/conftest.py b/tests/jax/conftest.py index db30f0ed3..74cb91202 100644 --- a/tests/jax/conftest.py +++ b/tests/jax/conftest.py @@ -86,6 +86,20 @@ def pytest_sessionfinish(self, session, exitstatus): print("=" * 80) +def pytest_addoption(parser): + """CLI options used by multiprocess JAX tests. + + ``--num-process`` and ``--process-id`` let a multiprocess launcher + (see ``tests/jax/run_multiprocess_moe_vjp.sh``) fork one pytest + process per GPU and tell each child its rank, so the test module + can call ``jax.distributed.initialize(...)`` with the right + ``local_device_ids``. Both default to 0; non-multiprocess tests + ignore them. + """ + parser.addoption("--num-process", action="store", default=0) + parser.addoption("--process-id", action="store", default=0) + + def pytest_configure(config): config.addinivalue_line( "markers", diff --git a/tests/jax/run_multiprocess_moe_vjp.sh b/tests/jax/run_multiprocess_moe_vjp.sh new file mode 100755 index 000000000..8dc1d2eb0 --- /dev/null +++ b/tests/jax/run_multiprocess_moe_vjp.sh @@ -0,0 +1,132 @@ +#!/usr/bin/env bash +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +# +# Multiprocess (one-GPU-per-process) launcher for the unified MoE VJP +# test suite. Forks one pytest invocation per visible GPU, passing each +# its own --num-process=N --process-id=i, and waits for all of them. +# Each child calls jax.distributed.initialize(..., local_device_ids= +# process_id) so each Python process only sees its one GPU as a local +# device and the participating processes form a global mesh. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TE_ROOT="$(cd "$SCRIPT_DIR/../.." && pwd)" +TEST_FILE="$TE_ROOT/tests/jax/test_multiprocess_moe_vjp.py" +PYTEST_INI="$TE_ROOT/tests/jax/pytest.ini" + +NUM_GPUS="${NUM_GPUS:-$(nvidia-smi -L | wc -l)}" +if [ "$NUM_GPUS" -lt 4 ]; then + echo "[run_multiprocess_moe_vjp.sh] need >=4 GPUs (got $NUM_GPUS); aborting" >&2 + exit 1 +fi + +export XLA_PYTHON_CLIENT_PREALLOCATE="${XLA_PYTHON_CLIENT_PREALLOCATE:-false}" +export XLA_PYTHON_CLIENT_MEM_FRACTION="${XLA_PYTHON_CLIENT_MEM_FRACTION:-0.5}" +export MOE_VJP_COORDINATOR_ADDRESS="${MOE_VJP_COORDINATOR_ADDRESS:-127.0.0.1:13456}" + +echo "============================================================" +echo "MoE VJP MULTIPROCESS test (one process per GPU, ${NUM_GPUS} GPUs)" +echo " test file : $TEST_FILE" +echo " coordinator : $MOE_VJP_COORDINATOR_ADDRESS" +echo " XLA_PYTHON_CLIENT_PREALLOCATE: $XLA_PYTHON_CLIENT_PREALLOCATE" +echo " XLA_PYTHON_CLIENT_MEM_FRACTION: $XLA_PYTHON_CLIENT_MEM_FRACTION" +echo "============================================================" + +# Per-process logs. MOE_VJP_MP_LOG_DIR can be set to a host-mounted dir +# (e.g. when running inside a container that throws away /tmp on exit) +# so logs survive for postmortem inspection. Defaults to a fresh /tmp. +if [ -n "${MOE_VJP_MP_LOG_DIR:-}" ]; then + LOG_DIR="$MOE_VJP_MP_LOG_DIR" + mkdir -p "$LOG_DIR" +else + LOG_DIR=$(mktemp -d -t moe_vjp_mp_XXXXXX) +fi +echo "Per-process logs: $LOG_DIR" + +PIDS=() + +cleanup() { + for pid in "${PIDS[@]:-}"; do + if kill -0 "$pid" 2>/dev/null; then + kill -TERM "$pid" 2>/dev/null || true + fi + done + sleep 1 + for pid in "${PIDS[@]:-}"; do + if kill -0 "$pid" 2>/dev/null; then + kill -KILL "$pid" 2>/dev/null || true + fi + done +} +trap cleanup EXIT INT TERM + +# Launch one pytest per GPU. Process 0 streams to stdout; others log +# only to file so the live output isn't a mosaic. +for i in $(seq 0 $((NUM_GPUS - 1))); do + LOG_FILE="$LOG_DIR/proc_${i}.log" + PYTEST_CMD=( + python3 -m pytest -c "$PYTEST_INI" + "$TEST_FILE" + -p no:typeguard + -v -s + --num-process="$NUM_GPUS" + --process-id="$i" + ) + if [ "$i" -eq 0 ]; then + echo "=== Live output from process 0 ===" + "${PYTEST_CMD[@]}" 2>&1 | tee "$LOG_FILE" & + else + "${PYTEST_CMD[@]}" > "$LOG_FILE" 2>&1 & + fi + PIDS+=("$!") +done + +# Wait for all and collect exit codes. +EXITS=() +for pid in "${PIDS[@]}"; do + if wait "$pid"; then + EXITS+=("0") + else + EXITS+=("$?") + fi +done + +# Summary. +echo +echo "============================================================" +echo "Per-process exit codes:" +for i in "${!EXITS[@]}"; do + echo " proc $i -> ${EXITS[$i]}" +done + +# Final pass/fail. Any non-zero in any process fails the suite, but +# we tolerate non-zero on the non-zero processes only if proc 0 +# reports PASS (this matches the encoder launcher's logic). Simplest +# Treat exit 0 (pass) and exit 5 (pytest "no tests collected", which +# the file emits via ``pytest.skip(allow_module_level=True)`` on +# pre-Blackwell GPUs) as success. Anything else is a failure. +FAILED=0 +for e in "${EXITS[@]}"; do + if [ "$e" != "0" ] && [ "$e" != "5" ]; then + FAILED=1 + break + fi +done + +echo +if [ "$FAILED" -eq 0 ]; then + echo "[run_multiprocess_moe_vjp.sh] all processes PASSED" + if [ -z "${MOE_VJP_MP_LOG_DIR:-}" ]; then + rm -rf "$LOG_DIR" + fi + exit 0 +fi + +echo "[run_multiprocess_moe_vjp.sh] at least one process FAILED" +echo " retaining logs at $LOG_DIR for diagnosis" +echo " process 0 tail:" +tail -20 "$LOG_DIR/proc_0.log" 2>/dev/null || true +exit 1 diff --git a/tests/jax/test_moe_vjp.py b/tests/jax/test_moe_vjp.py new file mode 100644 index 000000000..cc458d039 --- /dev/null +++ b/tests/jax/test_moe_vjp.py @@ -0,0 +1,443 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Single-device tests for the unified MoE custom_vjp at +``transformer_engine.jax.moe.moe`` (and its Flax wrapper +``transformer_engine.jax.flax._MoEBlock``). + +Strategy +-------- + +Rather than reproducing every internal kernel residual, we rely on a +single end-to-end pure-JAX *reference* implementation of the whole +MoE block (``_pure_jax_moe_reference`` below) and compare the TE +``moe(...)`` forward output AND parameter gradients against it. This +gives us coverage of: + +* the gate GEMM, +* the fused top-k routing primitive (and its bwd), +* the dispatch / per-expert FFN / combine pipeline (and their bwds + threaded through the absorbed primitives), +* the optional aux-loss path (and its bwd). + +The reference uses only ``jnp`` ops + ``jax.vjp``, so we get a +"definitive" pullback to compare against without needing the TE +primitive bwd kernels. + +Distributed (EP + FSDP) testing is intentionally NOT in this file -- +that needs a multi-device setup and lives in +``tests/jax/test_distributed_moe_vjp.py`` (follow-up). +""" + +from functools import partial +from typing import Optional, Tuple + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from transformer_engine_jax import get_device_compute_capability +from transformer_engine.jax.flax import _MoEBlock as MoEBlock +from transformer_engine.jax.moe import PermutationBackend, moe + +# The MoE custom_vjp uses grouped GEMM, which is currently +# Blackwell-only (sm_100+). Skip the whole file on older arches. +if get_device_compute_capability(0) < 100: + pytest.skip( + "MoE custom_vjp tests require Blackwell (sm_100+) for grouped GEMM", + allow_module_level=True, + ) + +# Parametrize values for the dispatch / combine backend. Only the +# ``triton`` variant is gated by the ``triton`` marker (so the +# ``pure_jax`` variant still runs on environments without Triton). +BACKEND_PARAMS = [ + pytest.param("pure_jax", id="pure_jax"), + pytest.param("triton", id="triton", marks=pytest.mark.triton), +] + + +# ----------------------------------------------------------------------------- +# Test config +# ----------------------------------------------------------------------------- + +DTYPE = jnp.float32 # use fp32 for tighter parity assertions +BATCH_SIZE = 2 +SEQUENCE_LENGTH = 16 +HIDDEN_SIZE = 32 +INTERMEDIATE_SIZE = 64 +NUM_EXPERTS = 8 +NUM_EXPERTS_PER_TOK = 2 + + +def _make_inputs(key: jax.Array, *, batch=BATCH_SIZE, seq=SEQUENCE_LENGTH) -> jax.Array: + return jax.random.normal(key, (batch, seq, HIDDEN_SIZE), dtype=DTYPE) + + +# ----------------------------------------------------------------------------- +# Pure-JAX reference MoE +# ----------------------------------------------------------------------------- +# +# Implements EXACTLY the same math as ``moe(...)`` for the no-EP, +# softmax-routing, no-bias, silu activation, no-quantization path. +# Returns ``(output, aux_loss_or_zero)``. Used as ground truth for both +# fwd and bwd parity. + + +@partial( + jax.jit, + static_argnames=("num_experts", "num_experts_per_tok", "aux_loss_coeff"), +) +def _pure_jax_moe_reference( + x: jnp.ndarray, + gate_kernel: jnp.ndarray, + wi_0: jnp.ndarray, + wi_1: jnp.ndarray, + wo: jnp.ndarray, + *, + num_experts: int, + num_experts_per_tok: int, + aux_loss_coeff: float = 0.0, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Reference no-EP MoE forward (pure JAX, no TE primitives). + + Mirrors :func:`transformer_engine.jax.moe._body_fwd` for the + PURE_JAX backend, no biases, softmax routing, silu activation, + no quantization. Linear ops only -- ``jax.vjp`` over this gives + the canonical bwd to compare against. + """ + B, S, H = x.shape + T = B * S + x_2d = x.reshape(T, H) + + # Gate + logits = x_2d @ gate_kernel # [T, E] + + # Softmax + topk (no expert_bias, no grouping, scale=1.0) + probs_full = jax.nn.softmax(logits, axis=-1) # [T, E] + # top-k by probability: + sorted_idx = jnp.argsort(probs_full, axis=-1) # ascending + selected = sorted_idx[:, -num_experts_per_tok:] # [T, K] + weights = jnp.take_along_axis(probs_full, selected, axis=-1) # [T, K] + # Normalize topk weights to sum to 1 (matches softmax->topk semantics + # of fused_topk_with_score_function with use_pre_softmax=False): + weights = weights / jnp.sum(weights, axis=-1, keepdims=True) + + # Build a sparse routing_map [T, E] with weights at selected positions + routing_weights_full = jnp.zeros_like(probs_full) + routing_weights_full = routing_weights_full.at[jnp.arange(T)[:, None], selected].set(weights) + + # Per-expert FFN: replicate each token K times, gather by expert, + # run through wi_0 / wi_1 / wo, gather back, weighted-sum. + # + # Vectorize the gather without sorting: for each (token, slot k), + # multiply the corresponding expert's FFN by routing_weights[t, k] + # and sum over experts. + # x_2d: [T, H], wi_0: [E, H, M], wi_1: [E, H, M], wo: [E, M, H] + # For each expert e: layer_w0_e = x_2d @ wi_0[e]; layer_w1_e = x_2d @ wi_1[e] + # intermediate_e = silu(layer_w0_e) * layer_w1_e + # expert_out_e = intermediate_e @ wo[e] + # output[t, h] = sum_e routing_weights_full[t, e] * expert_out_e[t, h] + layer_w0 = jnp.einsum("th,ehm->tem", x_2d, wi_0) # [T, E, M] + layer_w1 = jnp.einsum("th,ehm->tem", x_2d, wi_1) # [T, E, M] + intermediate = jax.nn.silu(layer_w0) * layer_w1 # [T, E, M] + expert_out = jnp.einsum("tem,emh->teh", intermediate, wo) # [T, E, H] + output_2d = jnp.einsum("te,teh->th", routing_weights_full, expert_out) # [T, H] + output = output_2d.reshape(B, S, H) + + if aux_loss_coeff > 0.0: + # aux scores: clean per-expert softmax (compute_aux_scores=True + # kernel uses a clean softmax, no bias, scale=1, no grouping). + aux_probs = jax.nn.softmax(logits.astype(jnp.float32), axis=-1) + # tokens_per_expert from REAL routing_map (post-grouping); here + # there's no grouping so == count of non-zero positions per expert. + routing_map = (routing_weights_full > 0).astype(jnp.int32) + tokens_per_expert = jnp.sum(routing_map, axis=0) # [E] + # aux_loss formula: (E * coeff / (k * T^2)) * sum_e + # (sum_t aux_probs[t, e]) * tokens_per_expert[e] + sum_probs_per_expert = jnp.sum(aux_probs, axis=0) # [E] + aux_loss = (num_experts * aux_loss_coeff / (num_experts_per_tok * (T**2))) * jnp.sum( + sum_probs_per_expert * tokens_per_expert.astype(jnp.float32) + ) + else: + aux_loss = jnp.zeros((), dtype=DTYPE) + + return output, aux_loss + + +# ----------------------------------------------------------------------------- +# Helpers +# ----------------------------------------------------------------------------- + + +def _init_params(key: jax.Array) -> dict: + k_g, k_w0, k_w1, k_wo = jax.random.split(key, 4) + init = jax.nn.initializers.variance_scaling(1.0, "fan_in", "truncated_normal") + return dict( + gate_kernel=init(k_g, (HIDDEN_SIZE, NUM_EXPERTS), DTYPE), + wi_0=init(k_w0, (NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE), DTYPE), + wi_1=init(k_w1, (NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE), DTYPE), + wo=init(k_wo, (NUM_EXPERTS, INTERMEDIATE_SIZE, HIDDEN_SIZE), DTYPE), + ) + + +@partial(jax.jit, static_argnames=("permutation_backend", "aux_loss_coeff")) +def _run_te_moe( + x: jnp.ndarray, + params: dict, + *, + permutation_backend, + aux_loss_coeff: float = 0.0, +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: + return moe( + x, + params["gate_kernel"], + params["wi_0"], + params["wi_1"], + params["wo"], + num_experts=NUM_EXPERTS, + num_experts_per_tok=NUM_EXPERTS_PER_TOK, + activation_type="silu", + score_function="softmax", + use_pre_softmax=False, + scaling_factor=1.0, + aux_loss_coeff=aux_loss_coeff, + permutation_backend=permutation_backend, + align_size=0, + dtype=DTYPE, + ) + + +@partial(jax.jit, static_argnames=("permutation_backend", "aux_loss_coeff")) +def _grads_te_main_loss(params, x, *, permutation_backend, aux_loss_coeff: float = 0.0): + """jit'd grad of ``mean(out**2)`` w.r.t. params (no aux contribution).""" + + def loss(params, x): + out, _ = _run_te_moe( + x, params, permutation_backend=permutation_backend, aux_loss_coeff=aux_loss_coeff + ) + return jnp.mean(out**2) + + return jax.grad(loss)(params, x) + + +@partial(jax.jit, static_argnames=("num_experts", "num_experts_per_tok", "aux_loss_coeff")) +def _grads_ref_main_loss(params, x, *, num_experts, num_experts_per_tok, aux_loss_coeff=0.0): + """jit'd grad of ``mean(out**2)`` w.r.t. params on the pure-JAX ref.""" + + def loss(params, x): + out, _ = _pure_jax_moe_reference( + x, + **params, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + aux_loss_coeff=aux_loss_coeff, + ) + return jnp.mean(out**2) + + return jax.grad(loss)(params, x) + + +@partial(jax.jit, static_argnames=("permutation_backend",)) +def _grad_te_aux_only(params, x, *, permutation_backend): + """jit'd grad of just the aux loss scalar (no main contribution).""" + + def aux_only(params, x): + _, aux = _run_te_moe( + x, params, permutation_backend=permutation_backend, aux_loss_coeff=1e-2 + ) + return aux.astype(jnp.float32) + + return jax.grad(aux_only)(params, x) + + +# ----------------------------------------------------------------------------- +# Tests +# ----------------------------------------------------------------------------- + + +class TestMoeVjpForward: + """Forward shape / finiteness / parity vs pure-JAX reference.""" + + @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) + def test_forward_shape_and_finite(self, backend_name): + backend = PermutationBackend(backend_name) + key = jax.random.PRNGKey(0) + kp, kx = jax.random.split(key) + params = _init_params(kp) + x = _make_inputs(kx) + out, aux = _run_te_moe(x, params, permutation_backend=backend) + assert out.shape == x.shape + assert out.dtype == x.dtype + assert jnp.all(jnp.isfinite(out)) + assert aux is None + + @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) + def test_forward_parity_vs_pure_jax_reference(self, backend_name): + backend = PermutationBackend(backend_name) + key = jax.random.PRNGKey(1) + kp, kx = jax.random.split(key) + params = _init_params(kp) + x = _make_inputs(kx) + out_te, _ = _run_te_moe(x, params, permutation_backend=backend) + out_ref, _ = _pure_jax_moe_reference( + x, + **params, + num_experts=NUM_EXPERTS, + num_experts_per_tok=NUM_EXPERTS_PER_TOK, + ) + # FP32, small shapes -> tight tolerance + np.testing.assert_allclose(np.array(out_te), np.array(out_ref), atol=2e-5, rtol=2e-5) + + def test_pure_jax_triton_equivalence(self): + key = jax.random.PRNGKey(2) + kp, kx = jax.random.split(key) + params = _init_params(kp) + x = _make_inputs(kx) + out_pj, _ = _run_te_moe(x, params, permutation_backend=PermutationBackend.PURE_JAX) + out_tr, _ = _run_te_moe(x, params, permutation_backend=PermutationBackend.TRITON) + np.testing.assert_allclose(np.array(out_pj), np.array(out_tr), atol=2e-5, rtol=2e-5) + + +class TestMoeVjpBackward: + """Backward parity vs pure-JAX reference (which uses ``jax.vjp`` over + plain JAX ops, giving us the canonical pullback).""" + + @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) + def test_grads_finite_and_nonzero(self, backend_name): + backend = PermutationBackend(backend_name) + key = jax.random.PRNGKey(3) + kp, kx = jax.random.split(key) + params = _init_params(kp) + x = _make_inputs(kx) + grads = _grads_te_main_loss(params, x, permutation_backend=backend) + for name in ("gate_kernel", "wi_0", "wi_1", "wo"): + g = grads[name] + assert jnp.all(jnp.isfinite(g)), f"{name} grad has NaN/Inf" + assert jnp.any(g != 0.0), f"{name} grad is identically zero" + + @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) + def test_grads_match_pure_jax_reference(self, backend_name): + backend = PermutationBackend(backend_name) + key = jax.random.PRNGKey(4) + kp, kx = jax.random.split(key) + params = _init_params(kp) + x = _make_inputs(kx) + grads_te = _grads_te_main_loss(params, x, permutation_backend=backend) + grads_ref = _grads_ref_main_loss( + params, + x, + num_experts=NUM_EXPERTS, + num_experts_per_tok=NUM_EXPERTS_PER_TOK, + ) + # Loose-ish tol on grads: routing path has discrete topk so the + # softmax cotangent paths through the non-topk experts diverge + # slightly between TE (which uses the fused topk bwd) and the + # reference (which uses argsort-based take_along_axis). + # Tighter than the bf16 tests. + for name in ("wi_0", "wi_1", "wo"): + np.testing.assert_allclose( + np.array(grads_te[name]), + np.array(grads_ref[name]), + atol=5e-5, + rtol=5e-5, + err_msg=f"grad mismatch on {name}", + ) + # Gate grad has more error budget because it propagates through + # the topk derivative kernel (which differs in zero-pattern + # treatment from a plain take_along_axis). + np.testing.assert_allclose( + np.array(grads_te["gate_kernel"]), + np.array(grads_ref["gate_kernel"]), + atol=5e-4, + rtol=5e-4, + err_msg="grad mismatch on gate_kernel", + ) + + +class TestMoeVjpAuxLoss: + """Aux-loss path: forward + grad parity.""" + + @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) + def test_aux_loss_returned_and_finite(self, backend_name): + backend = PermutationBackend(backend_name) + key = jax.random.PRNGKey(5) + kp, kx = jax.random.split(key) + params = _init_params(kp) + x = _make_inputs(kx) + _, aux = _run_te_moe(x, params, permutation_backend=backend, aux_loss_coeff=1e-2) + assert aux is not None + assert aux.shape == () + assert jnp.isfinite(aux) + assert jnp.abs(aux) < 1e2 + + @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) + def test_aux_loss_parity_vs_reference(self, backend_name): + backend = PermutationBackend(backend_name) + key = jax.random.PRNGKey(6) + kp, kx = jax.random.split(key) + params = _init_params(kp) + x = _make_inputs(kx) + _, aux_te = _run_te_moe(x, params, permutation_backend=backend, aux_loss_coeff=1e-2) + _, aux_ref = _pure_jax_moe_reference( + x, + **params, + num_experts=NUM_EXPERTS, + num_experts_per_tok=NUM_EXPERTS_PER_TOK, + aux_loss_coeff=1e-2, + ) + np.testing.assert_allclose(float(aux_te), float(aux_ref), atol=1e-5, rtol=1e-5) + + @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) + def test_aux_loss_grads_propagate_to_logits(self, backend_name): + """The aux-loss bwd path must produce non-zero gate-kernel grads + when only the aux-loss scalar is differentiated (no main-output + contribution).""" + backend = PermutationBackend(backend_name) + key = jax.random.PRNGKey(7) + kp, kx = jax.random.split(key) + params = _init_params(kp) + x = _make_inputs(kx) + g_gate = _grad_te_aux_only(params, x, permutation_backend=backend)["gate_kernel"] + assert jnp.all(jnp.isfinite(g_gate)) + assert jnp.any( + g_gate != 0.0 + ), "aux_loss bwd should propagate to gate_kernel via fused_topk bwd" + + +# ----------------------------------------------------------------------------- +# Flax wrapper smoke test +# ----------------------------------------------------------------------------- + + +class TestMoEBlockFlaxWrapper: + """Sanity-check the thin Flax wrapper: forward + grad on init.""" + + def test_init_and_apply(self): + block = MoEBlock( + num_experts=NUM_EXPERTS, + num_experts_per_tok=NUM_EXPERTS_PER_TOK, + intermediate_size=INTERMEDIATE_SIZE, + permutation_backend=PermutationBackend.PURE_JAX, + dtype=DTYPE, + ) + key = jax.random.PRNGKey(8) + ki, kx = jax.random.split(key) + x = _make_inputs(kx) + variables = jax.jit(block.init)(ki, x) + out, aux = jax.jit(block.apply)(variables, x) + assert out.shape == x.shape + assert aux is None + + @jax.jit + def grad_fn(variables, x): + return jax.grad(lambda v, x: jnp.mean(block.apply(v, x)[0] ** 2))(variables, x) + + grads = grad_fn(variables, x) + for name in ("gate_kernel", "wi_0", "wi_1", "wo"): + g = grads["params"][name] + g = g.value if hasattr(g, "value") else g + assert jnp.all(jnp.isfinite(g)), f"{name} grad NaN/Inf" + assert jnp.any(g != 0.0), f"{name} grad zero" diff --git a/tests/jax/test_multiprocess_moe_vjp.py b/tests/jax/test_multiprocess_moe_vjp.py new file mode 100644 index 000000000..97044780f --- /dev/null +++ b/tests/jax/test_multiprocess_moe_vjp.py @@ -0,0 +1,406 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Multi-process (one-GPU-per-process) tests for the unified MoE custom_vjp. + +The launcher ``tests/jax/run_multiprocess_moe_vjp.sh`` forks one pytest +process per visible GPU (mirroring +``examples/jax/encoder/run_test_multiprocessing_encoder.sh``). Each +process binds to exactly one device via +``jax.distributed.initialize(..., local_device_ids=process_id)``; the +participating processes form a global mesh through JAX's distributed +runtime. + +How to run +---------- + +You typically do NOT invoke pytest on this file directly -- use the +launcher, which passes ``--num-process=N --process-id=i`` to each +forked process. Driving it directly with only one process will skip +every test because :func:`jax.distributed.initialize` requires +multiple participants. + + bash tests/jax/run_multiprocess_moe_vjp.sh + +CI invocation lives in ``qa/L0_jax_distributed_unittest/test.sh``. +""" + +import os + +# NCCL needs HBM headroom that JAX's default 90% preallocation does +# not leave. Set before any jax import below. +os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false") +os.environ.setdefault("XLA_PYTHON_CLIENT_MEM_FRACTION", "0.5") + +import sys + +import jax +import jax.numpy as jnp +import numpy as np +import pytest + +from jax.experimental import mesh_utils +from jax.sharding import Mesh, NamedSharding, PartitionSpec as P +from flax.linen import partitioning as nn_partitioning + + +# Per-process distributed bootstrap. Each pytest invocation initializes +# JAX with exactly one local device (its assigned GPU). Once +# initialized, the four processes form one global mesh of 4 devices. +def _init_distributed(num_process: int, process_id: int) -> bool: + """Initialize jax.distributed for this pytest process. + + Returns True if initialization succeeded (i.e. this is a real + multi-process launch), False if num_process == 0 / 1 meaning the + file is being collected without a launcher and tests should be + skipped at module level. + """ + if num_process <= 1: + return False + coord = os.environ.get("MOE_VJP_COORDINATOR_ADDRESS", "127.0.0.1:1234") + jax.distributed.initialize( + coordinator_address=coord, + num_processes=num_process, + process_id=process_id, + local_device_ids=process_id, + ) + assert jax.local_device_count() == 1, "one GPU per process is the whole point" + assert ( + jax.device_count() == num_process + ), f"global device_count {jax.device_count()} != num_process {num_process}" + return True + + +# Read --num-process / --process-id BEFORE pytest collects any tests so +# we can fast-skip the whole module when not in a multiprocess launch. +def _read_mp_options(): + # Use pytest's option lookup via the request fixture isn't available + # at module top-level; parse argv ourselves the same way encoder + # test does. CLI form is e.g. "pytest ... --num-process=4 --process-id=0". + num = int(os.environ.get("MP_NUM_PROCESS", "0") or "0") + pid = int(os.environ.get("MP_PROCESS_ID", "0") or "0") + for i, a in enumerate(sys.argv): + if a.startswith("--num-process="): + num = int(a.split("=", 1)[1]) + elif a == "--num-process" and i + 1 < len(sys.argv): + num = int(sys.argv[i + 1]) + elif a.startswith("--process-id="): + pid = int(a.split("=", 1)[1]) + elif a == "--process-id" and i + 1 < len(sys.argv): + pid = int(sys.argv[i + 1]) + return num, pid + + +_MP_NUM_PROCESS, _MP_PROCESS_ID = _read_mp_options() +_MP_ACTIVE = _init_distributed(_MP_NUM_PROCESS, _MP_PROCESS_ID) + +if not _MP_ACTIVE: + # Skip the entire module if not launched via the multiprocess + # runner. Lets `pytest tests/jax/` collect this file harmlessly. + pytest.skip( + "test_multiprocess_moe_vjp.py requires the multiprocess launcher " + "(run_multiprocess_moe_vjp.sh). Skipping.", + allow_module_level=True, + ) + +from transformer_engine_jax import get_device_compute_capability + +# Grouped GEMM in the MoE custom_vjp currently requires Blackwell +# (sm_100+). Skip the whole file on older arches. +if get_device_compute_capability(0) < 100: + pytest.skip( + "MoE custom_vjp tests require Blackwell (sm_100+) for grouped GEMM", + allow_module_level=True, + ) + +import transformer_engine.jax as te +from transformer_engine.common import recipe as te_recipe +from transformer_engine.jax.flax import _MoEBlock as MoEBlock +from transformer_engine.jax.moe import PermutationBackend +from transformer_engine.jax.sharding import MeshResource, global_shard_guard + +# Parametrize values for the dispatch / combine backend. Only the +# ``triton`` variant carries the ``triton`` marker, so the +# ``pure_jax`` variant still runs on environments without Triton. +BACKEND_PARAMS = [ + pytest.param("pure_jax", id="pure_jax"), + pytest.param("triton", id="triton", marks=pytest.mark.triton), +] + + +EP_AXIS = "ep" +FSDP_AXIS = "fsdp" +EP_SIZE = 2 +# FSDP_SIZE adapts to whatever the launcher gave us: dlcluster GB200 +# gives 4 GPUs (FSDP=2), CI B200 gives 8 GPUs (FSDP=4). Both stay +# 128-aligned for MXFP8 and divide num_experts/topk cleanly. +assert ( + jax.device_count() % EP_SIZE == 0 +), f"device_count {jax.device_count()} must be divisible by EP_SIZE={EP_SIZE}" +FSDP_SIZE = jax.device_count() // EP_SIZE +NUM_DEVICES_REQUIRED = EP_SIZE * FSDP_SIZE + +LOGICAL_AXIS_RULES = ( + ("exp", EP_AXIS), + ("embed", FSDP_AXIS), + ("mlp", None), + ("batch", (EP_AXIS, FSDP_AXIS)), +) + + +@pytest.fixture(scope="module") +def mesh(): + if jax.device_count() < NUM_DEVICES_REQUIRED: + pytest.skip( + f"Need >={NUM_DEVICES_REQUIRED} devices for ep={EP_SIZE} x fsdp={FSDP_SIZE};" + f" have {jax.device_count()}" + ) + devices = mesh_utils.create_device_mesh((EP_SIZE, FSDP_SIZE)) + return Mesh(devices, axis_names=(EP_AXIS, FSDP_AXIS)) + + +# ``recipe`` parametrize values used across all tests below. ``None`` +# = plain bf16; the named recipes route through TE's autocast and +# exercise the FP8/MXFP8 quantization paths in _body_fwd/_body_bwd. +# Only recipes that work on TE Blackwell are included; older GPUs +# skip via the ``hardware_supports`` guard below. +RECIPE_NAMES = ("bf16", "MXFP8BlockScaling") + + +def _resolve_recipe(name): + """Return ``(use_fp8, recipe_instance)`` for the parametrize id.""" + if name == "bf16": + return False, None + if name == "MXFP8BlockScaling": + return True, te_recipe.MXFP8BlockScaling() + raise ValueError(f"unknown recipe name: {name!r}") + + +def _hardware_supports(recipe_name): + """Skip an FP8 recipe on GPUs that don't have the hw for it.""" + if recipe_name == "bf16": + return True + from transformer_engine_jax import get_device_compute_capability + + arch = get_device_compute_capability(0) + if recipe_name == "MXFP8BlockScaling": + return arch >= 100 + return False + + +def _autocast_ctx(recipe_name): + """Context manager that turns FP8 on for non-bf16 recipes.""" + use_fp8, recipe_inst = _resolve_recipe(recipe_name) + return te.autocast(enabled=use_fp8, recipe=recipe_inst) + + +def _tol_finite_grad(recipe_name): + """Per-recipe absolute tolerance for parity grad comparison.""" + if recipe_name == "bf16": + return 5e-2 + # MXFP8 grads carry block-scale quantization noise; loosen accordingly. + return 3e-1 + + +# ----------------------------------------------------------------------------- +# Helpers +# ----------------------------------------------------------------------------- + + +def _make_block( + *, + num_experts, + num_experts_per_tok, + intermediate_size, + permutation_backend, + aux_loss_coeff=0.0, + dtype=jnp.bfloat16, + align_size=0, +): + return MoEBlock( + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + intermediate_size=intermediate_size, + permutation_backend=permutation_backend, + data_parallelism_axes=(FSDP_AXIS,), + aux_loss_coeff=aux_loss_coeff, + dtype=dtype, + _align_size=align_size, + ) + + +def _shard_inputs(x, mesh): + return jax.lax.with_sharding_constraint( + x, NamedSharding(mesh, P((EP_AXIS, FSDP_AXIS), None, None)) + ) + + +def _init_apply(block, mesh, x, key): + with mesh, global_shard_guard( + MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) + ), nn_partitioning.axis_rules(LOGICAL_AXIS_RULES): + x = _shard_inputs(x, mesh) + variables = jax.jit(block.init)(key, x) + jax.block_until_ready(jax.tree_util.tree_leaves(variables)[0]) + output, aux = jax.jit(block.apply)(variables, x) + jax.block_until_ready(output) + return variables, output, aux + + +def _grad_step(block, variables, mesh, x): + with mesh, global_shard_guard( + MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) + ), nn_partitioning.axis_rules(LOGICAL_AXIS_RULES): + x = _shard_inputs(x, mesh) + + def loss_fn(variables, x): + output, aux = block.apply(variables, x) + main = jnp.mean(output.astype(jnp.float32) ** 2) + return main + (aux.astype(jnp.float32) if aux is not None else 0.0) + + grads = jax.jit(jax.grad(loss_fn))(variables, x) + jax.block_until_ready(jax.tree_util.tree_leaves(grads)[0]) + return grads + + +def _unwrap(x): + return x.value if hasattr(x, "value") else x + + +def _local_shard(x): + """Return the local (this-process) shard of a global JAX Array as numpy. + + Every assertion in this file is structural (finite-ness, non-zero, + parity within tolerance). For all of these, checking the local + shard on each process is sufficient and avoids any cross-process + collective in the test machinery. ``arr.addressable_data(0)`` + returns the local-device view of the sharded array -- with one + GPU per process there is exactly one addressable shard. + """ + return np.asarray(jax.device_get(x.addressable_data(0))) + + +# ----------------------------------------------------------------------------- +# Mixtral-style shapes, sized to fit on a single 4-GPU bf16 box (a +# 4-way data-parallel shard of a Mixtral-8 block). +# ----------------------------------------------------------------------------- + +BATCH = EP_SIZE * FSDP_SIZE * 4 # 16 on 4-GPU, 32 on 8-GPU +SEQ = 2048 +HIDDEN = 1024 +INTER = 4096 +NUM_EXPERTS = 8 +TOPK = 2 + + +class TestMoeVjpMultiprocess: + """Multiprocess (one-GPU-per-process) correctness checks for the + unified MoE custom_vjp. + """ + + @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) + @pytest.mark.parametrize("recipe_name", RECIPE_NAMES) + def test_fwd_and_bwd(self, mesh, backend_name, recipe_name): + if not _hardware_supports(recipe_name): + pytest.skip(f"recipe {recipe_name} not supported on this GPU") + backend = PermutationBackend(backend_name) + block = _make_block( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + intermediate_size=INTER, + permutation_backend=backend, + ) + x = jax.random.normal( + jax.random.PRNGKey(0), + (BATCH, SEQ, HIDDEN), + dtype=jnp.bfloat16, + ) + with _autocast_ctx(recipe_name): + variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(1)) + # Local-shard checks (see _local_shard docstring for why). + out_local = _local_shard(output) + assert output.dtype == x.dtype + assert np.all(np.isfinite(out_local)), "output has NaN/Inf" + assert aux is None + with _autocast_ctx(recipe_name): + grads = _grad_step(block, variables, mesh, x) + for name in ("gate_kernel", "wi_0", "wi_1", "wo"): + g_local = _local_shard(_unwrap(grads["params"][name])) + assert np.all(np.isfinite(g_local)), f"{name} grad has NaN/Inf" + assert np.any(g_local != 0.0), f"{name} grad is identically zero" + + @pytest.mark.parametrize("backend_name", BACKEND_PARAMS) + @pytest.mark.parametrize("recipe_name", RECIPE_NAMES) + def test_aux_loss(self, mesh, backend_name, recipe_name): + if not _hardware_supports(recipe_name): + pytest.skip(f"recipe {recipe_name} not supported on this GPU") + backend = PermutationBackend(backend_name) + block = _make_block( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + intermediate_size=INTER, + permutation_backend=backend, + aux_loss_coeff=1e-2, + ) + x = jax.random.normal( + jax.random.PRNGKey(4), + (BATCH, SEQ, HIDDEN), + dtype=jnp.bfloat16, + ) + with _autocast_ctx(recipe_name): + variables, output, aux = _init_apply(block, mesh, x, jax.random.PRNGKey(5)) + out_local = _local_shard(output) + assert np.all(np.isfinite(out_local)), "output has NaN/Inf under aux" + assert aux is not None + assert aux.shape == () + aux_local = _local_shard(aux) + assert np.isfinite(aux_local), "aux is NaN/Inf" + with _autocast_ctx(recipe_name): + grads = _grad_step(block, variables, mesh, x) + g_gate_local = _local_shard(_unwrap(grads["params"]["gate_kernel"])) + assert np.all(np.isfinite(g_gate_local)), "gate grad NaN/Inf under aux" + + @pytest.mark.parametrize("recipe_name", RECIPE_NAMES) + def test_pure_jax_triton_parity(self, mesh, recipe_name): + if not _hardware_supports(recipe_name): + pytest.skip(f"recipe {recipe_name} not supported on this GPU") + block_pj = _make_block( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + intermediate_size=INTER, + permutation_backend=PermutationBackend.PURE_JAX, + ) + block_tr = _make_block( + num_experts=NUM_EXPERTS, + num_experts_per_tok=TOPK, + intermediate_size=INTER, + permutation_backend=PermutationBackend.TRITON, + ) + x = jax.random.normal( + jax.random.PRNGKey(6), + (BATCH, SEQ, HIDDEN), + dtype=jnp.bfloat16, + ) + tol = _tol_finite_grad(recipe_name) + with _autocast_ctx(recipe_name): + variables, out_pj, _ = _init_apply(block_pj, mesh, x, jax.random.PRNGKey(7)) + with mesh, global_shard_guard( + MeshResource(ep_resource=EP_AXIS, fsdp_resource=FSDP_AXIS) + ), nn_partitioning.axis_rules(LOGICAL_AXIS_RULES): + x_sh = _shard_inputs(x, mesh) + out_tr, _ = jax.jit(block_tr.apply)(variables, x_sh) + + out_pj_local = _local_shard(out_pj) + out_tr_local = _local_shard(out_tr) + diff = float(np.max(np.abs(out_pj_local - out_tr_local))) + assert diff < tol, f"forward parity breach: max_abs_diff={diff} (tol={tol})" + + with _autocast_ctx(recipe_name): + grads_pj = _grad_step(block_pj, variables, mesh, x) + grads_tr = _grad_step(block_tr, variables, mesh, x) + for name in ("gate_kernel", "wi_0", "wi_1", "wo"): + g_pj = _local_shard(_unwrap(grads_pj["params"][name])) + g_tr = _local_shard(_unwrap(grads_tr["params"][name])) + d = float(np.max(np.abs(g_pj - g_tr))) + assert d < tol, f"grad parity breach on {name}: max_abs_diff={d} (tol={tol})" diff --git a/transformer_engine/common/triton/permutation.py b/transformer_engine/common/triton/permutation.py index 75bb85f5e..b3893843a 100644 --- a/transformer_engine/common/triton/permutation.py +++ b/transformer_engine/common/triton/permutation.py @@ -12,6 +12,16 @@ from packaging import version +_PERMUTATION_AUTOTUNE_BLOCK_SIZES = (64, 128, 256, 512, 1024, 2048, 4096) + + +def _permutation_autotune_configs(): + """Autotune ``configs`` list shared by every permutation Triton + kernel below. + """ + return [triton.Config({"BLOCK_SIZE": bs}) for bs in _PERMUTATION_AUTOTUNE_BLOCK_SIZES] + + # The following three argsort related kernels are adapted from # the issue https://github.com/triton-lang/triton/issues/3698 @@ -295,15 +305,7 @@ def _permute_kernel( try: _permute_kernel = triton.autotune( - configs=[ - triton.Config({"BLOCK_SIZE": 64}), - triton.Config({"BLOCK_SIZE": 128}), - triton.Config({"BLOCK_SIZE": 256}), - triton.Config({"BLOCK_SIZE": 512}), - triton.Config({"BLOCK_SIZE": 1024}), - triton.Config({"BLOCK_SIZE": 2048}), - triton.Config({"BLOCK_SIZE": 4096}), - ], + configs=_permutation_autotune_configs(), key=["hidden_size"], )(_permute_kernel) except RuntimeError: @@ -416,15 +418,7 @@ def _unpermute_kernel( try: _unpermute_kernel = triton.autotune( - configs=[ - triton.Config({"BLOCK_SIZE": 64}), - triton.Config({"BLOCK_SIZE": 128}), - triton.Config({"BLOCK_SIZE": 256}), - triton.Config({"BLOCK_SIZE": 512}), - triton.Config({"BLOCK_SIZE": 1024}), - triton.Config({"BLOCK_SIZE": 2048}), - triton.Config({"BLOCK_SIZE": 4096}), - ], + configs=_permutation_autotune_configs(), key=["hidden_size"], )(_unpermute_kernel) except RuntimeError: @@ -525,15 +519,7 @@ def _unpermute_bwd_with_merging_probs_kernel( try: _unpermute_bwd_with_merging_probs_kernel = triton.autotune( - configs=[ - triton.Config({"BLOCK_SIZE": 64}), - triton.Config({"BLOCK_SIZE": 128}), - triton.Config({"BLOCK_SIZE": 256}), - triton.Config({"BLOCK_SIZE": 512}), - triton.Config({"BLOCK_SIZE": 1024}), - triton.Config({"BLOCK_SIZE": 2048}), - triton.Config({"BLOCK_SIZE": 4096}), - ], + configs=_permutation_autotune_configs(), key=["hidden_size"], )(_unpermute_bwd_with_merging_probs_kernel) except RuntimeError: @@ -643,15 +629,7 @@ def _sort_chunks_by_map_kernel( try: _sort_chunks_by_map_kernel = triton.autotune( - configs=[ - triton.Config({"BLOCK_SIZE": 64}), - triton.Config({"BLOCK_SIZE": 128}), - triton.Config({"BLOCK_SIZE": 256}), - triton.Config({"BLOCK_SIZE": 512}), - triton.Config({"BLOCK_SIZE": 1024}), - triton.Config({"BLOCK_SIZE": 2048}), - triton.Config({"BLOCK_SIZE": 4096}), - ], + configs=_permutation_autotune_configs(), key=["hidden_size"], )(_sort_chunks_by_map_kernel) except RuntimeError: diff --git a/transformer_engine/jax/flax/__init__.py b/transformer_engine/jax/flax/__init__.py index 92a968f06..adf9c8911 100644 --- a/transformer_engine/jax/flax/__init__.py +++ b/transformer_engine/jax/flax/__init__.py @@ -9,6 +9,7 @@ make_dot_general_cls, make_grouped_dense_cls, ) +from .moe import _MoEBlock from .transformer import extend_logical_axis_rules from .transformer import DotProductAttention, MultiHeadAttention, RelativePositionBiases from .transformer import TransformerLayer, TransformerLayerType @@ -18,6 +19,7 @@ "LayerNorm", "LayerNormDenseGeneral", "LayerNormMLP", + "_MoEBlock", "wrap_function_in_te_state_module", "make_dot_general_cls", "make_grouped_dense_cls", diff --git a/transformer_engine/jax/flax/moe.py b/transformer_engine/jax/flax/moe.py new file mode 100644 index 000000000..91346a7a4 --- /dev/null +++ b/transformer_engine/jax/flax/moe.py @@ -0,0 +1,284 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + +"""Flax Linen MoE block for TransformerEngine JAX. + +This module exposes :class:`_MoEBlock`, an experimental Flax Linen layer +that is a thin wrapper around the framework-agnostic functional MoE entry +point :func:`transformer_engine.jax.moe.moe`. The wrapper's only job is +to: + +1. Register the gate kernel, per-expert FFN kernels, and optional biases + as ``self.param`` slots (with the right + :func:`flax.linen.with_logical_partitioning` annotations so JAX's + sharding layer FSDPs the params correctly). +2. Resolve the EP axis name from the active + :class:`transformer_engine.jax.sharding.MeshResource`. +3. Forward all knobs to :func:`moe`. + +All routing, dispatch, FFN, combine, and aux-loss logic lives in +``moe.py`` under a *single* ``jax.custom_vjp`` so future fusions +(FP8-on-the-wire EP, fused ``ragged_all_to_all + grouped_gemm``, gate + +route + dispatch fusion) can land without touching this wrapper. + +The class is intentionally underscore-prefixed; the public ``MoEBlock`` +alias will be introduced once TE's NCCL-backed EP component (and the +recipe-driven alignment follow-up) stabilises (target: the TE release +following the 2.16 code freeze). +""" + +from typing import Any, Callable, NewType, Optional, Tuple, Union + +import jax.numpy as jnp +from flax import linen as nn + +# Re-exported so downstream users can ``from transformer_engine.jax.flax.moe +# import P`` without a second jax.sharding import. +from jax.sharding import PartitionSpec as P # noqa: F401 # pylint: disable=unused-import + +from ..moe import PermutationBackend, moe +from ..quantize import noop_quantizer_set +from ..router import ScoreFunction +from ..sharding import get_active_resource_axis +from .module import TransformerEngineBase + +PRNGKey = Any +Shape = Tuple[int, ...] +DType = NewType("DType", jnp.dtype) +Array = NewType("Array", jnp.ndarray) +Initializer = Callable[[PRNGKey, Shape, DType], Array] + + +__all__ = ["PermutationBackend", "_MoEBlock"] + + +class _MoEBlock(TransformerEngineBase): + """Experimental Flax MoE layer over TransformerEngine. + + See module docstring for the design (this class is a thin Flax + wrapper around :func:`transformer_engine.jax.moe.moe`). Constructor + knob set kept compatible with the previous bespoke implementation so + existing call sites need no changes. + + Parameters + ---------- + num_experts : int + Total number of experts. Under EP this must be divisible by the + EP mesh axis size. + num_experts_per_tok : int + Top-k value for routing. + intermediate_size : int + Hidden dim of the per-expert FFN (the inner ``mlp`` axis). + activation_type : str + Activation between ``layer_w0 @ wi_0`` and the elementwise + product with ``layer_w0 @ wi_1``. Default ``"silu"``. + + score_function : Union[str, ScoreFunction] + ``"softmax"`` (default) or ``"sigmoid"`` for the routing scores. + use_pre_softmax : bool + Apply softmax before topk (vs. after). + num_groups, group_topk : Optional[int] + Grouped top-k knobs (DeepSeek-style). ``None`` disables grouping. + scaling_factor : float + Multiplier on the routing weights. + use_expert_bias : bool + If ``True``, registers a per-expert routing bias (shape ``[E]``). + Only meaningful with ``score_function="sigmoid"``; the underlying + primitive validates the pairing. + aux_loss_coeff : float + If ``> 0``, return the MoE auxiliary load-balancing loss scalar + in addition to the main output. + + gate_kernel_axes, wi_kernel_axes, wo_kernel_axes, input_axes : + Logical sharding axis tuples (consumed by Flax's + :func:`with_logical_partitioning` and our internal + :func:`with_sharding_constraint_by_logical_axes`). + data_parallelism_axes : tuple[str, ...] + FSDP axes over which the input *batch* dim is sharded IN + ADDITION to the EP axis. Empty (default) means activations are + replicated across non-EP axes within an EP group; set e.g. + ``("fsdp",)`` for true FSDP-of-batch where each device owns a + unique slice of the batch. + permutation_backend : PermutationBackend + ``PURE_JAX`` (default) or ``TRITON``. + _align_size : int + Per-expert group-size alignment (``0`` disables; required > 0 + for quantized grouped GEMM). Internal knob; will be inferred + from the active quantization recipe in a follow-up PR. + + dtype : jnp.dtype + Compute / parameter dtype. + kernel_init, bias_init, expert_bias_init : Initializers. + use_bias : bool + Register per-expert FFN biases. + + Quantization is currently configured via the standard TE autocast + context (``fp8_autocast``/``with_quantizer_set``); per-call + quantizer sets can also be passed through ``__call__``'s + ``quantizer_sets`` keyword once we stabilise the recipe pipeline. + """ + + # Architecture + num_experts: int = 8 + num_experts_per_tok: int = 2 + intermediate_size: int = 2048 + activation_type: str = "silu" + + # Routing + score_function: Union[str, ScoreFunction] = "softmax" + use_pre_softmax: bool = False + num_groups: Optional[int] = None + group_topk: Optional[int] = None + scaling_factor: float = 1.0 + use_expert_bias: bool = False + aux_loss_coeff: float = 0.0 + + # Sharding (logical axes) + gate_kernel_axes: Tuple[Optional[str], ...] = () + wi_kernel_axes: Tuple[Optional[str], ...] = ("exp", "embed", "mlp") + wo_kernel_axes: Tuple[Optional[str], ...] = ("exp", "mlp", "embed") + input_axes: Tuple[Optional[str], ...] = () + + # Parallelism + data_parallelism_axes: Tuple[str, ...] = () + + # Permutation + permutation_backend: PermutationBackend = PermutationBackend.PURE_JAX + _align_size: int = 0 + + # Dtypes / init / misc + dtype: DType = jnp.float32 + kernel_init: Optional[Initializer] = None + bias_init: Initializer = nn.initializers.zeros + expert_bias_init: Initializer = nn.initializers.zeros + use_bias: bool = False + + def __post_init__(self): + if self.kernel_init is None: + object.__setattr__( + self, + "kernel_init", + nn.initializers.variance_scaling( + 1.0, "fan_in", "truncated_normal", dtype=self.dtype + ), + ) + if not isinstance(self.permutation_backend, PermutationBackend): + raise TypeError( + "permutation_backend must be a PermutationBackend, got" + f" {self.permutation_backend!r}" + ) + super().__post_init__() + + @nn.compact + def __call__(self, inputs: Array) -> Tuple[Array, Optional[Array]]: + """Run the MoE forward pass. + + Parameters + ---------- + inputs : jnp.ndarray + ``[batch, sequence, hidden]``. + + Returns + ------- + output : jnp.ndarray + ``[batch, sequence, hidden]``. + aux_loss : Optional[jnp.ndarray] + Scalar load-balancing loss when ``aux_loss_coeff > 0``, + else ``None``. + """ + assert ( + inputs.ndim == 3 + ), f"_MoEBlock expects [batch, sequence, hidden] input, got shape {inputs.shape}" + _, _, hidden_size = inputs.shape + + # Param registrations -- must run OUTSIDE any JAX transform that + # alters the variable scope (e.g. shard_map). The functional + # ``moe(...)`` opens its own shard_map internally for the EP + # path, so registering params here is correct. + gate_kernel = self.param( + "gate_kernel", + nn.with_logical_partitioning(self.kernel_init, self.gate_kernel_axes), + (hidden_size, self.num_experts), + self.dtype, + ) + wi_0 = self.param( + "wi_0", + nn.with_logical_partitioning(self.kernel_init, self.wi_kernel_axes), + (self.num_experts, hidden_size, self.intermediate_size), + self.dtype, + ) + wi_1 = self.param( + "wi_1", + nn.with_logical_partitioning(self.kernel_init, self.wi_kernel_axes), + (self.num_experts, hidden_size, self.intermediate_size), + self.dtype, + ) + wo = self.param( + "wo", + nn.with_logical_partitioning(self.kernel_init, self.wo_kernel_axes), + (self.num_experts, self.intermediate_size, hidden_size), + self.dtype, + ) + wi_0_bias = wi_1_bias = wo_bias = None + if self.use_bias: + wi_0_bias = self.param( + "wi_0_bias", + nn.with_logical_partitioning(self.bias_init, ("exp", "mlp")), + (self.num_experts, self.intermediate_size), + self.dtype, + ) + wi_1_bias = self.param( + "wi_1_bias", + nn.with_logical_partitioning(self.bias_init, ("exp", "mlp")), + (self.num_experts, self.intermediate_size), + self.dtype, + ) + wo_bias = self.param( + "wo_bias", + nn.with_logical_partitioning(self.bias_init, ("exp", "embed")), + (self.num_experts, hidden_size), + self.dtype, + ) + expert_bias = None + if self.use_expert_bias: + expert_bias = self.param( + "expert_bias", + nn.with_logical_partitioning(self.expert_bias_init, ("exp",)), + (self.num_experts,), + self.dtype, + ) + + ep_axis = get_active_resource_axis("ep_resource") + + return moe( + inputs, + gate_kernel, + wi_0, + wi_1, + wo, + wi_0_bias, + wi_1_bias, + wo_bias, + expert_bias, + num_experts=self.num_experts, + num_experts_per_tok=self.num_experts_per_tok, + activation_type=self.activation_type, + score_function=self.score_function, + use_pre_softmax=self.use_pre_softmax, + num_groups=self.num_groups, + group_topk=self.group_topk, + scaling_factor=self.scaling_factor, + aux_loss_coeff=self.aux_loss_coeff, + permutation_backend=self.permutation_backend, + align_size=self._align_size, + gate_inside_vjp=True, + ep_axis=ep_axis, + data_parallelism_axes=self.data_parallelism_axes, + input_axes=self.input_axes, + gate_kernel_axes=self.gate_kernel_axes, + wi_kernel_axes=self.wi_kernel_axes, + wo_kernel_axes=self.wo_kernel_axes, + quantizer_sets=(noop_quantizer_set, noop_quantizer_set, noop_quantizer_set), + dtype=self.dtype, + ) diff --git a/transformer_engine/jax/moe.py b/transformer_engine/jax/moe.py new file mode 100644 index 000000000..2a1c818cb --- /dev/null +++ b/transformer_engine/jax/moe.py @@ -0,0 +1,2165 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. +"""Functional Mixture-of-Experts (MoE) entry point with a single fused VJP. + +This module exposes :func:`moe`, the framework-agnostic flat function that +implements an entire MoE block (gate -> top-k routing -> token dispatch -> +per-expert FFN -> token combine, plus optional expert parallelism via a +shard_map / ragged_all_to_all collective) under a *single* +``jax.custom_vjp``. It is the moral analog of +:func:`transformer_engine.jax.layernorm_mlp.layernorm_mlp` for MoE: one +custom_vjp boundary covers the whole block so future fusions (FP8 over the +EP wire, fused ``ragged_all_to_all + grouped_gemm``, gate+route+dispatch +fusion) can land without re-architecting the call site. + +Design rationale +---------------- + +The earlier MoE block (:class:`transformer_engine.jax.flax.moe._MoEBlock`) +composed many narrower custom_vjps -- one per :func:`grouped_dense`, one +per :func:`token_dispatch`, etc. Every nested custom_vjp is a place where +a quantized :class:`ScaledTensor` cannot survive (JAX requires custom_vjp +inputs / outputs to be plain ``jnp.ndarray`` ish pytrees). To enable +end-to-end FP8 flow -- in particular FP8 carried over the EP +ragged_all_to_all -- the dispatch's quantize, the a2a, the per-expert +FFN, the inverse a2a, and the combine all have to live inside the same +VJP. This file collapses them into one. + +Implementation conventions +-------------------------- + +* No nested ``custom_vjp``. Every primitive's ``_fwd`` and ``_bwd`` is + called directly (e.g. :func:`tex.fused_topk_with_score_function_fwd` / + ``_bwd``, :func:`unpermute_with_mask_map`, + :func:`unpermute_bwd_with_merging_probs`, + :func:`sort_chunks_by_map(is_forward=False)`, + forward + reverse :func:`jax.lax.ragged_all_to_all`) so the outer + ``_moe_bwd_rule`` controls the bwd graph end-to-end without invoking + ``jax.vjp`` for re-linearization. +* The fwd/bwd context (``ctx``) is a plain ``dict`` whose keys depend on + the static configuration (permutation backend, EP active or not, + presence of biases, aux loss enabled). The ``_moe_fwd_rule`` builds a + matching ``ctx_specs`` dict in lockstep when opening the EP shard_map + so ``out_specs`` structurally matches the body's return. +* :func:`_dispatch` is the helper that wraps + ``permute -> a2a -> local_permute`` (forward); :func:`_combine` is its + inverse. Their ``_bwd`` siblings drive the inverse collectives in the + bwd rule. None of these helpers form a custom_vjp boundary. +""" + +import math +from dataclasses import dataclass +from enum import Enum +from functools import partial +from typing import Any, NewType, Optional, Tuple, Union + +import jax +import jax.numpy as jnp +from flax import struct as flax_struct +from jax.sharding import PartitionSpec as P + +from . import cpp_extensions as tex +from .permutation import ( + PureJaxPermState, + compute_ragged_all_to_all_params, + compute_reverse_ragged_all_to_all_params, + pure_jax_token_combine, + pure_jax_token_dispatch, + routing_map_to_selected_experts, +) +from .quantize import ( + QuantizerSet, + ScaledTensor, + TensorUsage, + noop_quantizer_set, + with_sharding_constraint_by_logical_axes, +) +from .flax.module import _convert_to_activation_function +from .router import ScoreFunction, _validate_score_function +from .sharding import _get_mesh + +# Triton-backed primitives are imported lazily: callers on the PURE_JAX +# permutation backend should not need ``triton`` installed. The TRITON +# branches in this module call ``_require_triton()`` first to raise a +# clear error if the import failed. +try: + from .triton_extensions.permutation import ( + make_chunk_sort_map, + make_row_id_map, + permute_with_mask_map, + permute_with_mask_map_and_pad, + sort_chunks_by_map, + unpermute_bwd_with_merging_probs, + unpermute_bwd_with_merging_probs_and_unpad, + unpermute_with_mask_map, + unpermute_with_mask_map_and_unpad, + ) + + _TRITON_AVAILABLE = True +except ImportError: + _TRITON_AVAILABLE = False + make_chunk_sort_map = None + make_row_id_map = None + permute_with_mask_map = None + permute_with_mask_map_and_pad = None + sort_chunks_by_map = None + unpermute_bwd_with_merging_probs = None + unpermute_bwd_with_merging_probs_and_unpad = None + unpermute_with_mask_map = None + unpermute_with_mask_map_and_unpad = None + + +def _require_triton(): + """Raise a clear error if Triton permutation kernels are unavailable.""" + if not _TRITON_AVAILABLE: + raise ImportError( + "PermutationBackend.TRITON requires" + " ``transformer_engine.jax.triton_extensions`` (and ``triton``)." + " Install Triton or pass PermutationBackend.PURE_JAX." + ) + + +PRNGKey = Any +Shape = Tuple[int, ...] +DType = NewType("DType", jnp.dtype) +Array = NewType("Array", jnp.ndarray) + + +__all__ = ["moe", "PermutationBackend"] + + +# ============================================================================= +# Enums +# ============================================================================= + + +class PermutationBackend(Enum): + """Token-dispatch / combine backend used by :func:`moe`. + + * ``TRITON``: TE's fused Triton kernels. Faster than ``PURE_JAX`` + on current hardware and the recommended default. + * ``PURE_JAX``: ``jnp.argsort`` + gather paths compiled as plain + XLA; useful as a numerical reference and on builds without + Triton available. + """ + + PURE_JAX = "pure_jax" + TRITON = "triton" + + +# ============================================================================= +# Dispatch-state records (carried _dispatch -> _combine / *_bwd) +# ============================================================================= +# +# Two NamedTuples (one per permutation backend) so we get type +# discrimination at the consumer side via ``isinstance``. The backend- +# specific residuals are required fields; the EP-only residuals are +# Optional and are populated only when the run is EP-active. Each field +# is either an ``ndarray`` or ``None`` -- nothing static, since these +# values cross the shard_map pytree boundary and would otherwise be +# coerced into JitTracers. + + +@flax_struct.dataclass +class _PureJaxDispatchState: + """Residuals saved by :func:`_dispatch` on the PURE_JAX path. + + Registered as a JAX pytree via ``flax.struct.dataclass``: each + annotated field is a leaf, ``None`` is a non-leaf sentinel. The + matching spec built by :func:`_build_dispatch_specs` mirrors this + layout so shard_map's value and spec trees line up. + """ + + group_sizes: jnp.ndarray + sorted_indices: jnp.ndarray + routing_weights: jnp.ndarray + # EP-only: + all_shards_tokens_per_expert: Optional[jnp.ndarray] = None + local_perm_row_id_map: Optional[jnp.ndarray] = None + + +@flax_struct.dataclass +class _TritonDispatchState: + """Residuals saved by :func:`_dispatch` on the TRITON path.""" + + group_sizes: jnp.ndarray + row_id_map: jnp.ndarray + pad_offsets: Optional[jnp.ndarray] # populated only when align_size > 0 + merging_probs: jnp.ndarray + # EP-only: + all_shards_tokens_per_expert: Optional[jnp.ndarray] = None + local_perm_row_id_map: Optional[jnp.ndarray] = None + + +_DispatchState = Union[_PureJaxDispatchState, _TritonDispatchState] + + +@flax_struct.dataclass +class _BodyCtx: + """Residuals carried fwd_rule -> bwd_rule by :func:`_body_fwd`. + + Optional fields (``expert_bias``, ``aux_*``) are ``None`` when the + matching feature is disabled. :func:`_build_ctx_specs` mirrors that + layout so the shard_map spec and value trees match leaf-for-leaf. + """ + + # Always present. + x: Any + gate_kernel: Any + logits_2d: Any + saved_scores: Any + routing_map: Any + dispatch: Any # _DispatchState + casted_sorted_x_lhs_trans: Any + casted_wi_rhs_trans: Any # combined [E, H, 2M] residual for fused wi_0|wi_1 bwd + gate_proj_out: Any + up_proj_out: Any + casted_intermediate_lhs_trans: Any + casted_wo_rhs_trans: Any + expert_outputs: Any + local_group_sizes: Any + # Feature-gated. + expert_bias: Any = None + aux_const_buf: Any = None + aux_tokens_per_expert: Any = None + aux_logits_for_score: Any = None + aux_saved_scores: Any = None + + +# ============================================================================= +# ctx / dispatch-state key conventions +# ============================================================================= +# +# Both ``ctx`` (carried fwd_rule -> bwd_rule) and the dispatch state +# (carried _dispatch -> _combine / _dispatch_bwd / _combine_bwd) are plain +# python dicts. Using a dict (rather than a flax_struct.dataclass) lets us +# vary the populated keys with the static config without breaking +# ``shard_map``'s ``out_specs`` structural match: the spec dict and the +# value dict are built with the SAME keys via :func:`_build_ctx_specs`. +# +# Below is the key glossary so the rest of the file reads cleanly. +# +# DispatchState (dict): values are jnp.ndarray unless noted +# Always present: +# "group_sizes" [n_groups] per-expert token counts +# (n_groups = E for no-EP, +# E_local for EP) +# "ep_active" bool (carried as a Python flag, +# not in the dict; passed +# alongside) +# PURE_JAX backend: +# "sorted_indices" [num_real + padding] argsort indices +# "routing_weights" [num_tokens, topk] per-token-per-expert weights +# TRITON backend: +# "row_id_map" [num_tokens, 2*E + 1] +# "pad_offsets" [E] or None +# "merging_probs" [num_tokens, E] +# EP-only: +# "all_shards_tokens_per_expert" [num_ep, E] +# "local_perm_row_id_map" [recv_buffer_rows] +# "local_perm_inv_row_id_map" [recv_buffer_rows] +# +# NOTE: per-shard compile-time-constant shapes (num_real_tokens, +# padding_size, pre/post_a2a_buffer_shape) are NOT stored in this +# dict; they are recomputed in _body_fwd/_body_bwd via +# _compute_static_shape_info and passed as Python ints / int tuples to +# the dispatch/combine helpers. Storing them in the dict would cause +# JAX's pytree-flatten across the shard_map boundary to coerce them +# into JitTracer 0-d arrays, which breaks Python-level control flow +# (e.g. ``if padding > 0``) and ``jnp.zeros(shape)`` in the bwd. +# +# See :class:`_BodyCtx` (NamedTuple) for the ctx layout and field +# documentation. :func:`_build_ctx_specs` returns a matching ``_BodyCtx`` +# of ``P(...)`` specs so shard_map's value/spec trees line up +# leaf-for-leaf. + + +# ============================================================================= +# Static shape helper +# ============================================================================= +# +# A set of per-shard shape/size values that the dispatch and combine +# helpers (both fwd and bwd) need. They're all derivable from existing +# static args, so we recompute them in both ``_body_fwd`` and +# ``_body_bwd`` and pass them as Python ints / int-tuples through +# explicit kwargs. We MUST NOT stash them inside the dynamic +# ``state`` / ``ctx`` dict: when the dict crosses the EP shard_map's +# out_specs/in_specs boundary, JAX's pytree-flatten coerces any Python +# int leaves into traced 0-d arrays, which then breaks dependent Python +# code in the bwd (e.g. ``if padding > 0`` and ``jnp.zeros(shape)``). + + +@dataclass(frozen=True) +class _StaticShapeInfo: + """Per-shard compile-time-constant shape info used by dispatch / + combine fwd and bwd. Fields are Python ints / int tuples (NOT jnp + arrays) so they can be passed as ordinary static keyword args. + + Attributes + ---------- + num_real_tokens : int + Per-shard count of real (non-padding) permuted tokens, + i.e. ``per_shard_num_tokens * num_experts_per_tok``. + padding_size : int + Per-shard number of alignment-padding tokens appended to the + sort buffer (``num_experts * (align_size - 1)`` when + ``align_size > 0``, else ``0``). + pre_a2a_buffer_shape : tuple[int, int] + ``(num_real_tokens + padding_size, hidden)`` -- the per-shard + shape of the sorted-inputs buffer sent over the EP + ragged_all_to_all in the fwd direction. + post_a2a_buffer_shape : Optional[tuple[int, int]] + ``(recv_buffer_rows, hidden)`` when EP is active, ``None`` + otherwise. + """ + + num_real_tokens: int + padding_size: int + pre_a2a_buffer_shape: Tuple[int, int] + post_a2a_buffer_shape: Optional[Tuple[int, int]] + + +def _compute_static_shape_info( + *, + batch_size: int, + sequence_length: int, + hidden: int, + num_experts: int, + num_experts_per_tok: int, + align_size: int, + ep_active: bool, + num_ep: int = 1, + fsdp_sizes: Tuple[int, ...] = (), + recv_buffer_rows: int = 0, + batch_is_per_shard: bool = True, +) -> _StaticShapeInfo: + """Build a :class:`_StaticShapeInfo` for the current rank. + + ``batch_is_per_shard`` controls whether ``batch_size`` is already + sharded (True -- e.g. when this is called from inside a shard_map + body, where ``x.shape[0]`` reports the per-shard batch size) or + global (False -- e.g. when computing from x.shape outside the + shard_map body). + """ + if ep_active and not batch_is_per_shard: + dp_size = math.prod(fsdp_sizes) if fsdp_sizes else 1 + per_shard_batch = batch_size // (num_ep * dp_size) + else: + per_shard_batch = batch_size + per_shard_num_tokens = per_shard_batch * sequence_length + num_real_tokens = per_shard_num_tokens * num_experts_per_tok + padding_size = num_experts * (align_size - 1) if align_size > 0 else 0 + pre_a2a_buffer_shape = (num_real_tokens + padding_size, hidden) + post_a2a_buffer_shape = (recv_buffer_rows, hidden) if ep_active else None + return _StaticShapeInfo( + num_real_tokens=num_real_tokens, + padding_size=padding_size, + pre_a2a_buffer_shape=pre_a2a_buffer_shape, + post_a2a_buffer_shape=post_a2a_buffer_shape, + ) + + +# ============================================================================= +# Dispatch / combine helpers (no VJP boundary -- pure Python) +# ============================================================================= + + +def _dispatch( + inputs_2d: jnp.ndarray, + sparse_probs: jnp.ndarray, + routing_map: jnp.ndarray, + *, + backend: PermutationBackend, + num_experts: int, + num_experts_per_tok: int, + align_size: int, + # EP-only: + ep_active: bool, + ep_axis: Optional[str], + num_ep: int, + recv_buffer_rows: int, + shard_id: Optional[jnp.ndarray] = None, +) -> Tuple[jnp.ndarray, dict]: + """``permute -> (a2a -> local_permute) iff ep_active``. + + Returns ``(sorted_x, state)`` where ``sorted_x`` has shape + ``[buffer_rows, hidden]`` -- ``E`` groups (no-EP) or ``E_local`` groups + (EP) -- and ``state`` is a dict carrying everything :func:`_combine` + and the bwd helpers need to reverse the operation. + + Bypasses the ``custom_vjp``-wrapped public ``token_dispatch`` / + ``pure_jax_token_dispatch`` wrappers (well, mostly: PURE_JAX still + composes through ``pure_jax_token_dispatch`` because that helper has + no ``custom_vjp`` itself -- only its inner ``_sort_activations`` does, + which is fine since we never auto-diff through it from this layer). + For TRITON we call the underlying ``permute_with_mask_map`` / + ``permute_with_mask_map_and_pad`` primitives directly. + """ + num_tokens, hidden = inputs_2d.shape + topk = num_experts_per_tok + + # Backend-specific residuals collected here, then packaged into the + # appropriate _*DispatchState below. + sorted_indices = None + routing_weights_kept = None + row_id_map = None + pad_offsets = None + merging_probs = None + + # ------------------------------------------------------------------ + # Step 1: global permute (every shard routes its own tokens over the + # full expert axis). Backend-specific. + # ------------------------------------------------------------------ + if backend is PermutationBackend.PURE_JAX: + selected_experts, routing_weights = routing_map_to_selected_experts( + sparse_probs, routing_map, topk + ) + sorted_inputs, perm_state, group_sizes = pure_jax_token_dispatch( + inputs_2d, + selected_experts, + num_experts=num_experts, + num_experts_per_tok=topk, + align_size=align_size, + ) + # NOTE: ``perm_state.num_real_tokens`` and ``perm_state.padding_size`` + # are compile-time Python ints; intentionally NOT stored in the + # returned state (would be coerced to JitTracer 0-d arrays under + # the EP shard_map's pytree flatten). Recompute via + # ``_compute_static_shape_info`` in the bwd / EP-combine + # call sites that need them. + sorted_indices = perm_state.sorted_indices + routing_weights_kept = routing_weights + else: + # TRITON backend -- inline the underlying primitive sequence + # (mirrors ``_token_dispatch_fwd_rule`` but exposes the residuals + # to our ctx instead of saving them inside another custom_vjp). + num_out_tokens = num_tokens * topk + row_id_map = make_row_id_map(routing_map, num_tokens, num_experts) + tokens_per_expert = jnp.sum(routing_map, axis=0).astype(jnp.int32) + if align_size > 0: + target_tokens_per_expert = ( + jnp.ceil(tokens_per_expert / align_size) * align_size + ).astype(jnp.int32) + pad_lengths = target_tokens_per_expert - tokens_per_expert + cum_pad = jnp.cumsum(pad_lengths) + pad_offsets = jnp.concatenate([jnp.array([0], dtype=cum_pad.dtype), cum_pad[:-1]]) + worst_case_out_tokens = ( + (num_out_tokens + num_experts * (align_size - 1)) // align_size + ) * align_size + sorted_inputs, _ = permute_with_mask_map_and_pad( + inputs_2d, + row_id_map, + None, + pad_offsets, + num_tokens, + num_experts, + worst_case_out_tokens, + hidden, + align_size=align_size, + ) + group_sizes = target_tokens_per_expert + else: + sorted_inputs, _ = permute_with_mask_map( + inputs_2d, + row_id_map, + None, + num_tokens, + num_experts, + num_out_tokens, + hidden, + ) + pad_offsets = None + group_sizes = tokens_per_expert + merging_probs = sparse_probs + + def _build_state(group_sizes_val, ep_all=None, ep_local=None): + if backend is PermutationBackend.PURE_JAX: + return _PureJaxDispatchState( + group_sizes=group_sizes_val, + sorted_indices=sorted_indices, + routing_weights=routing_weights_kept, + all_shards_tokens_per_expert=ep_all, + local_perm_row_id_map=ep_local, + ) + return _TritonDispatchState( + group_sizes=group_sizes_val, + row_id_map=row_id_map, + pad_offsets=pad_offsets, + merging_probs=merging_probs, + all_shards_tokens_per_expert=ep_all, + local_perm_row_id_map=ep_local, + ) + + if not ep_active: + return sorted_inputs, _build_state(group_sizes) + + # ------------------------------------------------------------------ + # Step 2 (EP only): all_gather per-expert counts so every shard knows + # the [num_ep, num_experts] token-count matrix. + # ------------------------------------------------------------------ + all_shards_tokens_per_expert = jax.lax.all_gather( + group_sizes[None, :], + axis_name=ep_axis, + axis=0, + tiled=True, + ) + + # ------------------------------------------------------------------ + # Step 3 (EP only): forward ragged_all_to_all over the EP axis. + # ------------------------------------------------------------------ + in_off, send_sz, out_off, recv_sz = compute_ragged_all_to_all_params( + all_shards_tokens_per_expert, shard_id, num_ep + ) + post_a2a_buffer_shape = (recv_buffer_rows, hidden) + recv_buf = jnp.zeros(post_a2a_buffer_shape, dtype=sorted_inputs.dtype) + x_recv = jax.lax.ragged_all_to_all( + sorted_inputs, recv_buf, in_off, send_sz, out_off, recv_sz, axis_name=ep_axis + ) + + # ------------------------------------------------------------------ + # Step 4 (EP only): local permute -- (source_shard, expert) -> + # (expert, shard). Inlined ``local_permute_after_a2a`` so we control + # both the row_id_map and its inverse for the bwd. + # ------------------------------------------------------------------ + num_experts_local = num_experts // num_ep + local_expert_start = shard_id * num_experts_local + local_expert_columns = jax.lax.dynamic_slice( + all_shards_tokens_per_expert, + start_indices=(0, local_expert_start), + slice_sizes=(num_ep, num_experts_local), + ) + split_sizes = local_expert_columns.reshape(-1) # source-major + indices_matrix = jnp.arange(num_ep * num_experts_local, dtype=jnp.int32).reshape( + num_ep, num_experts_local + ) + sorted_chunk_indices = indices_matrix.T.reshape(-1) # source-major -> expert-major + num_chunks = num_ep * num_experts_local + # Build a SINGLE row_id_map. ``is_forward=True`` permutes + # source-major -> expert-major; ``is_forward=False`` is the exact + # inverse (this is exactly what ``_sort_chunks_by_index_bwd_rule`` + # uses on the saved residual). _MoEBlock builds two row_id_maps + # only because it calls ``sort_chunks_by_index`` twice -- once in + # ``local_permute_after_a2a`` and again in ``local_unpermute_before_a2a``; + # each of those wrappers calls ``make_chunk_sort_map`` internally. + # Here we share one map across (fwd permute, fwd inverse-permute, + # bwd permute, bwd inverse-permute). + local_perm_row_id_map = make_chunk_sort_map( + split_sizes, sorted_chunk_indices, recv_buffer_rows, num_chunks + ) + sorted_x, _ = sort_chunks_by_map( + x_recv, local_perm_row_id_map, None, recv_buffer_rows, hidden, is_forward=True + ) + local_group_sizes = jnp.sum(local_expert_columns, axis=0) + + # NOTE: pre_a2a_buffer_shape and post_a2a_buffer_shape are compile- + # time int tuples; intentionally NOT stored in the returned state + # (would be coerced to JitTracer 0-d arrays under the EP shard_map's + # pytree flatten). Recompute via ``_compute_static_shape_info`` in + # the bwd call sites that need them. For EP, ``group_sizes`` here is + # the per-local-expert count (the FFN runs over E_local groups, not + # E). The global ``group_sizes`` lives inside + # ``all_shards_tokens_per_expert`` if anyone needs it for + # diagnostics. + return sorted_x, _build_state( + local_group_sizes, + ep_all=all_shards_tokens_per_expert, + ep_local=local_perm_row_id_map, + ) + + +def _combine( + expert_outputs: jnp.ndarray, + state: _DispatchState, + *, + backend: PermutationBackend, + ep_active: bool, + batch_size: int, + sequence_length: int, + dtype: jnp.dtype, + num_experts_per_tok: int, + # Per-shard compile-time-constant shape info (Python ints / int tuples). + # Computed by _compute_static_shape_info in the caller, passed here + # rather than stored in ``state`` to survive shard_map crossings. + num_real_tokens: int, + padding_size: int, + pre_a2a_buffer_shape: Tuple[int, int], + # EP-only: + ep_axis: Optional[str], + shard_id: Optional[jnp.ndarray] = None, + num_ep: int = 1, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Inverse of :func:`_dispatch`. + + Returns ``(output, expert_outputs_post_ep)``. ``output`` is the + ``[B, S, H]`` combined activations. ``expert_outputs_post_ep`` is + the FFN-output tensor in the shape that Step 3 of the combine + actually consumed (i.e. after the reverse ragged_all_to_all on EP + runs, or the original input on non-EP). The caller stashes this as + the bwd residual so that ``_combine_bwd``'s Step-3 inverse sees + the same tensor the forward Step 3 used. + """ + if ep_active: + # Step 1 (EP): inverse local permute. Reuse the SAME row_id_map + # built in _dispatch by setting is_forward=False (this is the + # exact inverse, identical to what + # ``_sort_chunks_by_index_bwd_rule`` does with the saved residual). + recv_buffer_rows, hidden = expert_outputs.shape + x_send_back, _ = sort_chunks_by_map( + expert_outputs, + state.local_perm_row_id_map, + None, + recv_buffer_rows, + hidden, + is_forward=False, + ) + # Step 2 (EP): reverse ragged_all_to_all. + in_off_r, send_sz_r, out_off_r, recv_sz_r = compute_reverse_ragged_all_to_all_params( + state.all_shards_tokens_per_expert, shard_id, num_ep + ) + send_back_buf = jnp.zeros(pre_a2a_buffer_shape, dtype=expert_outputs.dtype) + expert_outputs = jax.lax.ragged_all_to_all( + x_send_back, + send_back_buf, + in_off_r, + send_sz_r, + out_off_r, + recv_sz_r, + axis_name=ep_axis, + ) + + # Step 3: global combine. ``expert_outputs`` here is the post-A2A + # tensor under EP, or the original input under non-EP -- whichever + # value Step 3 actually consumes. Returned as the second tuple + # element so the caller can stash it as the bwd residual. + if backend is PermutationBackend.PURE_JAX: + # Reuse the reference pure-jax implementation; it has no + # custom_vjp on its outer surface so we can call it freely. + perm_state = PureJaxPermState( + sorted_indices=state.sorted_indices, + num_real_tokens=num_real_tokens, + padding_size=padding_size, + ) + output = pure_jax_token_combine( + expert_outputs, + perm_state, + state.routing_weights, + num_experts_per_tok=num_experts_per_tok, + batch_size=batch_size, + sequence_length=sequence_length, + ) + return output, expert_outputs + # TRITON + num_tokens = state.row_id_map.shape[0] + num_experts = (state.row_id_map.shape[1] - 1) // 2 + hidden = expert_outputs.shape[-1] + if state.pad_offsets is not None: + out_2d, _ = unpermute_with_mask_map_and_unpad( + expert_outputs, + state.row_id_map, + state.merging_probs, + None, + state.pad_offsets, + num_tokens, + num_experts, + hidden, + ) + else: + out_2d, _ = unpermute_with_mask_map( + expert_outputs, + state.row_id_map, + state.merging_probs, + None, + num_tokens, + num_experts, + hidden, + ) + return out_2d.reshape(batch_size, sequence_length, hidden).astype(dtype), expert_outputs + + +def _combine_bwd( # pylint: disable=unused-argument + d_output: jnp.ndarray, + state: _DispatchState, + expert_outputs: jnp.ndarray, + *, + backend: PermutationBackend, + ep_active: bool, + batch_size: int, + sequence_length: int, + dtype: jnp.dtype, + num_experts: int, + num_experts_per_tok: int, + # Per-shard compile-time-constant shape info (Python ints / int tuples). + # See ``_compute_static_shape_info`` and the note in ``_dispatch`` + # for why these are kwargs rather than state-dict entries. + num_real_tokens: int, + padding_size: int, + post_a2a_buffer_shape: Optional[Tuple[int, int]], + # EP-only: + ep_axis: Optional[str], + shard_id: Optional[jnp.ndarray] = None, + num_ep: int = 1, +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: + """Inverse of :func:`_combine` on the cotangent. + + Returns ``(d_expert_outputs, d_routing_weights_or_merging_probs)``. + + ``expert_outputs`` is the *forward* output of the FFN (same value the + fwd handed to :func:`_combine`). It's required by the TRITON + combine_bwd kernel; for PURE_JAX we don't need it but accept it for + a symmetric signature. + """ + # Step 3 inverse: global combine bwd. + d_output_2d = d_output.reshape(-1, d_output.shape[-1]) + if backend is PermutationBackend.PURE_JAX: + # The pure-jax combine is: + # unsort = _sort_activations(expert_outputs, argsort(sorted_indices)) + # if pad: unsort = unsort[:num_real] + # reshape -> einsum BKE,BK -> BE -> reshape to BSE + # Hand-derive the bwd in plain JAX (no custom_vjp involved): + unsort_indices = jnp.argsort(state.sorted_indices) + topk = num_experts_per_tok + num_real = num_real_tokens + padding = padding_size + # Recover the unsorted intermediate that the fwd produced (we + # need it for the d_routing_weights pullback). Apply the same + # gather the fwd did. + unsort_intermediate = expert_outputs[unsort_indices] + if padding > 0: + unsort_intermediate = unsort_intermediate[:num_real] + # Bwd of einsum/reshape: + # output[B, E] = sum_K intermediate[B, K, E] * weights[B, K] + # d_intermediate[B, K, E] = d_output[B, E] * weights[B, K] + # d_weights[B, K] = sum_E d_output[B, E] * intermediate[B, K, E] + rw = state.routing_weights.reshape(-1, topk) + intermediate_3d = unsort_intermediate.reshape(rw.shape[0], topk, -1) + rw_cast = rw.astype(intermediate_3d.dtype) + d_intermediate_3d = jnp.einsum("BE,BK -> BKE", d_output_2d, rw_cast) + d_routing_weights = jnp.einsum("BE,BKE -> BK", d_output_2d, intermediate_3d).astype( + state.routing_weights.dtype + ) + d_routing_weights = d_routing_weights.reshape(state.routing_weights.shape) + d_unsort_intermediate = d_intermediate_3d.reshape(num_real, -1) + # Pad back with zeros if the fwd stripped padding. + if padding > 0: + d_unsort_intermediate = jnp.concatenate( + [ + d_unsort_intermediate, + jnp.zeros( + (padding, d_unsort_intermediate.shape[-1]), + dtype=d_unsort_intermediate.dtype, + ), + ], + axis=0, + ) + # Bwd of the gather is gather-by-original-indices: + # sorted = unsort[argsort(sorted_indices)] + # d_sorted = scatter d_unsort via argsort(sorted_indices) + # = d_unsort[sorted_indices] (gather by original sorted_indices, + # which is the inverse of argsort(sorted_indices)). + d_expert_outputs_global = d_unsort_intermediate[state.sorted_indices] + else: + # TRITON combine bwd: requires fwd_input (expert_outputs). + num_tokens = state.row_id_map.shape[0] + n_experts = (state.row_id_map.shape[1] - 1) // 2 + hidden = d_output_2d.shape[-1] + num_out_tokens = expert_outputs.shape[0] + if state.pad_offsets is not None: + d_expert_outputs_global, d_merging_probs = unpermute_bwd_with_merging_probs_and_unpad( + d_output_2d, + state.row_id_map, + expert_outputs, + state.merging_probs, + state.pad_offsets, + num_tokens, + n_experts, + num_out_tokens, + hidden, + ) + # The kernel only writes positions tokens map to; padded + # positions may contain NaN. Replace with zeros (matches + # ``_token_combine_bwd_rule``). + d_expert_outputs_global = jnp.where( + jnp.isnan(d_expert_outputs_global), 0.0, d_expert_outputs_global + ) + else: + d_expert_outputs_global, d_merging_probs = unpermute_bwd_with_merging_probs( + d_output_2d, + state.row_id_map, + expert_outputs, + state.merging_probs, + num_tokens, + n_experts, + num_out_tokens, + hidden, + ) + d_routing_weights = d_merging_probs + + if not ep_active: + return d_expert_outputs_global, d_routing_weights + + # Step 2 (EP) inverse: bwd of reverse ragged_all_to_all is a forward + # ragged_all_to_all using the SAME forward parameters (sender / + # receiver roles swap from the reverse direction back to forward). + in_off_f, send_sz_f, out_off_f, recv_sz_f = compute_ragged_all_to_all_params( + state.all_shards_tokens_per_expert, shard_id, num_ep + ) + recv_buf_for_bwd = jnp.zeros(post_a2a_buffer_shape, dtype=d_expert_outputs_global.dtype) + d_x_send_back = jax.lax.ragged_all_to_all( + d_expert_outputs_global, + recv_buf_for_bwd, + in_off_f, + send_sz_f, + out_off_f, + recv_sz_f, + axis_name=ep_axis, + ) + # Step 1 (EP) inverse: combine fwd applied is_forward=False; the + # bwd is is_forward=True with the SAME row_id_map. + recv_buffer_rows, hidden = d_x_send_back.shape + d_expert_outputs, _ = sort_chunks_by_map( + d_x_send_back, + state.local_perm_row_id_map, + None, + recv_buffer_rows, + hidden, + is_forward=True, + ) + return d_expert_outputs, d_routing_weights + + +def _dispatch_bwd( + d_sorted_x: jnp.ndarray, + state: _DispatchState, + inputs_2d_shape: Tuple[int, ...], + *, + backend: PermutationBackend, + ep_active: bool, + num_experts: int, + num_experts_per_tok: int, + # Per-shard compile-time-constant shape info (Python ints / int tuples). + # See ``_compute_static_shape_info`` and the note in ``_dispatch`` + # for why these are kwargs rather than state-dict entries. + num_real_tokens: int, + padding_size: int, + pre_a2a_buffer_shape: Tuple[int, int], + # EP-only: + ep_axis: Optional[str], + shard_id: Optional[jnp.ndarray] = None, + num_ep: int = 1, +) -> jnp.ndarray: + """Inverse of :func:`_dispatch` on the cotangent. Returns ``d_inputs_2d``. + + The probs path through dispatch is always discarded (PURE_JAX never + threads probs through dispatch; TRITON technically does but the + caller drops ``permuted_probs``, so its cotangent is structurally + zero). The probs gradient instead flows back through + :func:`_combine_bwd`. + """ + if ep_active: + # Step 4 inverse: dispatch fwd applied is_forward=True; bwd is + # is_forward=False with the SAME row_id_map. + recv_buffer_rows, hidden = d_sorted_x.shape + d_x_recv, _ = sort_chunks_by_map( + d_sorted_x, + state.local_perm_row_id_map, + None, + recv_buffer_rows, + hidden, + is_forward=False, + ) + # Step 3 inverse: bwd of forward ragged_a2a is the reverse-direction + # ragged_a2a using the SAME params with sender/receiver swapped. + in_off_r, send_sz_r, out_off_r, recv_sz_r = compute_reverse_ragged_all_to_all_params( + state.all_shards_tokens_per_expert, shard_id, num_ep + ) + recv_buf_pre = jnp.zeros(pre_a2a_buffer_shape, dtype=d_x_recv.dtype) + d_sorted_x = jax.lax.ragged_all_to_all( + d_x_recv, + recv_buf_pre, + in_off_r, + send_sz_r, + out_off_r, + recv_sz_r, + axis_name=ep_axis, + ) + + # Step 1 inverse: global permute bwd. + if backend is PermutationBackend.PURE_JAX: + # Fwd was: replicated = repeat(inputs_2d, topk, axis=0) + # padded = pad(replicated, (0, padding_size)) + # sorted = padded[sorted_indices] + # Bwd: d_padded = scatter via sorted_indices + # = d_sorted[argsort(sorted_indices)] + # d_replicated = d_padded[:num_real] + # d_inputs_2d = d_replicated.reshape(T, topk, H).sum(axis=1) + sorted_indices = state.sorted_indices + num_real = num_real_tokens + padding = padding_size + topk = num_experts_per_tok + unsort_indices = jnp.argsort(sorted_indices) + d_padded = d_sorted_x[unsort_indices] + if padding > 0: + d_replicated = d_padded[:num_real] + else: + d_replicated = d_padded + num_tokens = inputs_2d_shape[0] + hidden = inputs_2d_shape[-1] + d_inputs_2d = d_replicated.reshape(num_tokens, topk, hidden).sum(axis=1) + return d_inputs_2d + + # TRITON: bwd is unpermute_with_mask_map[_and_unpad]. + num_tokens = inputs_2d_shape[0] + hidden = inputs_2d_shape[-1] + if state.pad_offsets is not None: + d_inputs_2d, _ = unpermute_with_mask_map_and_unpad( + d_sorted_x, + state.row_id_map, + None, + None, + state.pad_offsets, + num_tokens, + num_experts, + hidden, + ) + else: + d_inputs_2d, _ = unpermute_with_mask_map( + d_sorted_x, + state.row_id_map, + None, + None, + num_tokens, + num_experts, + hidden, + ) + return d_inputs_2d + + +# ============================================================================= +# Per-shard body +# ============================================================================= + + +def _body_fwd( # pylint: disable=unused-argument + captured: dict, + *, + # Statics + num_experts: int, + num_experts_per_tok: int, + activation_type: str, + score_function: ScoreFunction, + use_pre_softmax: bool, + num_groups: Optional[int], + group_topk: Optional[int], + scaling_factor: float, + aux_loss_coeff: float, + permutation_backend: PermutationBackend, + align_size: int, + gate_inside_vjp: bool, + quantizer_sets: Tuple[QuantizerSet, QuantizerSet, QuantizerSet], + dtype: jnp.dtype, + # EP-only statics + ep_active: bool, + ep_axis: Optional[str], + data_parallelism_axes: Tuple[str, ...], + fsdp_sizes: Tuple[int, ...], + num_ep: int, + num_experts_local: int, + recv_buffer_rows: int, +) -> Tuple[jnp.ndarray, jnp.ndarray, dict]: + """Per-shard forward body. Returns ``(output, aux_loss, ctx_dict)``. + + ``aux_loss`` is always materialized (zeros scalar when disabled) so + the ``shard_map``'s ``out_specs`` has a static structure. + """ + if not gate_inside_vjp: + raise NotImplementedError( + "gate_inside_vjp=False is deferred to a follow-up PR; for now" + " the gate GEMM lives inside the MoE VJP." + ) + + x = captured["inputs"] + gate_kernel = captured["gate_kernel"] + wi_0 = captured["wi_0"] + wi_1 = captured["wi_1"] + wo = captured["wo"] + wi_0_bias = captured.get("wi_0_bias") + wi_1_bias = captured.get("wi_1_bias") + wo_bias = captured.get("wo_bias") + expert_bias = captured.get("expert_bias") + + batch_size, sequence_length, hidden = x.shape + + # ---------------- Stage 1: gate ---------------- + gate_kernel_cast = gate_kernel.astype(x.dtype) + gate_logits = jnp.einsum("bsh,he->bse", x, gate_kernel_cast) + logits_2d = gate_logits.reshape(-1, num_experts) + inputs_2d = x.reshape(-1, hidden) + + # ---------------- Stage 2: routing ---------------- + # Under EP, expert_bias is sharded P(ep_axis); the router needs the + # full E-dim view, so all_gather it. + if ep_active and expert_bias is not None: + full_expert_bias = jax.lax.all_gather(expert_bias, axis_name=ep_axis, tiled=True) + else: + full_expert_bias = expert_bias + # Pass an empty array sentinel when expert_bias is unused (the + # underlying primitive expects a real ndarray, not None). + eb_arg = ( + full_expert_bias if full_expert_bias is not None else jnp.zeros((0,), dtype=jnp.float32) + ) + sparse_probs, routing_map, saved_scores = tex.fused_topk_with_score_function_fwd( + logits_2d, + topk=num_experts_per_tok, + use_pre_softmax=use_pre_softmax, + num_groups=-1 if num_groups is None else num_groups, + group_topk=-1 if group_topk is None else group_topk, + scaling_factor=scaling_factor, + score_function=score_function, + expert_bias=eb_arg, + compute_aux_scores=False, + ) + sparse_probs = sparse_probs.astype(dtype) + + # ---------------- Stage 2b: aux loss ---------------- + if aux_loss_coeff > 0.0: + if ep_active: + collective_axes: Any = ( + ep_axis if not data_parallelism_axes else (ep_axis, *data_parallelism_axes) + ) + global_logits_2d = jax.lax.all_gather( + logits_2d, axis_name=collective_axes, axis=0, tiled=True + ) + _, global_routing_map, _ = tex.fused_topk_with_score_function_fwd( + global_logits_2d, + topk=num_experts_per_tok, + use_pre_softmax=use_pre_softmax, + num_groups=-1 if num_groups is None else num_groups, + group_topk=-1 if group_topk is None else group_topk, + scaling_factor=scaling_factor, + score_function=score_function, + expert_bias=eb_arg, + compute_aux_scores=False, + ) + aux_tokens_per_expert = jnp.sum(global_routing_map.astype(jnp.int32), axis=0) + aux_logits_for_score = global_logits_2d + else: + aux_tokens_per_expert = jnp.sum(routing_map.astype(jnp.int32), axis=0) + aux_logits_for_score = logits_2d + # Aux-side scores: clean per-expert scores (no grouped routing, + # no bias). compute_aux_scores=True takes a separate path that + # ignores the grouping knobs. + aux_probs, _aux_routing_map, aux_saved_scores = tex.fused_topk_with_score_function_fwd( + aux_logits_for_score.astype(jnp.float32), + topk=num_experts_per_tok, + use_pre_softmax=False, + num_groups=-1, + group_topk=-1, + scaling_factor=1.0, + score_function=score_function, + expert_bias=jnp.zeros((0,), dtype=jnp.float32), + compute_aux_scores=True, + ) + aux_loss, aux_const_buf = tex.fused_moe_aux_loss_fwd( + aux_probs.astype(jnp.float32), + aux_tokens_per_expert.astype(jnp.int32), + topk=num_experts_per_tok, + coeff=aux_loss_coeff, + ) + else: + aux_loss = jnp.zeros((), dtype=dtype) + aux_const_buf = None + aux_tokens_per_expert = None + aux_logits_for_score = None + aux_saved_scores = None + + # ---------------- Stage 3: dispatch ---------------- + shard_id = jax.lax.axis_index(ep_axis) if ep_active else None + sorted_x, dispatch_state = _dispatch( + inputs_2d, + sparse_probs, + routing_map, + backend=permutation_backend, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + align_size=align_size, + ep_active=ep_active, + ep_axis=ep_axis, + num_ep=num_ep, + recv_buffer_rows=recv_buffer_rows, + shard_id=shard_id, + ) + local_group_sizes = dispatch_state.group_sizes + + # ---------------- Stage 4: per-expert FFN (inlined) ---------------- + q_set_w0, q_set_w1, q_set_wo = quantizer_sets + if q_set_w0 == noop_quantizer_set: + wi_0 = wi_0.astype(sorted_x.dtype) + if q_set_w1 == noop_quantizer_set: + wi_1 = wi_1.astype(sorted_x.dtype) + if q_set_wo == noop_quantizer_set: + wo = wo.astype(sorted_x.dtype) + + # GEMM 1+2 (fused): up_proj_combined = sorted_x @ wi where + # wi := concat([wi_0, wi_1], axis=-1) -> shape [E, H, 2M] + # combined_out := sorted_x @ wi -> shape [T, 2M] + # Splitting the output back into ``gate_proj_out`` / ``up_proj_out`` + # is free (it's a slicing reshape). This collapses two grouped + # GEMMs and two grouped quantizes of ``sorted_x`` (one per kernel) + # into one of each. Bias is concatenated the same way. + # + # FP8/MXFP8 caveat: per-expert amax is now computed over [H, 2M] + # rather than [H, M] for each of wi_0 / wi_1 separately, so the + # representable range for one of the two halves may shift slightly + # vs. the pre-fusion code. Numerics tests cover this. + inter_M = wi_0.shape[-1] + wi_combined = jnp.concatenate([wi_0, wi_1], axis=-1) + wi_combined_bias = ( + jnp.concatenate([wi_0_bias, wi_1_bias], axis=-1) if wi_0_bias is not None else None + ) + casted_sorted_x = tex.grouped_quantize(sorted_x, q_set_w0.x, local_group_sizes, flatten_axis=-1) + casted_wi = tex.grouped_quantize(wi_combined, q_set_w0.kernel, flatten_axis=-1) + combined_out = tex.grouped_gemm( + casted_sorted_x.get_tensor(usage=TensorUsage.LHS), + casted_wi.get_tensor(usage=TensorUsage.RHS), + contracting_dims=((1,), (1,)), + bias=wi_combined_bias, + ) + gate_proj_out = combined_out[..., :inter_M] + up_proj_out = combined_out[..., inter_M:] + casted_sorted_x_lhs_trans = casted_sorted_x.get_tensor(usage=TensorUsage.LHS_TRANS) + casted_wi_rhs_trans = casted_wi.get_tensor(usage=TensorUsage.RHS_TRANS) + if isinstance(casted_sorted_x_lhs_trans, ScaledTensor): + casted_sorted_x_lhs_trans = casted_sorted_x_lhs_trans.checkpoint(q_set_w0.x) + if isinstance(casted_wi_rhs_trans, ScaledTensor): + casted_wi_rhs_trans = casted_wi_rhs_trans.checkpoint(q_set_w0.kernel) + + # Activation: intermediate = act(gate_proj_out) * up_proj_out + act_fn = _convert_to_activation_function(activation_type) + intermediate = act_fn(gate_proj_out) * up_proj_out + + # GEMM 3: expert_outputs = intermediate @ wo + casted_intermediate = tex.grouped_quantize( + intermediate, q_set_wo.x, local_group_sizes, flatten_axis=-1 + ) + casted_wo = tex.grouped_quantize(wo, q_set_wo.kernel, flatten_axis=-1) + expert_outputs = tex.grouped_gemm( + casted_intermediate.get_tensor(usage=TensorUsage.LHS), + casted_wo.get_tensor(usage=TensorUsage.RHS), + contracting_dims=((1,), (1,)), + bias=wo_bias, + ) + casted_intermediate_lhs_trans = casted_intermediate.get_tensor(usage=TensorUsage.LHS_TRANS) + casted_wo_rhs_trans = casted_wo.get_tensor(usage=TensorUsage.RHS_TRANS) + if isinstance(casted_intermediate_lhs_trans, ScaledTensor): + casted_intermediate_lhs_trans = casted_intermediate_lhs_trans.checkpoint(q_set_wo.x) + if isinstance(casted_wo_rhs_trans, ScaledTensor): + casted_wo_rhs_trans = casted_wo_rhs_trans.checkpoint(q_set_wo.kernel) + + # ---------------- Stage 5: combine ---------------- + # Compute per-shard static shape info once and pass through both + # _combine and (later) the bwd helpers via kwargs -- never via the + # state dict, which gets pytree-flattened across shard_map and would + # coerce Python ints into JitTracer 0-d arrays. + _static_shape = _compute_static_shape_info( + batch_size=batch_size, + sequence_length=sequence_length, + hidden=hidden, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + align_size=align_size, + ep_active=ep_active, + num_ep=num_ep, + fsdp_sizes=fsdp_sizes, + recv_buffer_rows=recv_buffer_rows, + ) + # ``expert_outputs_residual`` is the post-A2A FFN-output tensor that + # Step 3 of the combine actually consumed. Saving this (rather than + # the pre-A2A shard-local FFN output) is what makes + # ``_combine_bwd``'s Step-3 inverse see the same value the forward + # Step 3 saw -- otherwise EP + TRITON yields wrong d_expert_outputs. + output, expert_outputs_residual = _combine( + expert_outputs, + dispatch_state, + backend=permutation_backend, + ep_active=ep_active, + batch_size=batch_size, + sequence_length=sequence_length, + dtype=dtype, + num_experts_per_tok=num_experts_per_tok, + num_real_tokens=_static_shape.num_real_tokens, + padding_size=_static_shape.padding_size, + pre_a2a_buffer_shape=_static_shape.pre_a2a_buffer_shape, + ep_axis=ep_axis, + shard_id=shard_id, + num_ep=num_ep, + ) + + # ---------------- Build ctx ---------------- + aux_enabled = aux_loss_coeff > 0.0 + ctx = _BodyCtx( + x=x, + gate_kernel=gate_kernel, + logits_2d=logits_2d, + saved_scores=saved_scores, + routing_map=routing_map, + dispatch=dispatch_state, + casted_sorted_x_lhs_trans=casted_sorted_x_lhs_trans, + casted_wi_rhs_trans=casted_wi_rhs_trans, + gate_proj_out=gate_proj_out, + up_proj_out=up_proj_out, + casted_intermediate_lhs_trans=casted_intermediate_lhs_trans, + casted_wo_rhs_trans=casted_wo_rhs_trans, + expert_outputs=expert_outputs_residual, + local_group_sizes=local_group_sizes, + expert_bias=expert_bias if expert_bias is not None else None, + aux_const_buf=aux_const_buf if aux_enabled else None, + aux_tokens_per_expert=aux_tokens_per_expert if aux_enabled else None, + aux_logits_for_score=aux_logits_for_score if aux_enabled else None, + aux_saved_scores=aux_saved_scores if aux_enabled else None, + ) + + return output, aux_loss, ctx + + +def _body_bwd( # pylint: disable=unused-argument + ctx: _BodyCtx, + dy_pair: Tuple[jnp.ndarray, jnp.ndarray], + *, + num_experts: int, + num_experts_per_tok: int, + activation_type: str, + score_function: ScoreFunction, + use_pre_softmax: bool, + num_groups: Optional[int], + group_topk: Optional[int], + scaling_factor: float, + aux_loss_coeff: float, + permutation_backend: PermutationBackend, + align_size: int, + gate_inside_vjp: bool, + quantizer_sets: Tuple[QuantizerSet, QuantizerSet, QuantizerSet], + dtype: jnp.dtype, + ep_active: bool, + ep_axis: Optional[str], + data_parallelism_axes: Tuple[str, ...], + fsdp_sizes: Tuple[int, ...], + num_ep: int, + num_experts_local: int, + recv_buffer_rows: int, + # Static side info (kept here rather than inside ctx because they're + # python flags / shapes, not array leaves): + has_wi_bias: bool, + has_wo_bias: bool, + has_expert_bias: bool, + x_shape: Tuple[int, ...], +) -> dict: + """Per-shard backward body. Returns a dict of grads keyed identically + to the ``captured`` dict consumed by :func:`_body_fwd`.""" + if not gate_inside_vjp: + raise NotImplementedError("gate_inside_vjp=False is deferred to a follow-up PR.") + + d_output, d_aux_loss = dy_pair + # The fused FFN bwd quantizes via ``q_set_w0`` only (one quantize for + # the [E, H, 2M] fused wi tensor and one for the [T, 2M] fused dgrad), + # so ``q_set_w1`` is intentionally unused here. + q_set_w0, _q_set_w1, q_set_wo = quantizer_sets + batch_size, sequence_length, hidden = x_shape + shard_id = jax.lax.axis_index(ep_axis) if ep_active else None + + # Recompute per-shard static shape info from existing statics + # (Python ints / int tuples). Plumbed via kwargs to _combine_bwd + # and _dispatch_bwd -- NOT through the ctx dict, because the + # dict gets pytree-flattened across the bwd shard_map's in_specs + # and Python ints would be coerced into JitTracer 0-d arrays + # (breaking ``if padding > 0`` and ``jnp.zeros(shape)`` callsites). + # ``batch_size`` here is the GLOBAL batch size (captured in + # ``x_shape`` by the outer fwd rule), hence ``batch_is_per_shard=False``. + _static_shape = _compute_static_shape_info( + batch_size=batch_size, + sequence_length=sequence_length, + hidden=hidden, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + align_size=align_size, + ep_active=ep_active, + num_ep=num_ep, + fsdp_sizes=fsdp_sizes, + recv_buffer_rows=recv_buffer_rows, + batch_is_per_shard=False, + ) + + # Compute per-shard input shape: under the EP shard_map body, the + # gradient tensors live at per-shard shape, so the dispatch_bwd + # reshape target and ``d_x_from_dispatch.reshape(x_shape)`` below + # must use the per-shard shape rather than the captured global + # ``x_shape``. + if ep_active: + dp_size = math.prod(fsdp_sizes) if fsdp_sizes else 1 + per_shard_batch = batch_size // (num_ep * dp_size) + per_shard_x_shape: Tuple[int, ...] = (per_shard_batch, sequence_length, hidden) + else: + per_shard_x_shape = x_shape + + # ---------------- Combine bwd ---------------- + d_expert_outputs, d_routing_weights = _combine_bwd( + d_output, + ctx.dispatch, + ctx.expert_outputs, + backend=permutation_backend, + ep_active=ep_active, + batch_size=batch_size, + sequence_length=sequence_length, + dtype=dtype, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + num_real_tokens=_static_shape.num_real_tokens, + padding_size=_static_shape.padding_size, + post_a2a_buffer_shape=_static_shape.post_a2a_buffer_shape, + ep_axis=ep_axis, + shard_id=shard_id, + num_ep=num_ep, + ) + + # ---------------- FFN bwd: GEMM 3 (wo) ---------------- + casted_d_eo = tex.grouped_quantize( + d_expert_outputs, q_set_wo.dgrad, ctx.local_group_sizes, flatten_axis=-1 + ) + d_intermediate = tex.grouped_gemm( + casted_d_eo.get_tensor(usage=TensorUsage.LHS), + ctx.casted_wo_rhs_trans, + contracting_dims=((1,), (2,)), + ) + d_wo = tex.grouped_gemm( + ctx.casted_intermediate_lhs_trans, + casted_d_eo.get_tensor(usage=TensorUsage.RHS), + contracting_dims=((0,), (0,)), + ) + d_wo_bias = tex.grouped_dbias(d_expert_outputs, ctx.local_group_sizes) if has_wo_bias else None + + # ---------------- Activation bwd ---------------- + # intermediate = act(gate_proj_out) * up_proj_out + # d(gate_proj_out) = vjp(act, gate_proj_out)(d_intermediate * up_proj_out) + # d(up_proj_out) = d_intermediate * act(gate_proj_out) + act_fn = _convert_to_activation_function(activation_type) + act_gate_proj_out, dact_gate_proj_pullback = jax.vjp(act_fn, ctx.gate_proj_out) + d_up_proj_out = d_intermediate * act_gate_proj_out + (d_gate_proj_out,) = dact_gate_proj_pullback(d_intermediate * ctx.up_proj_out) + + # ---------------- FFN bwd: GEMM 1+2 fused (wi_0 | wi_1) ---------------- + # Concat the two upstream grads along the output (M) axis, do one + # grouped quantize + one dgrad GEMM + one wgrad GEMM, then split. + # ``ctx.casted_wi_rhs_trans`` has shape [E, H, 2M] from the fwd + # fused quantize, so the dgrad math is: + # d_sorted_x = [d_gate | d_up] @ wi_rhs_trans + # = d_gate @ wi_0^T + d_up @ wi_1^T + inter_M = d_gate_proj_out.shape[-1] + d_combined = jnp.concatenate([d_gate_proj_out, d_up_proj_out], axis=-1) + casted_d_combined = tex.grouped_quantize( + d_combined, q_set_w0.dgrad, ctx.local_group_sizes, flatten_axis=-1 + ) + d_sorted_x = tex.grouped_gemm( + casted_d_combined.get_tensor(usage=TensorUsage.LHS), + ctx.casted_wi_rhs_trans, + contracting_dims=((1,), (2,)), + ) + d_wi_combined = tex.grouped_gemm( + ctx.casted_sorted_x_lhs_trans, + casted_d_combined.get_tensor(usage=TensorUsage.RHS), + contracting_dims=((0,), (0,)), + ) + d_wi_0 = d_wi_combined[..., :inter_M] + d_wi_1 = d_wi_combined[..., inter_M:] + if has_wi_bias: + d_wi_combined_bias = tex.grouped_dbias(d_combined, ctx.local_group_sizes) + d_wi_0_bias = d_wi_combined_bias[..., :inter_M] + d_wi_1_bias = d_wi_combined_bias[..., inter_M:] + else: + d_wi_0_bias = None + d_wi_1_bias = None + + # ---------------- Dispatch bwd ---------------- + inputs_2d_shape = (per_shard_x_shape[0] * per_shard_x_shape[1], hidden) + d_inputs_2d = _dispatch_bwd( + d_sorted_x, + ctx.dispatch, + inputs_2d_shape=inputs_2d_shape, + backend=permutation_backend, + ep_active=ep_active, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + num_real_tokens=_static_shape.num_real_tokens, + padding_size=_static_shape.padding_size, + pre_a2a_buffer_shape=_static_shape.pre_a2a_buffer_shape, + ep_axis=ep_axis, + shard_id=shard_id, + num_ep=num_ep, + ) + d_x_from_dispatch = d_inputs_2d.reshape(per_shard_x_shape) + + # ---------------- Routing bwd ---------------- + # The probs cotangent comes from _combine_bwd. For PURE_JAX it's the + # cotangent of routing_weights (post-routing_map_to_selected_experts); + # we need to bridge back to sparse_probs. For TRITON it's already the + # cotangent of merging_probs == sparse_probs. + if d_routing_weights is not None: + if permutation_backend is PermutationBackend.PURE_JAX: + # routing_map_to_selected_experts: + # selected_experts = argsort(routing_map)[..., -topk:] + # weights = take_along_axis(sparse_probs, selected_experts, axis=-1) + # routing_map is bool (non-diff); the gradient of weights + # w.r.t. sparse_probs is a scatter-into-zero along the + # selected_experts indices. + selected_experts = jnp.argsort(ctx.routing_map, axis=-1)[..., -num_experts_per_tok:] + d_sparse_probs = jnp.zeros_like(ctx.saved_scores).astype(d_routing_weights.dtype) + d_sparse_probs = jnp.take_along_axis(d_sparse_probs, selected_experts, axis=-1) + # Actually scatter: build via jnp.zeros + .at[].set + d_sparse_probs = jnp.zeros(ctx.routing_map.shape, dtype=d_routing_weights.dtype) + d_sparse_probs = d_sparse_probs.at[ + jnp.arange(ctx.routing_map.shape[0])[:, None], selected_experts + ].set(d_routing_weights) + else: + d_sparse_probs = d_routing_weights.astype(jnp.float32) + else: + d_sparse_probs = jnp.zeros(ctx.routing_map.shape, dtype=jnp.float32) + + # Topk bwd primitive: returns d_logits (no d_expert_bias). + d_logits_2d_main = tex.fused_topk_with_score_function_bwd( + ctx.routing_map, + ctx.saved_scores, + d_sparse_probs.astype(ctx.saved_scores.dtype), + topk=num_experts_per_tok, + use_pre_softmax=use_pre_softmax, + scaling_factor=scaling_factor, + score_function=score_function, + compute_aux_scores=False, + ) + + # ---------------- Aux loss bwd ---------------- + if aux_loss_coeff > 0.0: + # Step 1: aux_loss bwd -> d_aux_probs + aux_num_tokens = ctx.aux_logits_for_score.shape[0] + d_aux_probs = tex.fused_moe_aux_loss_bwd( + ctx.aux_const_buf, + ctx.aux_tokens_per_expert.astype(jnp.int32), + d_aux_loss.reshape(()), + num_tokens=aux_num_tokens, + ) + # Step 2: aux-side topk bwd (compute_aux_scores=True path). + # The routing_map argument is ignored in this branch (the kernel + # uses saved_scores); pass any shape-correct integer tensor. + d_aux_logits = tex.fused_topk_with_score_function_bwd( + jnp.zeros(ctx.aux_logits_for_score.shape, dtype=jnp.bool_), + ctx.aux_saved_scores, + d_aux_probs.astype(ctx.aux_saved_scores.dtype), + topk=num_experts_per_tok, + use_pre_softmax=False, + scaling_factor=1.0, + score_function=score_function, + compute_aux_scores=True, + ) + # Step 3: under EP the aux logits were all_gathered along + # ``(ep_axis, *data_parallelism_axes)`` (the latter being FSDP + # axes that shard the batch). The bwd is the inverse of that + # multi-axis tiled all_gather: ``dynamic_slice`` to pick out + # this shard's local rows from the global cotangent. + # + # JAX's convention for tiled ``all_gather(axis_name=(a, b, ...))`` + # is row-major over the tuple: the shard at mesh position + # ``(i_a, i_b, ...)`` writes to rows + # ``[(i_a * size_b * ... + i_b * ... + ...) * local_T : + # + local_T)``. We invert that by computing the same flat + # index here and slicing. + if ep_active: + local_T_aux = ctx.logits_2d.shape[0] + flat_shard = shard_id # ep is the outermost axis in the gather tuple + for ax, sz in zip(data_parallelism_axes, fsdp_sizes): + flat_shard = flat_shard * sz + jax.lax.axis_index(ax) + d_aux_logits_local = jax.lax.dynamic_slice( + d_aux_logits.astype(ctx.logits_2d.dtype), + start_indices=(flat_shard * local_T_aux, 0), + slice_sizes=(local_T_aux, num_experts), + ) + else: + d_aux_logits_local = d_aux_logits.astype(d_logits_2d_main.dtype) + d_logits_2d = d_logits_2d_main + d_aux_logits_local.astype(d_logits_2d_main.dtype) + else: + d_logits_2d = d_logits_2d_main + + # ---------------- Gate bwd ---------------- + d_gate_logits = d_logits_2d.reshape(per_shard_x_shape[0], per_shard_x_shape[1], num_experts) + gate_kernel_cast = ctx.gate_kernel.astype(ctx.x.dtype) + d_x_from_gate = jnp.einsum("bse,he->bsh", d_gate_logits, gate_kernel_cast) + d_gate_kernel = jnp.einsum("bsh,bse->he", ctx.x, d_gate_logits).astype(ctx.gate_kernel.dtype) + d_x = d_x_from_gate + d_x_from_dispatch + + # Reduce per-rank partial contributions to match the out_specs + # declared by _build_grads_specs: + # gate_kernel : P() -> psum across (ep, *fsdp) + # wi_0/wi_1/wo : P(ep_axis, ...) -> psum across (*fsdp) only + # inputs : P((ep, fsdp), ...) -> already shard-local, no reduction + if ep_active: + replicate_all = (ep_axis,) + tuple(data_parallelism_axes) + d_gate_kernel = jax.lax.psum(d_gate_kernel, axis_name=replicate_all) + if data_parallelism_axes: + replicate_fsdp = tuple(data_parallelism_axes) + d_wi_0 = jax.lax.psum(d_wi_0, axis_name=replicate_fsdp) + d_wi_1 = jax.lax.psum(d_wi_1, axis_name=replicate_fsdp) + d_wo = jax.lax.psum(d_wo, axis_name=replicate_fsdp) + if has_wi_bias: + d_wi_0_bias = jax.lax.psum(d_wi_0_bias, axis_name=replicate_fsdp) + d_wi_1_bias = jax.lax.psum(d_wi_1_bias, axis_name=replicate_fsdp) + if has_wo_bias: + d_wo_bias = jax.lax.psum(d_wo_bias, axis_name=replicate_fsdp) + + grads: dict = { + "inputs": d_x, + "gate_kernel": d_gate_kernel, + "wi_0": d_wi_0, + "wi_1": d_wi_1, + "wo": d_wo, + } + if has_wi_bias: + grads["wi_0_bias"] = d_wi_0_bias + grads["wi_1_bias"] = d_wi_1_bias + if has_wo_bias: + grads["wo_bias"] = d_wo_bias + if has_expert_bias: + # expert_bias has no gradient through topk (the topk bwd returns + # None for it). Emit a structural zero so the outer rule has + # something to package. + grads["expert_bias"] = jnp.zeros_like(ctx.expert_bias) + return grads + + +# ============================================================================= +# Spec builders for shard_map (lockstep with ctx_dict / captured_dict) +# ============================================================================= + + +def _build_in_specs( + ep_axis: str, + batch_pspec_axis: Any, + *, + has_bias: bool, + has_expert_bias: bool, +) -> dict: + """Build the ``in_specs`` dict for the EP fwd shard_map.""" + specs: dict = { + "inputs": P(batch_pspec_axis, None, None), + "gate_kernel": P(), + "wi_0": P(ep_axis, None, None), + "wi_1": P(ep_axis, None, None), + "wo": P(ep_axis, None, None), + } + if has_bias: + for name in ("wi_0_bias", "wi_1_bias", "wo_bias"): + specs[name] = P(ep_axis, None) + if has_expert_bias: + specs["expert_bias"] = P(ep_axis) + return specs + + +def _build_dispatch_specs( # pylint: disable=unused-argument + ep_axis: str, + *, + backend: PermutationBackend, + ep_active: bool, + align_size: int, +) -> _DispatchState: + """Build the shard_map ``out_specs`` for the dispatch state. + + Returns a :data:`_DispatchState` (either :class:`_PureJaxDispatchState` + or :class:`_TritonDispatchState`) whose fields are + :class:`PartitionSpec` placeholders. Optional fields are set to + ``P()`` when populated by :func:`_dispatch` and to ``None`` when + intentionally omitted, so the spec's pytree structure mirrors the + value's structure leaf-for-leaf. + """ + ep_all = P() if ep_active else None + ep_local = P() if ep_active else None + if backend is PermutationBackend.PURE_JAX: + return _PureJaxDispatchState( + group_sizes=P(), + sorted_indices=P(), + routing_weights=P(), + all_shards_tokens_per_expert=ep_all, + local_perm_row_id_map=ep_local, + ) + return _TritonDispatchState( + group_sizes=P(), + row_id_map=P(), + pad_offsets=P() if align_size > 0 else None, + merging_probs=P(), + all_shards_tokens_per_expert=ep_all, + local_perm_row_id_map=ep_local, + ) + + +def _build_ctx_specs( # pylint: disable=unused-argument + ep_axis: str, + batch_pspec_axis: Any, + *, + backend: PermutationBackend, + ep_active: bool, + has_bias: bool, + has_expert_bias: bool, + aux_loss_enabled: bool, + align_size: int, +) -> _BodyCtx: + """Build the spec :class:`_BodyCtx` mirroring :func:`_body_fwd`'s ctx. + + Fields gated off by the static config (``expert_bias``, ``aux_*``) + are ``None`` here so the spec pytree matches the value pytree + leaf-for-leaf. + """ + return _BodyCtx( + # Per-shard local activations along the batch axis. + x=P(batch_pspec_axis, None, None), + gate_kernel=P(), + logits_2d=P(batch_pspec_axis, None), + saved_scores=P(batch_pspec_axis, None), + routing_map=P(batch_pspec_axis, None), + dispatch=_build_dispatch_specs( + ep_axis, backend=backend, ep_active=ep_active, align_size=align_size + ), + # FFN residuals: the LHS_TRANS / RHS_TRANS variants of + # grouped_quantize have leading "rows"/"experts" dims that are + # already shard-local (post-dispatch). Use P(ep_axis,...) on + # leading dim; that works whether the leaf is a plain ndarray + # or a ScaledTensor (shard_map applies the spec leaf-wise to + # the registered ScaledTensor pytree). + casted_sorted_x_lhs_trans=P(), + casted_wi_rhs_trans=P(ep_axis, None, None), + gate_proj_out=P(), + up_proj_out=P(), + casted_intermediate_lhs_trans=P(), + casted_wo_rhs_trans=P(ep_axis, None, None), + expert_outputs=P(), + local_group_sizes=P(), + expert_bias=P(ep_axis) if has_expert_bias else None, + aux_const_buf=P() if aux_loss_enabled else None, + aux_tokens_per_expert=P() if aux_loss_enabled else None, + aux_logits_for_score=P() if aux_loss_enabled else None, + aux_saved_scores=P() if aux_loss_enabled else None, + ) + + +def _build_grads_specs( + ep_axis: str, + batch_pspec_axis: Any, + *, + has_bias: bool, + has_expert_bias: bool, +) -> dict: + """Spec dict for the grads dict returned by :func:`_body_bwd`.""" + return _build_in_specs( + ep_axis, + batch_pspec_axis, + has_bias=has_bias, + has_expert_bias=has_expert_bias, + ) + + +# ============================================================================= +# Top-level VJP rules +# ============================================================================= + + +def _moe_fwd_rule( # pylint: disable=unused-argument + # Args MUST match the positional order of ``_moe`` (diff first, + # then nondiff). See ``_moe_bwd_rule`` for the opposite convention. + x, + gate_kernel, + wi_0, + wi_1, + wo, + wi_0_bias, + wi_1_bias, + wo_bias, + expert_bias, + num_experts, + num_experts_per_tok, + activation_type, + score_function, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + aux_loss_coeff, + permutation_backend, + align_size, + gate_inside_vjp, + ep_axis, + data_parallelism_axes, + input_axes, + gate_kernel_axes, + wi_kernel_axes, + wo_kernel_axes, + quantizer_sets, + dtype, +): + x = with_sharding_constraint_by_logical_axes(x, input_axes) + ep_active = ep_axis is not None + body_kwargs = { + "num_experts": num_experts, + "num_experts_per_tok": num_experts_per_tok, + "activation_type": activation_type, + "score_function": score_function, + "use_pre_softmax": use_pre_softmax, + "num_groups": num_groups, + "group_topk": group_topk, + "scaling_factor": scaling_factor, + "aux_loss_coeff": aux_loss_coeff, + "permutation_backend": permutation_backend, + "align_size": align_size, + "gate_inside_vjp": gate_inside_vjp, + "quantizer_sets": quantizer_sets, + "dtype": dtype, + "ep_axis": ep_axis, + "data_parallelism_axes": data_parallelism_axes, + } + captured: dict = { + "inputs": x, + "gate_kernel": gate_kernel, + "wi_0": wi_0, + "wi_1": wi_1, + "wo": wo, + } + has_bias = wi_0_bias is not None + has_expert_bias = expert_bias is not None + if has_bias: + captured["wi_0_bias"] = wi_0_bias + captured["wi_1_bias"] = wi_1_bias + captured["wo_bias"] = wo_bias + if has_expert_bias: + captured["expert_bias"] = expert_bias + + if not ep_active: + output, aux_loss, ctx = _body_fwd( + captured, + **body_kwargs, + ep_active=False, + fsdp_sizes=(), + num_ep=1, + num_experts_local=num_experts, + recv_buffer_rows=0, + ) + # Carry static side info to the bwd rule alongside ctx. These + # are Python ints/bools/tuples (NOT pytree leaves), so we + # bundle them as a plain dict rather than putting them on the + # ``_BodyCtx`` NamedTuple where shard_map would try to flatten + # them into JitTracers. + static = { + "has_wi_bias": has_bias, + "has_wo_bias": has_bias, + "has_expert_bias": has_expert_bias, + "x_shape": x.shape, + "num_experts_local": num_experts, + "recv_buffer_rows": 0, + } + return (output, aux_loss), (ctx, static) + + # ---------------- EP path ---------------- + from jax.experimental.shard_map import shard_map + + mesh = _get_mesh() + if mesh is None or mesh.empty: + raise ValueError("moe(...) requires an active jax.sharding.Mesh when ep_axis is set.") + num_ep = mesh.shape[ep_axis] + if num_experts % num_ep != 0: + raise ValueError(f"num_experts={num_experts} must be divisible by EP size={num_ep}") + num_experts_local = num_experts // num_ep + + # Reject overlapping EP / FSDP axes. Listing ep_axis in + # data_parallelism_axes would produce a duplicate-axis PartitionSpec + # ((ep, ep, ...)) which JAX rejects, and would also double-count + # num_ep in dp_size (under-sizing recv_buffer_rows by a factor of + # num_ep). Catch it up front with a clear error. + for ax in data_parallelism_axes: + if ax not in mesh.shape: + raise ValueError( + f"data_parallelism_axes contains {ax!r} but mesh has" + f" axes {tuple(mesh.shape.keys())}" + ) + if ax == ep_axis: + raise ValueError( + f"data_parallelism_axes={data_parallelism_axes!r} contains the EP" + f" axis {ep_axis!r}; EP is implicit in the batch sharding and must" + " not also be listed as a data-parallel axis." + ) + + if not data_parallelism_axes: + batch_pspec_axis: Any = ep_axis + else: + batch_pspec_axis = (ep_axis, *data_parallelism_axes) + dp_size = 1 + for ax in data_parallelism_axes: + dp_size *= mesh.shape[ax] + + global_batch_size, sequence_length, _hidden = x.shape + topk = num_experts_per_tok + if global_batch_size % (num_ep * dp_size) != 0: + raise ValueError(f"batch={global_batch_size} not divisible by ep*dp={num_ep * dp_size}") + recv_buffer_rows = (global_batch_size // dp_size) * sequence_length * topk + if align_size > 0: + recv_buffer_rows += num_experts * (align_size - 1) + + in_specs = _build_in_specs( + ep_axis, + batch_pspec_axis, + has_bias=has_bias, + has_expert_bias=has_expert_bias, + ) + output_spec = P(batch_pspec_axis, None, None) + aux_spec = P() + ctx_spec = _build_ctx_specs( + ep_axis, + batch_pspec_axis, + backend=permutation_backend, + ep_active=True, + has_bias=has_bias, + has_expert_bias=has_expert_bias, + aux_loss_enabled=(aux_loss_coeff > 0.0), + align_size=align_size, + ) + + _fsdp_sizes: Tuple[int, ...] = tuple(mesh.shape[ax] for ax in data_parallelism_axes) + + def _shardmap_body(captured_local): + return _body_fwd( + captured_local, + **body_kwargs, + ep_active=True, + fsdp_sizes=_fsdp_sizes, + num_ep=num_ep, + num_experts_local=num_experts_local, + recv_buffer_rows=recv_buffer_rows, + ) + + output, aux_loss, ctx = shard_map( + _shardmap_body, + mesh=mesh, + in_specs=(in_specs,), + out_specs=(output_spec, aux_spec, ctx_spec), + check_rep=False, + )(captured) + static = { + "has_wi_bias": has_bias, + "has_wo_bias": has_bias, + "has_expert_bias": has_expert_bias, + "x_shape": x.shape, + "num_experts_local": num_experts_local, + "recv_buffer_rows": recv_buffer_rows, + } + return (output, aux_loss), (ctx, static) + + +def _moe_bwd_rule( + num_experts, + num_experts_per_tok, + activation_type, + score_function, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + aux_loss_coeff, + permutation_backend, + align_size, + gate_inside_vjp, + ep_axis, + data_parallelism_axes, + input_axes, + gate_kernel_axes, + wi_kernel_axes, + wo_kernel_axes, + quantizer_sets, + dtype, + ctx, + dy_pair, +): + ctx, static = ctx # split tensor residuals from static side info + has_wi_bias = static["has_wi_bias"] + has_wo_bias = static["has_wo_bias"] + has_expert_bias = static["has_expert_bias"] + x_shape = static["x_shape"] + num_experts_local = static["num_experts_local"] + recv_buffer_rows = static["recv_buffer_rows"] + + ep_active = ep_axis is not None + mesh = _get_mesh() if ep_active else None + fsdp_sizes: Tuple[int, ...] = ( + tuple(mesh.shape[ax] for ax in data_parallelism_axes) if ep_active else () + ) + body_kwargs = { + "num_experts": num_experts, + "num_experts_per_tok": num_experts_per_tok, + "activation_type": activation_type, + "score_function": score_function, + "use_pre_softmax": use_pre_softmax, + "num_groups": num_groups, + "group_topk": group_topk, + "scaling_factor": scaling_factor, + "aux_loss_coeff": aux_loss_coeff, + "permutation_backend": permutation_backend, + "align_size": align_size, + "gate_inside_vjp": gate_inside_vjp, + "quantizer_sets": quantizer_sets, + "dtype": dtype, + "ep_axis": ep_axis, + "data_parallelism_axes": data_parallelism_axes, + "fsdp_sizes": fsdp_sizes, + "num_ep": 1 if not ep_active else mesh.shape[ep_axis], + "num_experts_local": num_experts_local, + "recv_buffer_rows": recv_buffer_rows, + "has_wi_bias": has_wi_bias, + "has_wo_bias": has_wo_bias, + "has_expert_bias": has_expert_bias, + "x_shape": x_shape, + } + + if not ep_active: + grads = _body_bwd(ctx, dy_pair, ep_active=False, **body_kwargs) + # Apply sharding constraints on grads. + grads["gate_kernel"] = with_sharding_constraint_by_logical_axes( + grads["gate_kernel"], gate_kernel_axes + ) + grads["wi_0"] = with_sharding_constraint_by_logical_axes(grads["wi_0"], wi_kernel_axes) + grads["wi_1"] = with_sharding_constraint_by_logical_axes(grads["wi_1"], wi_kernel_axes) + grads["wo"] = with_sharding_constraint_by_logical_axes(grads["wo"], wo_kernel_axes) + grads["inputs"] = with_sharding_constraint_by_logical_axes(grads["inputs"], input_axes) + return _grads_dict_to_tuple(grads, has_wi_bias, has_wo_bias, has_expert_bias) + + from jax.experimental.shard_map import shard_map + + if not data_parallelism_axes: + batch_pspec_axis: Any = ep_axis + else: + batch_pspec_axis = (ep_axis, *data_parallelism_axes) + ctx_spec = _build_ctx_specs( + ep_axis, + batch_pspec_axis, + backend=permutation_backend, + ep_active=True, + has_bias=has_wi_bias, + has_expert_bias=has_expert_bias, + aux_loss_enabled=(aux_loss_coeff > 0.0), + align_size=align_size, + ) + dy_specs = (P(batch_pspec_axis, None, None), P()) + grads_spec = _build_grads_specs( + ep_axis, batch_pspec_axis, has_bias=has_wi_bias, has_expert_bias=has_expert_bias + ) + + def _bwd_body(ctx_local, dy_local): + return _body_bwd(ctx_local, dy_local, ep_active=True, **body_kwargs) + + grads = shard_map( + _bwd_body, + mesh=mesh, + in_specs=(ctx_spec, dy_specs), + out_specs=grads_spec, + check_rep=False, + )(ctx, dy_pair) + + grads["gate_kernel"] = with_sharding_constraint_by_logical_axes( + grads["gate_kernel"], gate_kernel_axes + ) + grads["wi_0"] = with_sharding_constraint_by_logical_axes(grads["wi_0"], wi_kernel_axes) + grads["wi_1"] = with_sharding_constraint_by_logical_axes(grads["wi_1"], wi_kernel_axes) + grads["wo"] = with_sharding_constraint_by_logical_axes(grads["wo"], wo_kernel_axes) + grads["inputs"] = with_sharding_constraint_by_logical_axes(grads["inputs"], input_axes) + return _grads_dict_to_tuple(grads, has_wi_bias, has_wo_bias, has_expert_bias) + + +def _grads_dict_to_tuple( + grads: dict, has_wi_bias: bool, has_wo_bias: bool, has_expert_bias: bool +) -> Tuple: + """Pack the body_bwd's grads dict into the positional tuple JAX expects.""" + return ( + grads["inputs"], + grads["gate_kernel"], + grads["wi_0"], + grads["wi_1"], + grads["wo"], + grads.get("wi_0_bias") if has_wi_bias else None, + grads.get("wi_1_bias") if has_wi_bias else None, + grads.get("wo_bias") if has_wo_bias else None, + grads.get("expert_bias") if has_expert_bias else None, + ) + + +# ============================================================================= +# custom_vjp + public entry +# ============================================================================= + + +@partial(jax.custom_vjp, nondiff_argnums=tuple(range(9, 29))) +def _moe( + x, + gate_kernel, + wi_0, + wi_1, + wo, + wi_0_bias, + wi_1_bias, + wo_bias, + expert_bias, + num_experts, + num_experts_per_tok, + activation_type, + score_function, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + aux_loss_coeff, + permutation_backend, + align_size, + gate_inside_vjp, + ep_axis, + data_parallelism_axes, + input_axes, + gate_kernel_axes, + wi_kernel_axes, + wo_kernel_axes, + quantizer_sets, + dtype, +): + # Call in `_moe`'s own signature order to match what JAX will pass + # the fwd rule via ``_argnums_partial``. See the comment block at + # the top of ``_moe_fwd_rule`` for why this differs from + # ``_moe_bwd_rule``'s convention. + output_pair, _ = _moe_fwd_rule( + x, + gate_kernel, + wi_0, + wi_1, + wo, + wi_0_bias, + wi_1_bias, + wo_bias, + expert_bias, + num_experts, + num_experts_per_tok, + activation_type, + score_function, + use_pre_softmax, + num_groups, + group_topk, + scaling_factor, + aux_loss_coeff, + permutation_backend, + align_size, + gate_inside_vjp, + ep_axis, + data_parallelism_axes, + input_axes, + gate_kernel_axes, + wi_kernel_axes, + wo_kernel_axes, + quantizer_sets, + dtype, + ) + return output_pair + + +_moe.defvjp(_moe_fwd_rule, _moe_bwd_rule) + + +def moe( + x: jnp.ndarray, + gate_kernel: jnp.ndarray, + wi_0: jnp.ndarray, + wi_1: jnp.ndarray, + wo: jnp.ndarray, + wi_0_bias: Optional[jnp.ndarray] = None, + wi_1_bias: Optional[jnp.ndarray] = None, + wo_bias: Optional[jnp.ndarray] = None, + expert_bias: Optional[jnp.ndarray] = None, + *, + # Architecture + num_experts: int, + num_experts_per_tok: int, + activation_type: str = "silu", + # Routing + score_function: Union[str, ScoreFunction] = "softmax", + use_pre_softmax: bool = False, + num_groups: Optional[int] = None, + group_topk: Optional[int] = None, + scaling_factor: float = 1.0, + aux_loss_coeff: float = 0.0, + # Permutation + permutation_backend: PermutationBackend = PermutationBackend.PURE_JAX, + align_size: int = 0, + # Gate placement (Phuong: "perhaps as an option") + gate_inside_vjp: bool = True, + # Parallelism (resolved by caller from MeshResource) + ep_axis: Optional[str] = None, + data_parallelism_axes: Tuple[str, ...] = (), + # Logical axes for sharding constraints + input_axes: Tuple[Optional[str], ...] = (), + gate_kernel_axes: Tuple[Optional[str], ...] = (), + wi_kernel_axes: Tuple[Optional[str], ...] = ("exp", "embed", "mlp"), + wo_kernel_axes: Tuple[Optional[str], ...] = ("exp", "mlp", "embed"), + # Quantization + quantizer_sets: Tuple[QuantizerSet, QuantizerSet, QuantizerSet] = ( + noop_quantizer_set, + noop_quantizer_set, + noop_quantizer_set, + ), + dtype: jnp.dtype = jnp.float32, +) -> Tuple[jnp.ndarray, Optional[jnp.ndarray]]: + """Run a full MoE block under a single fused custom_vjp. + + Parameters and return are documented at the call site of + ``_MoEBlock.__call__``. See module docstring for design rationale. + """ + if not isinstance(permutation_backend, PermutationBackend): + raise TypeError( + f"permutation_backend must be a PermutationBackend, got {permutation_backend!r}" + ) + if permutation_backend is PermutationBackend.TRITON: + _require_triton() + # Normalize string score_function ("softmax" / "sigmoid") to the + # ScoreFunction enum once here. The underlying primitive + # ``tex.fused_topk_with_score_function_fwd`` expects an int-coercible + # value (the enum has integer .value), and the public router wrapper + # we bypass also normalizes here. + score_function = _validate_score_function(score_function) + + output, aux_loss = _moe( + x, + gate_kernel, + wi_0, + wi_1, + wo, + wi_0_bias, + wi_1_bias, + wo_bias, + expert_bias, + num_experts=num_experts, + num_experts_per_tok=num_experts_per_tok, + activation_type=activation_type, + score_function=score_function, + use_pre_softmax=use_pre_softmax, + num_groups=num_groups, + group_topk=group_topk, + scaling_factor=scaling_factor, + aux_loss_coeff=aux_loss_coeff, + permutation_backend=permutation_backend, + align_size=align_size, + gate_inside_vjp=gate_inside_vjp, + ep_axis=ep_axis, + data_parallelism_axes=data_parallelism_axes, + input_axes=input_axes, + gate_kernel_axes=gate_kernel_axes, + wi_kernel_axes=wi_kernel_axes, + wo_kernel_axes=wo_kernel_axes, + quantizer_sets=quantizer_sets, + dtype=dtype, + ) + if aux_loss_coeff <= 0.0: + aux_loss = None + return output, aux_loss diff --git a/transformer_engine/jax/permutation.py b/transformer_engine/jax/permutation.py index 81972aac0..68283e234 100644 --- a/transformer_engine/jax/permutation.py +++ b/transformer_engine/jax/permutation.py @@ -7,6 +7,19 @@ This module provides high-level token dispatch and combine operations for Mixture of Experts (MoE) models with proper automatic differentiation support. +Two backends are offered: + +* Triton-backed ``token_dispatch`` / ``token_combine`` - uses the + Triton kernels in ``transformer_engine.jax.triton_extensions.permutation``. +* Pure-JAX ``pure_jax_token_dispatch`` / ``pure_jax_token_combine`` - uses + only ``jnp.argsort`` + gather and is therefore compiled as plain XLA. + Despite the name, this path is often *faster* than the Triton kernels in + current testing because XLA can fuse the ops with surrounding work. + +Both backends support optional alignment padding (``align_size > 0``) so each +expert's group size is a multiple of ``align_size``, which is required for +quantized grouped GEMMs. + Token Dispatch (Permute): - Forward: Permute tokens according to routing map (scatter to experts) - Backward: Unpermute gradients (gather from experts) @@ -17,27 +30,73 @@ """ from functools import partial -from typing import Optional, Tuple +from typing import NamedTuple, Optional, Tuple import jax import jax.numpy as jnp -from transformer_engine.jax.triton_extensions.permutation import ( - make_row_id_map, - permute_with_mask_map, - permute_with_mask_map_and_pad, - unpermute_with_mask_map, - unpermute_with_mask_map_and_unpad, - unpermute_bwd_with_merging_probs, - unpermute_bwd_with_merging_probs_and_unpad, - make_chunk_sort_map, - sort_chunks_by_map, -) +# Triton-backed primitives are imported lazily: they require ``triton`` +# which we do not want as a hard install dependency for the pure-JAX +# permutation backend. Anything that touches one of these symbols must +# either be guarded by a TRITON-backend check or live inside a function +# that is only reachable from the TRITON path. +try: + from transformer_engine.jax.triton_extensions.permutation import ( + make_row_id_map, + permute_with_mask_map, + permute_with_mask_map_and_pad, + unpermute_with_mask_map, + unpermute_with_mask_map_and_unpad, + unpermute_bwd_with_merging_probs, + unpermute_bwd_with_merging_probs_and_unpad, + make_chunk_sort_map, + sort_chunks_by_map, + ) + + _TRITON_PERMUTATION_AVAILABLE = True +except ImportError: + _TRITON_PERMUTATION_AVAILABLE = False + make_row_id_map = None + permute_with_mask_map = None + permute_with_mask_map_and_pad = None + unpermute_with_mask_map = None + unpermute_with_mask_map_and_unpad = None + unpermute_bwd_with_merging_probs = None + unpermute_bwd_with_merging_probs_and_unpad = None + make_chunk_sort_map = None + sort_chunks_by_map = None + + +def _require_triton_permutation(): + """Raise a clear error if Triton permutation kernels are unavailable. + + Callers in the TRITON branch must invoke this before touching any of + the names imported above; on a Triton-less install those names are + ``None`` and would otherwise produce a confusing ``TypeError`` deep + in the call stack. + """ + if not _TRITON_PERMUTATION_AVAILABLE: + raise ImportError( + "TRITON permutation backend requires" + " ``transformer_engine.jax.triton_extensions.permutation`` (which" + " in turn requires ``triton``). Either install Triton or use" + " PermutationBackend.PURE_JAX." + ) + __all__ = [ "token_dispatch", "token_combine", "sort_chunks_by_index", + "pure_jax_token_dispatch", + "pure_jax_token_combine", + "PureJaxPermState", + # Ragged-all-to-all expert-parallelism helpers + "compute_ragged_all_to_all_params", + "compute_reverse_ragged_all_to_all_params", + "local_permute_after_a2a", + "local_unpermute_before_a2a", + "routing_map_to_selected_experts", ] @@ -655,3 +714,640 @@ def _sort_chunks_by_index_bwd_rule( _sort_chunks_by_index.defvjp(_sort_chunks_by_index_fwd_rule, _sort_chunks_by_index_bwd_rule) + + +# ============================================================================= +# Pure-JAX token dispatch / combine +# ============================================================================= +# +# The following implementations use only ``jnp.argsort`` + gather and compile +# to plain XLA. They are a drop-in alternative to ``token_dispatch`` / +# ``token_combine`` above, differing only in input/output conventions (the +# Triton path takes ``routing_map`` and ``sparse_probs`` over all experts; the +# pure-JAX path takes dense ``selected_experts`` and per-token ``weights`` of +# shape ``[..., topk]``). +# +# Note: despite Triton being fused and pure-JAX being a sequence of XLA ops, +# the pure-JAX backend is often *faster* in current testing because XLA can +# fuse these ops into the surrounding work. + + +# ----------------------------------------------------------------------------- +# Custom-VJP argsort-based gather. +# +# ``inputs[sort_indices]`` has a known inverse: ``output[argsort(sort_indices)]``. +# Using a custom VJP lets the backward pass exploit that inverse instead of +# relying on the compiler to discover it from the scatter-style default +# gradient of a gather, which is typically less efficient. + + +@jax.custom_vjp +def _sort_activations(inputs: jax.Array, sort_indices: jax.Array) -> jax.Array: + """Sort ``inputs`` along the leading dim by ``sort_indices``.""" + assert ( + inputs.shape[0] == sort_indices.shape[0] + ), f"inputs.shape[0]={inputs.shape[0]} must match sort_indices.shape[0]={sort_indices.shape[0]}" + with jax.named_scope("pure_jax_sort_activations"): + return inputs[sort_indices, ...] + + +def _sort_activations_fwd( + inputs: jax.Array, sort_indices: jax.Array +) -> Tuple[jax.Array, jax.Array]: + return _sort_activations(inputs, sort_indices), sort_indices + + +def _sort_activations_bwd(residuals: jax.Array, grads: jax.Array) -> Tuple[jax.Array, None]: + sort_indices = residuals + # Inverse permutation: gather-by-argsort undoes the forward gather. + return _sort_activations(grads, jnp.argsort(sort_indices)), None + + +_sort_activations.defvjp(_sort_activations_fwd, _sort_activations_bwd) + + +def routing_map_to_selected_experts( + sparse_probs: jnp.ndarray, + routing_map: jnp.ndarray, + topk: int, +) -> Tuple[jnp.ndarray, jnp.ndarray]: + """Convert ``(sparse_probs, routing_map)`` from TE's fused router to the + ``(selected_experts, weights)`` format consumed by + :func:`pure_jax_token_dispatch`. + + ``routing_map`` is a boolean mask of shape ``[num_tokens, num_experts]`` + with exactly ``topk`` ``True`` positions per row. + """ + # Argsort on a bool tensor places ``True`` rows last (False=0 < True=1), + # so the last ``topk`` indices are the selected expert IDs. + selected_experts = jnp.argsort(routing_map, axis=-1)[..., -topk:] + weights = jnp.take_along_axis(sparse_probs, selected_experts, axis=-1) + return selected_experts, weights + + +# ----------------------------------------------------------------------------- +# Permutation state carried from dispatch to combine. + + +class PureJaxPermState(NamedTuple): + """Opaque state produced by :func:`pure_jax_token_dispatch`. + + Attributes + ---------- + sorted_indices : jnp.ndarray + The argsort indices used in the forward sort. Needed to reverse the + permutation in :func:`pure_jax_token_combine`. Shape + ``[num_real_tokens + padding_size]``. + num_real_tokens : int + Number of real (non-padding) permuted tokens, i.e. + ``batch_size * sequence_length * num_experts_per_tok``. Compile-time + constant. + padding_size : int + Number of alignment-padding tokens appended to the sort buffer. Equals + ``num_experts * (align_size - 1)`` when ``align_size > 0``, else ``0``. + Compile-time constant. + """ + + sorted_indices: jax.Array + num_real_tokens: int + padding_size: int + + +# ----------------------------------------------------------------------------- +# Dispatch (permute) + + +def pure_jax_token_dispatch( + inputs: jnp.ndarray, + selected_experts: jnp.ndarray, + num_experts: int, + num_experts_per_tok: int, + align_size: int = 0, + roll_to_expert_id: Optional[int] = None, +) -> Tuple[jnp.ndarray, PureJaxPermState, jnp.ndarray]: + """Pure-JAX ``argsort``-based token dispatch. + + Parameters + ---------- + inputs : jnp.ndarray + Input tensor of shape ``[num_tokens, hidden_size]`` (or + ``[batch, seq, hidden]``; it will be flattened). + selected_experts : jnp.ndarray + Per-token expert IDs, shape ``[num_tokens, num_experts_per_tok]`` (or + ``[batch, seq, num_experts_per_tok]``). Integer dtype. + num_experts : int + Total number of experts. + num_experts_per_tok : int + Top-k. Must equal ``selected_experts.shape[-1]``. + align_size : int, default 0 + Alignment for each expert's group size. ``0`` disables padding; a value + ``> 0`` appends a static-size padding buffer so each resulting group + size is a multiple of ``align_size`` (required for quantized grouped + GEMM). + roll_to_expert_id : Optional[int] + If provided, rotates expert IDs by ``-roll_to_expert_id`` modulo + ``num_experts`` before the sort (ring-of-experts EP). The returned + ``group_sizes`` is rolled to match. + + Returns + ------- + sorted_inputs : jnp.ndarray + Permuted tokens grouped by expert, shape + ``[num_real_tokens + padding_size, hidden_size]``. + perm_state : PureJaxPermState + State needed by :func:`pure_jax_token_combine`. + group_sizes : jnp.ndarray + Token count per expert, shape ``[num_experts]``. Each entry is a + multiple of ``align_size`` when ``align_size > 0``. + """ + assert num_experts_per_tok == selected_experts.shape[-1], ( + f"num_experts_per_tok={num_experts_per_tok} must match" + f" selected_experts.shape[-1]={selected_experts.shape[-1]}" + ) + assert align_size >= 0, f"align_size must be >= 0, got {align_size}" + + hidden_size = inputs.shape[-1] + inputs_2d = inputs.reshape(-1, hidden_size) + num_tokens = inputs_2d.shape[0] + num_real_tokens = num_tokens * num_experts_per_tok + + flatten_selected_experts = jnp.ravel(selected_experts) + + if align_size > 0: + # Per-expert token count, and how many extra tokens each expert needs + # to become aligned to ``align_size``. Using + # ``(align - count % align) % align`` gives 0 (not ``align``) when + # already aligned, so we never exceed the per-expert slot capacity of + # ``align_size - 1``. + token_count_per_expert = jnp.bincount(flatten_selected_experts, length=num_experts) + padding_tokens_required_per_expert = ( + align_size - (token_count_per_expert % align_size) + ) % align_size + + # Build a static-size padding buffer of shape + # ``[num_experts * (align_size - 1)]``. Each expert ``i`` owns a slot + # of ``align_size - 1`` positions (worst-case padding, which occurs + # when ``token_count[i] % align_size == 1``). Within slot ``i``, + # positions ``[0, padding_needed)`` are assigned expert ``i`` and act + # as real padding; the rest are assigned to ``num_experts - 1`` as + # overflow placeholders that keep the buffer statically sized for JIT. + max_padding_per_expert = align_size - 1 + max_total_padding_size = num_experts * max_padding_per_expert + positions = jnp.arange(max_total_padding_size) + expert_for_pos = positions // max_padding_per_expert + offset_in_slot = positions % max_padding_per_expert + padding_needed = padding_tokens_required_per_expert[expert_for_pos] + flatten_padding_selected_experts = jnp.where( + offset_in_slot < padding_needed, + expert_for_pos, + num_experts - 1, + ) + + flatten_selected_experts = jnp.concatenate( + [flatten_selected_experts, flatten_padding_selected_experts], axis=0 + ) + + if roll_to_expert_id is not None: + flatten_selected_experts = (flatten_selected_experts - roll_to_expert_id) % num_experts + + sorted_selected_experts = jnp.argsort(flatten_selected_experts) + + replicated_inputs_2d = jnp.repeat(inputs_2d, num_experts_per_tok, axis=0) + # Pad inputs with zeros so the sort operand shape matches the expanded + # selected-experts vector. + replicated_inputs_2d = jnp.pad( + replicated_inputs_2d, + pad_width=((0, max_total_padding_size), (0, 0)), + mode="constant", + constant_values=0.0, + ) + + sorted_inputs = _sort_activations(replicated_inputs_2d, sorted_selected_experts) + + # Compute ``group_sizes`` directly from counts rather than via + # ``bincount(flatten_selected_experts)``: the overflow placeholder + # tokens would inflate ``group_sizes[num_experts - 1]``, breaking the + # alignment guarantee. Direct computation gives each expert exactly + # ``ceil(count / align) * align`` tokens. + group_sizes = token_count_per_expert + padding_tokens_required_per_expert + + if roll_to_expert_id is not None: + group_sizes = jnp.roll(group_sizes, -roll_to_expert_id) + + padding_size = max_total_padding_size + else: + if roll_to_expert_id is not None: + flatten_selected_experts = (flatten_selected_experts - roll_to_expert_id) % num_experts + + sorted_selected_experts = jnp.argsort(flatten_selected_experts) + + replicated_inputs_2d = jnp.repeat(inputs_2d, num_experts_per_tok, axis=0) + sorted_inputs = _sort_activations(replicated_inputs_2d, sorted_selected_experts) + + group_sizes = jnp.bincount(flatten_selected_experts, length=num_experts) + if roll_to_expert_id is not None: + group_sizes = jnp.roll(group_sizes, -roll_to_expert_id) + + padding_size = 0 + + perm_state = PureJaxPermState( + sorted_indices=sorted_selected_experts, + num_real_tokens=num_real_tokens, + padding_size=padding_size, + ) + return sorted_inputs, perm_state, group_sizes + + +# ----------------------------------------------------------------------------- +# Combine (unpermute + weighted sum) + + +def pure_jax_token_combine( + expert_outputs: jnp.ndarray, + perm_state: PureJaxPermState, + routing_weights: jnp.ndarray, + num_experts_per_tok: int, + batch_size: int, + sequence_length: int, +) -> jnp.ndarray: + """Pure-JAX ``argsort``-based token combine. + + Reverses the permutation performed by :func:`pure_jax_token_dispatch`, + strips any alignment-padding rows appended during dispatch, and applies a + per-token weighted sum across the top-k experts. + + Parameters + ---------- + expert_outputs : jnp.ndarray + Output of the expert FFN, shape + ``[num_real_tokens + padding_size, hidden_size]``. + perm_state : PureJaxPermState + State returned by :func:`pure_jax_token_dispatch`. + routing_weights : jnp.ndarray + Top-k routing weights, shape ``[batch*seq, num_experts_per_tok]`` + (or broadcastable to it after a ``reshape``). + num_experts_per_tok : int + Top-k. + batch_size : int + Original batch size. + sequence_length : int + Original sequence length. + + Returns + ------- + output : jnp.ndarray + Combined output tensor of shape ``[batch_size, sequence_length, hidden_size]``. + """ + # Reverse the permutation: ``output[argsort(sorted_indices)]`` undoes + # ``input[sorted_indices]``. + unsort_intermediate = _sort_activations( + expert_outputs, + jnp.argsort(perm_state.sorted_indices), + ) + + # Strip alignment padding tokens appended during dispatch. After unsorting, + # the first ``num_real_tokens`` rows hold the real per-(token, top-k) + # outputs; any trailing rows are padding placeholders (zeros) and must be + # discarded before the reshape below. + if perm_state.padding_size > 0: + unsort_intermediate = unsort_intermediate[: perm_state.num_real_tokens] + + hidden_size = unsort_intermediate.shape[-1] + reshaped_weights = jnp.reshape(routing_weights, (-1, num_experts_per_tok)) + reshaped_intermediate = jnp.reshape( + unsort_intermediate, (reshaped_weights.shape[0], num_experts_per_tok, hidden_size) + ) + + # Cast weights to match intermediate dtype (weighted sum happens in + # intermediate dtype; callers can upcast before calling if higher + # precision weight-sum is desired). + reshaped_weights = reshaped_weights.astype(reshaped_intermediate.dtype) + with jax.named_scope("pure_jax_weight_sum"): + output = jnp.einsum( + "BKE,BK -> BE", + reshaped_intermediate, + reshaped_weights, + ) + return output.reshape(batch_size, sequence_length, hidden_size) + + +# ============================================================================= +# Ragged-all-to-all expert-parallelism helpers +# ============================================================================= +# +# These helpers support the ragged-all-to-all (A2A / A2Av) EP strategy used by +# :class:`transformer_engine.jax.flax._MoEBlock`. The forward EP path looks +# like:: +# +# route -> global_permute -> AG(group_sizes, ep) +# -> ragged_all_to_all(fwd, ep) +# -> local_permute_after_a2a +# -> grouped_dense x3 + activation +# -> local_unpermute_before_a2a +# -> ragged_all_to_all(reverse, ep) +# -> global_combine +# +# The two ``compute_*_ragged_all_to_all_params`` functions translate +# ``all_shards_tokens_per_expert`` (an EP-axis ``all_gather`` of each shard's +# global ``group_sizes``) into the four ``ragged_all_to_all`` arguments +# (``input_offsets``, ``send_sizes``, ``output_offsets``, ``recv_sizes``). +# ``shard_id`` may be a traced value (e.g. from :func:`jax.lax.axis_index`), +# which is why every slice into ``all_shards_tokens_per_expert`` uses +# :func:`jax.lax.dynamic_slice`. +# +# These functions are pure JAX (no TE-internal dependencies). + + +def compute_ragged_all_to_all_params( + all_shards_tokens_per_expert: jnp.ndarray, + shard_id: jnp.ndarray, + num_expert_shards: int, +) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Forward-direction ragged_all_to_all parameters. + + Computes the four index/size arrays that :func:`jax.lax.ragged_all_to_all` + consumes for the **forward** EP shuffle, where each shard sends its + expert-grouped tokens to the shard that owns those experts. + + Parameters + ---------- + all_shards_tokens_per_expert : jnp.ndarray + Per-shard, per-expert token counts gathered across the EP axis. Shape + ``[num_expert_shards, num_experts]`` and integer dtype. + shard_id : jnp.ndarray + Index of the current shard along the EP axis (typically + :func:`jax.lax.axis_index` of the EP axis). Must be a 0-d integer. + num_expert_shards : int + Static EP-axis size. Must match + ``all_shards_tokens_per_expert.shape[0]``. + + Returns + ------- + input_offsets : jnp.ndarray + Shape ``[num_expert_shards]``. Cumulative ``send_sizes`` (with a + leading 0) -- where in the local source buffer each destination + shard's chunk begins. + send_sizes : jnp.ndarray + Shape ``[num_expert_shards]``. ``send_sizes[i]`` is the number of + tokens this shard sends to shard ``i`` (= the sum of token counts + for the experts owned by shard ``i``). + output_offsets : jnp.ndarray + Shape ``[num_expert_shards]``. ``output_offsets[i]`` is the row in + shard ``i``'s receive buffer where this shard's contribution should + land. Sender-side semantics, per :func:`jax.lax.ragged_all_to_all`. + recv_sizes : jnp.ndarray + Shape ``[num_expert_shards]``. ``recv_sizes[i]`` is the number of + tokens shard ``i`` sends to this shard. + """ + num_experts = all_shards_tokens_per_expert.shape[1] + assert ( + num_experts % num_expert_shards == 0 + ), f"num_experts={num_experts} must be divisible by num_expert_shards={num_expert_shards}" + local_expert_size = num_experts // num_expert_shards + + # This shard's row of the gathered table, reshaped so axis 0 indexes the + # destination shard and axis 1 indexes its local experts. + local_tokens_per_expert = jax.lax.dynamic_slice( + all_shards_tokens_per_expert, + start_indices=(shard_id, 0), + slice_sizes=(1, num_experts), + ).squeeze(0) + local_reshaped = local_tokens_per_expert.reshape(num_expert_shards, local_expert_size) + + # send_sizes[i] = sum of token counts for shard i's experts in our buffer. + send_sizes = jnp.sum(local_reshaped, axis=1) + input_offsets = jnp.concatenate( + [ + jnp.array([0], dtype=send_sizes.dtype), + jnp.cumsum(send_sizes)[:-1], + ] + ) + + # recv_sizes[i] = how many tokens shard i sends to this shard, i.e. the + # sum across our local-expert columns of shard i's row. + local_expert_start = shard_id * local_expert_size + local_expert_columns = jax.lax.dynamic_slice( + all_shards_tokens_per_expert, + start_indices=(0, local_expert_start), + slice_sizes=(num_expert_shards, local_expert_size), + ) + recv_sizes = jnp.sum(local_expert_columns, axis=1) + + # output_offsets uses sender-side semantics for ragged_all_to_all: + # output_offsets[j] = row in shard j's buffer where THIS shard's chunk + # should be placed. That's the cumulative sum (over source shards 0..j-1) + # of how many tokens those earlier source shards already sent to shard j. + sends_to_target = jnp.sum( + all_shards_tokens_per_expert.reshape( + num_expert_shards, num_expert_shards, local_expert_size + ), + axis=2, + ) # [src_shard, dst_shard] + zero_row = jnp.zeros((1, num_expert_shards), dtype=sends_to_target.dtype) + cumulated = jnp.cumsum( + jnp.concatenate([zero_row, sends_to_target], axis=0), + axis=0, + dtype=sends_to_target.dtype, + ) # [src_shard + 1, dst_shard]; row r = total sent by sources 0..r-1 + output_offsets = jax.lax.dynamic_slice( + cumulated, + start_indices=(shard_id, 0), + slice_sizes=(1, num_expert_shards), + ).squeeze(0) + + return input_offsets, send_sizes, output_offsets, recv_sizes + + +def compute_reverse_ragged_all_to_all_params( + all_shards_tokens_per_expert: jnp.ndarray, + shard_id: jnp.ndarray, + num_expert_shards: int, +) -> Tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]: + """Reverse-direction ragged_all_to_all parameters. + + Mirror of :func:`compute_ragged_all_to_all_params` for the **reverse** + EP shuffle that returns expert outputs to their source shards. The + sender / receiver roles are swapped: what we received in the forward + shuffle we now send back, and vice versa. + + Parameters and shapes are identical to + :func:`compute_ragged_all_to_all_params`. + """ + num_experts = all_shards_tokens_per_expert.shape[1] + assert ( + num_experts % num_expert_shards == 0 + ), f"num_experts={num_experts} must be divisible by num_expert_shards={num_expert_shards}" + local_expert_size = num_experts // num_expert_shards + + local_expert_start = shard_id * local_expert_size + + # In reverse, what we received becomes what we send. send_sizes[i] is how + # many tokens we send back to source shard i (= what shard i originally + # sent us, summed across our local experts). + local_expert_columns = jax.lax.dynamic_slice( + all_shards_tokens_per_expert, + start_indices=(0, local_expert_start), + slice_sizes=(num_expert_shards, local_expert_size), + ) + send_sizes = jnp.sum(local_expert_columns, axis=1) + input_offsets = jnp.concatenate( + [ + jnp.array([0], dtype=send_sizes.dtype), + jnp.cumsum(send_sizes)[:-1], + ] + ) + + # recv_sizes[i] = how many tokens we receive back from shard i (= what + # we originally sent to shard i in the forward). + local_tokens_per_expert = jax.lax.dynamic_slice( + all_shards_tokens_per_expert, + start_indices=(shard_id, 0), + slice_sizes=(1, num_experts), + ).squeeze(0) + local_reshaped = local_tokens_per_expert.reshape(num_expert_shards, local_expert_size) + recv_sizes = jnp.sum(local_reshaped, axis=1) + + # output_offsets: the reverse sends-to-target matrix is the transpose of + # the forward one (row i = what shard i sends in reverse = what shard i + # received in forward). Cumsum down source-shard axis, then index our row. + fwd_sends_to = jnp.sum( + all_shards_tokens_per_expert.reshape( + num_expert_shards, num_expert_shards, local_expert_size + ), + axis=2, + ) # forward: [src, dst] + rev_sends_to = jnp.transpose(fwd_sends_to) # reverse: [src, dst] + zero_row = jnp.zeros((1, num_expert_shards), dtype=rev_sends_to.dtype) + rev_cumulated = jnp.cumsum( + jnp.concatenate([zero_row, rev_sends_to], axis=0), + axis=0, + dtype=rev_sends_to.dtype, + ) + output_offsets = jax.lax.dynamic_slice( + rev_cumulated, + start_indices=(shard_id, 0), + slice_sizes=(1, num_expert_shards), + ).squeeze(0) + + return input_offsets, send_sizes, output_offsets, recv_sizes + + +# ----------------------------------------------------------------------------- +# Local permute / unpermute +# ----------------------------------------------------------------------------- +# +# After the forward ragged_all_to_all the receive buffer is laid out as +# ``[from_shard_0_chunk | from_shard_1_chunk | ... ]`` and within each chunk +# tokens are sorted by local-expert id. To feed ``grouped_dense`` we want +# ``[expert_0_block | expert_1_block | ... ]`` where each expert's block +# contains tokens from every source shard. ``local_permute_after_a2a`` +# performs that reorder; ``local_unpermute_before_a2a`` undoes it before the +# reverse ragged_all_to_all. +# +# Implementation uses :func:`sort_chunks_by_index`, which is Triton-backed +# (see ``transformer_engine.jax.triton_extensions.permutation``) and has a +# paired custom-VJP backward. There is no pure-JAX alternative here -- the +# global :func:`pure_jax_token_dispatch` / :func:`token_dispatch` choice is +# unaffected by this; only the (small) post-A2A chunk reorder uses Triton +# unconditionally. + + +def local_permute_after_a2a( + x_recv: jnp.ndarray, + all_shards_tokens_per_expert: jnp.ndarray, + shard_id: jnp.ndarray, + num_expert_shards: int, +) -> Tuple[jnp.ndarray, jnp.ndarray, dict]: + """Reorder tokens received via ragged_all_to_all so each local expert's + tokens are contiguous. + + This is the EP-side complement to the global :func:`token_dispatch` / + :func:`pure_jax_token_dispatch`. Internally uses + :func:`sort_chunks_by_index` (Triton-backed) for both the forward sort + and -- via :func:`local_unpermute_before_a2a` -- the inverse. + + Parameters + ---------- + x_recv : jnp.ndarray + Output of the forward ``ragged_all_to_all`` of shape + ``[buffer_size, hidden_size]``. Layout: source-shard major, then + local-expert id within each source chunk. + all_shards_tokens_per_expert : jnp.ndarray + Per-shard, per-expert token counts of shape + ``[num_expert_shards, num_experts]``. + shard_id : jnp.ndarray + Current EP shard index (typically a traced + :func:`jax.lax.axis_index`). + num_expert_shards : int + Static EP-axis size. + + Returns + ------- + sorted_x : jnp.ndarray + Tokens reordered into expert-major layout. Same shape as ``x_recv``. + local_group_sizes : jnp.ndarray + Per-local-expert token counts of shape ``[local_expert_size]``. + state : dict + Opaque state for :func:`local_unpermute_before_a2a`. + """ + num_experts = all_shards_tokens_per_expert.shape[1] + assert ( + num_experts % num_expert_shards == 0 + ), f"num_experts={num_experts} must be divisible by num_expert_shards={num_expert_shards}" + local_expert_size = num_experts // num_expert_shards + local_expert_start = shard_id * local_expert_size + local_expert_columns = jax.lax.dynamic_slice( + all_shards_tokens_per_expert, + start_indices=(0, local_expert_start), + slice_sizes=(num_expert_shards, local_expert_size), + ) + + # Flat sizes in source-major order, matching the receive buffer layout: + # [(s0,e0), (s0,e1), ..., (s1,e0), (s1,e1), ...] + split_sizes = local_expert_columns.reshape(-1) + + # Permutation that maps source-major -> expert-major: + # original index = s * E_local + e + # target index = e * num_shards + s + indices_matrix = jnp.arange(num_expert_shards * local_expert_size, dtype=jnp.int32).reshape( + num_expert_shards, local_expert_size + ) + sorted_chunk_indices = indices_matrix.T.reshape(-1) + + sorted_x, _ = sort_chunks_by_index(x_recv, split_sizes, sorted_chunk_indices) + sorted_split_sizes = split_sizes[sorted_chunk_indices] + inverse_chunk_indices = jnp.argsort(sorted_chunk_indices) + local_group_sizes = jnp.sum(local_expert_columns, axis=0) + state = { + "sorted_split_sizes": sorted_split_sizes, + "inverse_chunk_indices": inverse_chunk_indices, + } + return sorted_x, local_group_sizes, state + + +def local_unpermute_before_a2a( + expert_outputs: jnp.ndarray, + state: dict, +) -> jnp.ndarray: + """Inverse of :func:`local_permute_after_a2a`. + + Parameters + ---------- + expert_outputs : jnp.ndarray + Output of the local expert FFN of shape ``[buffer_size, hidden_size]``, + in expert-major layout. + state : dict + Opaque state returned by :func:`local_permute_after_a2a`. + + Returns + ------- + unsorted_x : jnp.ndarray + Tokens reordered back into source-shard-major layout, ready for the + reverse ``ragged_all_to_all``. Same shape as ``expert_outputs``. + """ + out, _ = sort_chunks_by_index( + expert_outputs, + state["sorted_split_sizes"], + state["inverse_chunk_indices"], + ) + return out diff --git a/transformer_engine/jax/sharding.py b/transformer_engine/jax/sharding.py index 9b13412c1..182a4a2e0 100644 --- a/transformer_engine/jax/sharding.py +++ b/transformer_engine/jax/sharding.py @@ -332,6 +332,7 @@ class MeshResource: fsdp_resource: Axis name for full-sharded data parallelism, default is None pp_resource: Axis name for pipeline parallelism (layer sharding), default is None cp_resource: Axis name for context parallelism (sequence sharding), default is None + ep_resource: Axis name for expert parallelism (MoE expert sharding), default is None """ dp_resource: str = None @@ -340,6 +341,7 @@ class MeshResource: fsdp_resource: str = None pp_resource: str = None cp_resource: str = None + ep_resource: str = None _GLOBAL_MESH_RESOURCE = None @@ -379,6 +381,38 @@ def global_mesh_resource() -> MeshResource: return _GLOBAL_MESH_RESOURCE +def get_active_resource_axis(resource_name: str) -> Optional[str]: + """Resolve a :class:`MeshResource` attribute to its mesh axis name, + or return ``None`` if that resource is not active. + + "Active" means all three are true: + + * a physical mesh is set (``is_mesh_available()``), + * the ``MeshResource`` attribute is non-``None``, + * the corresponding mesh axis has more than 1 device. + + Mirrors the three-step ``is_X_enabled`` idiom in + :func:`get_sharding_map_logic_axis_to_mesh_axis` but returns the + axis name itself (or ``None``) so callers can use it directly in + collectives / ``shard_map`` specs. + + Args: + resource_name: Attribute name on :class:`MeshResource`, e.g. + ``"fsdp_resource"`` or ``"ep_resource"``. + + Returns: + The mesh axis name when active, else ``None``. + """ + if not is_mesh_available(): + return None + if _GLOBAL_MESH_RESOURCE is None: + return None + axis = getattr(_GLOBAL_MESH_RESOURCE, resource_name) + if axis is None or get_mesh_axis_size(axis) <= 1: + return None + return axis + + def all_reduce_sum_along_dp_fsdp(x: jnp.array, mesh: jax.sharding.Mesh): """Perform all-reduce sum operation along data parallelism and FSDP axes. From 79821e2b0f7af2eb2ba7e2e36f311d4c05bb0027 Mon Sep 17 00:00:00 2001 From: vthumbe1503 Date: Fri, 29 May 2026 15:38:00 -0700 Subject: [PATCH 093/180] [Pytorch] Skip the Single Grouped Param Test if NVTE_GROUPED_LINEAR_SINGLE_PARAM=0 (#3061) * skip the test if the env variable is not turned on for single grouped weight Signed-off-by: Varun Thumbe * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Apply suggestion from @greptile-apps[bot] Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: vthumbe1503 --------- Signed-off-by: Varun Thumbe Signed-off-by: vthumbe1503 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/pytorch/test_sanity.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/pytorch/test_sanity.py b/tests/pytorch/test_sanity.py index 27eafbecd..404fae85f 100644 --- a/tests/pytorch/test_sanity.py +++ b/tests/pytorch/test_sanity.py @@ -598,7 +598,8 @@ def test_sanity_grouped_linear( # Small batch size used to catch bug from https://github.com/NVIDIA/TransformerEngine/pull/1527. bs = bs * 16 num_tokens = bs * config.max_seqlen_q * (num_gemms - 1) - + if os.environ.get("NVTE_GROUPED_LINEAR_SINGLE_PARAM", "0") == "0" and single_param: + pytest.skip("single parameter grouped linear requires NVTE_GROUPED_LINEAR_SINGLE_PARAM=1") skip_unsupported_backward_override("grouped_linear", fp8_recipe, backward_override) if fp8_recipe is not None: fp8_recipe = copy.deepcopy(fp8_recipe) From 2055c6d2edd996363925633af76d594e058ded69 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Pawe=C5=82=20Gadzi=C5=84ski?= <62263673+pggPL@users.noreply.github.com> Date: Mon, 1 Jun 2026 11:10:14 +0200 Subject: [PATCH 094/180] [PyTorch Debug] Fix scale_inv_min returning 0 for MXFP8/NVFP4 (#3041) * [PyTorch Debug] Fix scale_inv_min always returning 0 for MXFP8/NVFP4 MXFP8/NVFP4 quantizers pad scale_inv to multiples of [128, 4] (or [4, 128] columnwise) with zeros, so a plain .min() over the whole tensor was always returning 0. Mask zeros out before computing the minimum. Fixes #2628 Signed-off-by: Pawel Gadzinski * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Clarify scale_inv padding comment The previous wording said the padding was always [128, 4] / [4, 128], which is true for MXFP8 but inaccurate for NVFP4 columnwise (padded to [128, 4], not [4, 128]). Also note that scale_inv is never naturally 0 (compute_scale_from_amax returns 1.0 for all-zero blocks), so masking zeros is exact rather than heuristic. Signed-off-by: Pawel Gadzinski --------- Signed-off-by: Pawel Gadzinski Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../debug/features/utils/stats_computation.py | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/transformer_engine/debug/features/utils/stats_computation.py b/transformer_engine/debug/features/utils/stats_computation.py index b0002ffee..666840001 100644 --- a/transformer_engine/debug/features/utils/stats_computation.py +++ b/transformer_engine/debug/features/utils/stats_computation.py @@ -348,6 +348,21 @@ def get_scale_inv(quantized_tensor, columnwise): return getattr(quantized_tensor, "_columnwise_scale_inv") return getattr(quantized_tensor, "_rowwise_scale_inv") + def nonzero_min(scale_inv): + # MXFP8/NVFP4 quantizers round the scale_inv shape up to multiples of + # 128 along one axis and 4 along the other and fill the extra slots + # with zeros (via torch.nn.functional.pad with the default value=0), + # so a plain .min() always returns 0 for shapes that needed padding. + # A real scale_inv entry is never 0: compute_scale_from_amax returns + # scale=1.0 for all-zero blocks and clamps the inf case to a finite + # fallback, so zeros uniquely identify padding and masking them out + # gives the true minimum. The empty-after-mask branch is a safety + # net for the (in practice unreachable) all-zero tensor. + nz = scale_inv[scale_inv != 0] + if nz.numel() == 0: + return scale_inv.new_zeros(()) + return nz.min() + columnwise_suffix = "_columnwise" if columnwise else "" # Prepare stat names. stat_name_min = ( @@ -363,7 +378,9 @@ def get_scale_inv(quantized_tensor, columnwise): # Capture the attribute name inside lambdas via default args to avoid late binding. STATS[stat_name_min] = ( - lambda x, aux_dict, _col=columnwise: get_scale_inv(aux_dict[recipe_name], _col).min(), + lambda x, aux_dict, _col=columnwise: nonzero_min( + get_scale_inv(aux_dict[recipe_name], _col) + ), lambda buffers, _sn=stat_name_min: min(_get(buffers, _sn)), ) STATS[stat_name_max] = ( From 920a7db1a03c854a93d8f9f51995e788d5c425e3 Mon Sep 17 00:00:00 2001 From: Siddhartha Raman Sundara Raman Date: Mon, 1 Jun 2026 14:50:55 -0500 Subject: [PATCH 095/180] Enable NVFP4 fused grouped MLP (#3048) * Enable NVFP4 fused grouped MLP follow-up Signed-off-by: sraman-rgb <270218152+sraman-rgb@users.noreply.github.com> * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * Address NVFP4 grouped MLP review comments Signed-off-by: Siddhartha Raman S * Update transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Siddhartha Raman Sundara Raman * Update transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Signed-off-by: Siddhartha Raman Sundara Raman * Address grouped MLP NVFP4 review feedback Signed-off-by: Siddhartha Raman S * Add NVFP4 RHT grouped MLP coverage Route the NVFP4 RHT grouped MLP test through the shared recipe helpers, compare the fused path against a TE unfused reference, and keep plain NVFP4 on the non-RHT fallback path. Signed-off-by: Siddhartha Raman S * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * [PyTorch] Drop `_MXFP8` suffix from fused grouped MLP op classes These fused ops now support both MXFP8 and NVFP4, so the recipe-specific suffix is misleading. Rename the four `CuTeGEMM` GLU/Unary forward/backward classes and update callsites and tests. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Tim Moon * [PyTorch] Add `ceil_div` utility and use it in fused grouped MLP ops Adds `ceil_div` next to `round_up_to_nearest_multiple` in `pytorch/utils.py` and replaces the `(x + d - 1) // d` patterns in the fused grouped MLP ops with it. Also fixes asymmetric floor-divs in the NVFP4 scale-view paths (`data_(in_)k // k_sf_divisor`) that would underestimate the scale-block count for padded layouts; the MXFP8 branches already used ceil-div for the same dimension. Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Tim Moon * [PyTorch] Drop defensive `getattr` defaults in fused grouped MLP ops Functions like `_group_quantize_for_grouped_mlp` already constrain their inputs (e.g. an NVFP4Quantizer always yields an NVFP4Tensor), so the `getattr(..., default)` fallbacks for `_rowwise_data`/`_with_gemm_swizzled_scales`/etc. were dead code that obscured intent. Replace those with direct attribute access, drop the dead double-`getattr` for the non-underscored public name that doesn't exist, and add brief comments on the type contract. The `getattr` calls that remain are legitimate (polymorphic inputs, dynamic attribute names, user-stamped optional flags on `torch.nn.Parameter`). Co-Authored-By: Claude Opus 4.7 (1M context) Signed-off-by: Tim Moon * Fix NVFP4 RHT grouped MLP reference Use explicit grouped linear quantizer roles and keep the NVFP4 RHT grouped MLP reference in plain PyTorch, with RHT applied only to reference wgrad. Signed-off-by: Siddhartha Raman S * Simplify NVFP4 RHT grouped MLP reference Signed-off-by: Siddhartha Raman S * Fix PyTorch ops lint Signed-off-by: Siddhartha Raman S --------- Signed-off-by: sraman-rgb <270218152+sraman-rgb@users.noreply.github.com> Signed-off-by: Siddhartha Raman S Signed-off-by: Siddhartha Raman Sundara Raman Signed-off-by: Tim Moon Co-authored-by: sraman-rgb <270218152+sraman-rgb@users.noreply.github.com> Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> Co-authored-by: Tim Moon Co-authored-by: Claude Opus 4.7 (1M context) --- tests/pytorch/test_fusible_ops.py | 190 ++++---- tests/pytorch/utils.py | 22 +- transformer_engine/pytorch/ops/_common.py | 149 +++++- .../pytorch/ops/basic/grouped_linear.py | 21 +- .../pytorch/ops/fused/__init__.py | 8 +- .../pytorch/ops/fused/backward_grouped_mlp.py | 456 +++++++++++++----- .../pytorch/ops/fused/forward_grouped_mlp.py | 414 +++++++++++----- transformer_engine/pytorch/utils.py | 9 +- 8 files changed, 909 insertions(+), 360 deletions(-) diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py index 1ced32e1a..df607bd6d 100644 --- a/tests/pytorch/test_fusible_ops.py +++ b/tests/pytorch/test_fusible_ops.py @@ -44,7 +44,6 @@ is_bf16_available, ) from transformer_engine.pytorch.tensor.grouped_tensor import GroupedTensor -from transformer_engine.pytorch.cpp_extensions.gemm import general_grouped_gemm_for_grouped_tensor import transformer_engine_torch as tex # Import utility functions @@ -80,6 +79,9 @@ if nvfp4_available: _quantization_list.append("nvfp4") _quantization_list.append("nvfp4_4over6") +_grouped_mlp_quantization_list = list(_quantization_list) +if nvfp4_available: + _grouped_mlp_quantization_list.append("nvfp4_rht") @pytest.fixture(autouse=True, scope="function") @@ -109,7 +111,10 @@ def maybe_skip_quantization( pytest.skip(reason_for_no_fp8) if quantization == "mxfp8" and not mxfp8_available: pytest.skip(reason_for_no_mxfp8) - if quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6") and not nvfp4_available: + if ( + quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht") + and not nvfp4_available + ): pytest.skip(reason_for_no_nvfp4) # Check dims @@ -122,14 +127,14 @@ def maybe_skip_quantization( elif quantization == "mxfp8": if math.prod(dims[:-1]) % 32 != 0 or dims[-1] % 32 != 0: pytest.skip("MXFP8 GEMMs require dims that are divisible by 32") - elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): + elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"): if math.prod(dims[:-1]) % 16 != 0 or dims[-1] % 16 != 0: pytest.skip("NVFP4 GEMMs require dims that are divisible by 16") # Check dtype if dtype is not None: if ( - quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6") + quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht") and dtype != torch.bfloat16 ): pytest.skip("NVFP4 quantization is only supported with BF16 data") @@ -187,10 +192,14 @@ def make_reference_and_test_tensors( test = quantizer(test) elif quantization == "mxfp8": test = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3)(test) - elif quantization in ("nvfp4", "nvfp4_row_scaled"): + elif quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_rht"): + tensor_type = "input" + if quantizer_role is not None: + tensor_type = quantizer_role.tensor_type + with_rht = quantization == "nvfp4_rht" and tensor_type != "weight" test = NVFP4Quantizer( - with_rht=False, - with_post_rht_amax=False, + with_rht=with_rht, + with_post_rht_amax=with_rht, with_2d_quantization=False, stochastic_rounding=False, with_random_sign_mask=False, @@ -3685,7 +3694,7 @@ def test_layernorm_mlp( @pytest.mark.parametrize("bias", (False, True)) @pytest.mark.parametrize("dtype", _dtypes) - @pytest.mark.parametrize("quantization", _quantization_list) + @pytest.mark.parametrize("quantization", _grouped_mlp_quantization_list) @pytest.mark.parametrize("single_grouped_weight", (False, True)) @pytest.mark.parametrize("single_grouped_bias", (False, True)) @pytest.mark.parametrize("accumulate_into_main_grad", (False, True)) @@ -3753,16 +3762,19 @@ def test_grouped_mlp( pytest.skip("Unary activations do not use GLU interleaving") if quantization == "nvfp4_4over6": pytest.skip("NVFP4 4over6 grouped quantization is not supported") + if quantization == "nvfp4_rht" and ( + activation != "scaled_swiglu" or bias or glu_interleave_size != 32 + ): + pytest.skip("NVFP4 RHT grouped MLP coverage is limited to fused no-bias SwiGLU") if ( with_quantization - and quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6") + and quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht") and activation.startswith("scaled_clamped_qgeglu") and bias ): # TODO: ksivaman: Need to debug numerics for this case. pytest.skip("Bias/dbias not yet supported in NVFP4 fused grouped MLP with GeGLU") fc1_out_features = 2 * hidden_size if activation_is_glu else hidden_size - # Activation parameters for clamped QGeGLU variants if activation == "scaled_clamped_qgeglu_custom": geglu_limit = 5.0 @@ -3845,13 +3857,7 @@ def test_grouped_mlp( fc2_ws_test.append(fc2_w_test) fc2_bs_test.append(fc2_b_test) - # Reference implementation - xs = torch.split(x_ref, split_sizes.tolist()) - probs = torch.split(probs_ref, split_sizes.tolist()) - ys = [] - for group_idx in range(group_size): - x = xs[group_idx] - x = torch.nn.functional.linear(x, fc1_ws_ref[group_idx], bias=fc1_bs_ref[group_idx]) + def _apply_activation(x: torch.Tensor) -> torch.Tensor: if activation_is_glu and glu_interleave_size is not None: x = x.reshape( -1, @@ -3863,66 +3869,85 @@ def test_grouped_mlp( x = x.reshape(-1, 2 * hidden_size) if activation == "scaled_swiglu": x1, x2 = x.chunk(2, dim=-1) - x = torch.nn.functional.silu(x1) * x2 - elif activation.startswith("scaled_clamped_qgeglu"): + return torch.nn.functional.silu(x1) * x2 + if activation.startswith("scaled_clamped_qgeglu"): x1, x2 = x.chunk(2, dim=-1) lim = torch.tensor(geglu_limit, device=x1.device, dtype=x1.dtype) x1c = torch.minimum(x1, lim) x2c = torch.clamp(x2, -lim, lim) - x = (x2c + geglu_offset) * (x1c * torch.sigmoid(geglu_alpha * x1c)) - elif activation == "scaled_srelu": - x = torch.nn.functional.relu(x).square() - else: - raise ValueError(f"Unexpected grouped MLP activation ({activation})") - x = x * probs[group_idx].unsqueeze(-1) - x = torch.nn.functional.linear(x, fc2_ws_ref[group_idx]) + return (x2c + geglu_offset) * (x1c * torch.sigmoid(geglu_alpha * x1c)) + if activation == "scaled_srelu": + return torch.nn.functional.relu(x).square() + raise ValueError(f"Unexpected grouped MLP activation ({activation})") + + # Reference implementation + xs = torch.split(x_ref, split_sizes.tolist()) + probs = torch.split(probs_ref, split_sizes.tolist()) + ys = [] + for group_idx in range(group_size): + x = xs[group_idx] + fc1_out = torch.nn.functional.linear( + x, fc1_ws_ref[group_idx], bias=fc1_bs_ref[group_idx] + ) + fc2_in = _apply_activation(fc1_out) + fc2_in = fc2_in * probs[group_idx].unsqueeze(-1) + y = torch.nn.functional.linear(fc2_in, fc2_ws_ref[group_idx]) if bias: - x = x + fc2_bs_ref[group_idx] * probs[group_idx].unsqueeze(-1) - ys.append(x) + y = y + fc2_bs_ref[group_idx] * probs[group_idx].unsqueeze(-1) + ys.append(y) y_ref = torch.cat(ys) y_ref.backward(dy_ref) # Construct operations recipe = make_recipe(quantization) - if activation == "scaled_clamped_qgeglu_custom": - scaled_act = te_ops.ScaledClampedQGeGLU( - glu_interleave_size=glu_interleave_size, - limit=geglu_limit, - alpha=geglu_alpha, - glu_linear_offset=geglu_offset, - ) - with te.quantized_model_init(enabled=with_quantization, recipe=recipe): - fc1 = te_ops.GroupedLinear( - group_size, - hidden_size, - fc1_out_features, - bias=bias, - device=device, - dtype=dtype, - single_grouped_weight=single_grouped_weight, - single_grouped_bias=single_grouped_bias, - accumulate_into_main_grad=accumulate_into_main_grad, - delay_wgrad_compute=delay_wgrad_compute, - ) - fc2 = te_ops.GroupedLinear( - group_size, - hidden_size, - hidden_size, - bias=bias, - device=device, - dtype=dtype, - single_grouped_weight=single_grouped_weight, - single_grouped_bias=single_grouped_bias, - accumulate_into_main_grad=accumulate_into_main_grad, - delay_wgrad_compute=delay_wgrad_compute, - scale_bias=bias, - ) - module = te_ops.Sequential( - fc1, - scaled_act, - fc2, - ) + def _make_scaled_act(): + if activation == "scaled_swiglu": + return te_ops.ScaledSwiGLU(glu_interleave_size=glu_interleave_size) + if activation == "scaled_clamped_qgeglu_custom": + return te_ops.ScaledClampedQGeGLU( + glu_interleave_size=glu_interleave_size, + limit=geglu_limit, + alpha=geglu_alpha, + glu_linear_offset=geglu_offset, + ) + if activation.startswith("scaled_clamped_qgeglu"): + return te_ops.ScaledClampedQGeGLU(glu_interleave_size=glu_interleave_size) + if activation == "scaled_srelu": + return te_ops.ScaledSReLU() + raise ValueError(f"Unexpected grouped MLP activation ({activation})") + + def _make_module(): + with te.quantized_model_init(enabled=with_quantization, recipe=recipe): + fc1_op = te_ops.GroupedLinear( + group_size, + hidden_size, + fc1_out_features, + bias=bias, + device=device, + dtype=dtype, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, + accumulate_into_main_grad=accumulate_into_main_grad, + delay_wgrad_compute=delay_wgrad_compute, + ) + + fc2_op = te_ops.GroupedLinear( + group_size, + hidden_size, + hidden_size, + bias=bias, + device=device, + dtype=dtype, + single_grouped_weight=single_grouped_weight, + single_grouped_bias=single_grouped_bias, + accumulate_into_main_grad=accumulate_into_main_grad, + delay_wgrad_compute=delay_wgrad_compute, + scale_bias=bias, + ) + return te_ops.Sequential(fc1_op, _make_scaled_act(), fc2_op), fc1_op, fc2_op + + module, fc1, fc2 = _make_module() # Copy weights with torch.no_grad(): @@ -3976,7 +4001,7 @@ def test_grouped_mlp( fc2.backward_dw() # Check for expected fusions - if ( + expected_grouped_mlp_fusion = ( quantization == "mxfp8" and dtype in (torch.bfloat16, torch.float16) and ( @@ -3984,13 +4009,14 @@ def test_grouped_mlp( or (activation_is_glu and glu_interleave_size == 32) ) and _cudnn_frontend_version_supported() - ): + ) + if expected_grouped_mlp_fusion: if activation_is_glu: - forward_cls = te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU_MXFP8 - backward_cls = te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8 + forward_cls = te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU + backward_cls = te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU else: - forward_cls = te_ops.fused.ForwardGroupedMLP_CuTeGEMMUnary_MXFP8 - backward_cls = te_ops.fused.BackwardGroupedMLP_CuTeGEMMDUnary_MXFP8 + forward_cls = te_ops.fused.ForwardGroupedMLP_CuTeGEMMUnary + backward_cls = te_ops.fused.BackwardGroupedMLP_CuTeGEMMDUnary if forward_cls.is_supported(): forward_ops = module._module_groups[0]._forward_ops assert len(forward_ops) == 1 @@ -4008,7 +4034,7 @@ def test_grouped_mlp( # Loose tols for sanity checking tols = {"rtol": 0.125, "atol": 0.25} - if quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): + if quantization in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"): tols = {"rtol": 0.25, "atol": 0.5} # Check values @@ -4088,9 +4114,9 @@ def test_grouped_mlp_single_weight_numerics( ) -> None: """single_grouped_weight=True/False should match exactly for fused MXFP8 grouped MLP.""" - if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU_MXFP8.is_supported(): + if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP forward is not supported on this system") - if not te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8.is_supported(): + if not te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP backward is not supported on this system") split_sizes = [split_alignment * (i + 1) for i in range(group_size)] @@ -4192,12 +4218,12 @@ def _run_case(single_grouped_weight: bool) -> tuple[torch.Tensor, ...]: assert len(forward_ops) == 1 assert isinstance( forward_ops[0][0], - te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU_MXFP8, + te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU, ) assert len(backward_ops) == 1 assert isinstance( backward_ops[0][0], - te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8, + te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU, ) if single_grouped_weight: @@ -4310,9 +4336,9 @@ def test_grouped_mlp_overwrite_main_grad( that read ``.grad`` don't see stale bytes from the cached dummy). """ - if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU_MXFP8.is_supported(): + if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP forward is not supported on this system") - if not te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8.is_supported(): + if not te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP backward is not supported on this system") recipe = make_recipe("mxfp8") @@ -4444,7 +4470,7 @@ def test_grouped_mlp_cuda_graph_safe_mxfp8( ) -> None: """Grouped MLP forward+backward should be CUDA graph capturable (MXFP8).""" - if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU_MXFP8.is_supported(): + if not te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU.is_supported(): pytest.skip("MXFP8 fused grouped MLP is not supported on this system") if dtype not in (torch.bfloat16, torch.float16): pytest.skip("MXFP8 fused grouped MLP is only supported with BF16/FP16") @@ -4586,12 +4612,12 @@ def train_step( assert len(forward_ops) == 1 assert isinstance( forward_ops[0][0], - te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU_MXFP8, + te_ops.fused.ForwardGroupedMLP_CuTeGEMMGLU, ) assert len(backward_ops) == 1 assert isinstance( backward_ops[0][0], - te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8, + te_ops.fused.BackwardGroupedMLP_CuTeGEMMDGLU, ) fresh_x = torch.randn_like(static_x) diff --git a/tests/pytorch/utils.py b/tests/pytorch/utils.py index 19cc118a9..84489f30c 100644 --- a/tests/pytorch/utils.py +++ b/tests/pytorch/utils.py @@ -118,7 +118,7 @@ def quantization_tols(name: str) -> dict[str, float]: "mxfp8_block_scaling", ): return dtype_tols(tex.DType.kFloat8E4M3) - if name in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): + if name in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"): return dtype_tols(tex.DType.kFloat4E2M1) raise ValueError(f"Unsupported quantization scheme ({name})") @@ -145,10 +145,10 @@ def make_recipe(name: Optional[str], **recipe_kwargs: Any) -> Optional[Recipe]: ) if name == "fp8_block_scaling": return transformer_engine.common.recipe.Float8BlockScaling(**recipe_kwargs) - if name in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6"): + if name in ("nvfp4", "nvfp4_row_scaled", "nvfp4_4over6", "nvfp4_rht"): use_4over6 = name == "nvfp4_4over6" kwargs = { - "disable_rht": True, + "disable_rht": name != "nvfp4_rht", "disable_stochastic_rounding": True, "disable_2d_quantization": not use_4over6, "row_scaled_activation": name == "nvfp4_row_scaled", @@ -163,12 +163,16 @@ def recipe_id(recipe: Optional[Recipe]) -> str: """Readable pytest id for a quantization recipe.""" if not isinstance(recipe, Recipe): return "None" - if recipe.nvfp4() and recipe.row_scaled_activation and recipe.nvfp4_4over6 != "none": - return "NVFP4RowScaled4Over6BlockScaling" - if recipe.nvfp4() and recipe.nvfp4_4over6 != "none": - return "NVFP44Over6BlockScaling" - if recipe.nvfp4() and recipe.row_scaled_activation: - return "NVFP4RowScaledBlockScaling" + if recipe.nvfp4(): + nvfp4_features = [] + if recipe.row_scaled_activation: + nvfp4_features.append("RowScaled") + if recipe.nvfp4_4over6 != "none": + nvfp4_features.append("4Over6") + if not recipe.disable_rht: + nvfp4_features.append("RHT") + if nvfp4_features: + return f"NVFP4{''.join(nvfp4_features)}BlockScaling" return type(recipe).__name__ diff --git a/transformer_engine/pytorch/ops/_common.py b/transformer_engine/pytorch/ops/_common.py index 717d87201..87911d76f 100644 --- a/transformer_engine/pytorch/ops/_common.py +++ b/transformer_engine/pytorch/ops/_common.py @@ -5,6 +5,7 @@ """Helper functions used in fusible operations.""" from __future__ import annotations +from collections.abc import Iterable import functools import math from importlib.metadata import PackageNotFoundError, version as get_pkg_version @@ -13,10 +14,13 @@ import torch from packaging.version import Version as PkgVersion +import transformer_engine_torch as tex from transformer_engine_torch import FP8TensorMeta from ..torch_version import torch_version from ..quantization import FP8GlobalStateManager +from ..tensor import NVFP4Quantizer, NVFP4Tensor, NVFP4TensorStorage, Quantizer from ..tensor.float8_tensor import Float8Tensor +from ..tensor.grouped_tensor import GroupedTensor from ..quantized_tensor import QuantizedTensorStorage from ..utils import canonicalize_dtype @@ -57,6 +61,146 @@ def _nvidia_cudnn_frontend_supports_wgrad() -> bool: return _cudnn_frontend_version_supported() +def _group_quantize_for_grouped_mlp( + tensor: torch.Tensor, + quantizer: Quantizer, + num_groups: int, + split_sizes: Optional[torch.Tensor], + *, + tensor_offsets: Optional[torch.Tensor] = None, +) -> GroupedTensor: + """Quantize into grouped storage.""" + + # Typical case: group-quantize + if num_groups != 1 or not isinstance(quantizer, NVFP4Quantizer): + return tex.group_quantize(tensor, quantizer, num_groups, split_sizes) + + # -------------------------------------------------- + # Special case: single-tensor NVFP4 quantize + # -------------------------------------------------- + + quantized = tex.quantize(tensor, quantizer) + with_gemm_swizzled_scales = quantized._with_gemm_swizzled_scales + if quantizer.optimize_for_gemm: + tex.swizzle_scales_for_gemm_(quantized) + with_gemm_swizzled_scales = True + + rowwise_data = quantized._rowwise_data + rowwise_scale = quantized._rowwise_scale_inv + columnwise_data = quantized._columnwise_data + columnwise_scale = quantized._columnwise_scale_inv + amax = quantized._amax_rowwise + columnwise_amax = quantized._amax_columnwise + + if split_sizes is None: + split_sizes = torch.full((1,), tensor.shape[0], dtype=torch.int64, device=tensor.device) + else: + split_sizes = split_sizes.to(dtype=torch.int64, device=tensor.device) + + m_dim = tensor.shape[0] + if rowwise_data is not None: + k_dim = rowwise_data.shape[-1] * 2 + elif columnwise_data is not None: + k_dim = columnwise_data.shape[0] + else: + k_dim = tensor.shape[-1] + + if tensor_offsets is None: + tensor_offsets = torch.cat( + [ + torch.zeros(1, dtype=torch.int64, device=tensor.device), + torch.cumsum(split_sizes * k_dim, dim=0), + ], + ) + + return GroupedTensor( + shape=(m_dim, k_dim), + dtype=tensor.dtype, + quantizer=quantizer, + num_tensors=1, + data=rowwise_data.reshape(-1) if rowwise_data is not None else None, + columnwise_data=columnwise_data.reshape(-1) if columnwise_data is not None else None, + scale_inv=rowwise_scale.reshape(-1) if rowwise_scale is not None else None, + columnwise_scale_inv=columnwise_scale.reshape(-1) if columnwise_scale is not None else None, + amax=amax, + columnwise_amax=columnwise_amax, + first_dims=split_sizes, + tensor_offsets=tensor_offsets, + with_gemm_swizzled_scales=with_gemm_swizzled_scales, + ) + + +def _nvfp4_amax( + tensors: GroupedTensor | Iterable[NVFP4TensorStorage], + *, + columnwise: bool, +) -> torch.Tensor: + """Get one NVFP4 amax value per group.""" + grouped_attr = "columnwise_amax" if columnwise else "amax" + tensor_attr = "_amax_columnwise" if columnwise else "_amax_rowwise" + + if hasattr(tensors, grouped_attr): + amax = getattr(tensors, grouped_attr) + if amax is None: + raise RuntimeError(f"NVFP4 GroupedTensor is missing {grouped_attr}.") + return amax.view(-1) + + amaxes = [getattr(tensor, tensor_attr) for tensor in tensors] + if any(amax is None for amax in amaxes): + raise RuntimeError(f"NVFP4 tensor list is missing {tensor_attr}.") + return torch.cat([amax.view(-1) for amax in amaxes], dim=0) + + +def _nvfp4_single_tensor_from_grouped( + grouped: GroupedTensor, + quantizer: Optional[NVFP4Quantizer] = None, + *, + fp4_dtype: Optional[torch.dtype] = None, +) -> NVFP4Tensor: + """Build a single NVFP4Tensor view over a one-member grouped storage.""" + if quantizer is None: + quantizer = grouped.quantizer + if not isinstance(quantizer, NVFP4Quantizer): + raise TypeError("Expected an NVFP4 GroupedTensor.") + + shape = tuple(grouped.logical_shape) + rowwise_data = None + if grouped.rowwise_data is not None: + rowwise_data = grouped.rowwise_data.view(quantizer.convert_shape_for_fp4(shape)) + + rowwise_scale_inv = None + if grouped.scale_inv is not None: + rowwise_scale_inv = grouped.scale_inv.view(quantizer.get_scale_shape(shape, False)) + + columnwise_data = None + if grouped.columnwise_data is not None: + columnwise_shape = quantizer.get_columnwise_shape(shape) + columnwise_data = grouped.columnwise_data.view( + quantizer.convert_shape_for_fp4(columnwise_shape) + ) + + columnwise_scale_inv = None + if grouped.columnwise_scale_inv is not None: + columnwise_scale_inv = grouped.columnwise_scale_inv.view( + quantizer.get_scale_shape(shape, True) + ) + + return NVFP4Tensor( + shape=shape, + dtype=grouped.get_dtype(), + rowwise_data=rowwise_data, + rowwise_scale_inv=rowwise_scale_inv, + columnwise_data=columnwise_data, + columnwise_scale_inv=columnwise_scale_inv, + amax_rowwise=grouped.amax, + amax_columnwise=grouped.columnwise_amax, + fp4_dtype=fp4_dtype or quantizer.dtype, + quantizer=quantizer, + requires_grad=False, + with_gemm_swizzled_scales=grouped._with_gemm_swizzled_scales, + ) + + def is_quantized_tensor(tensor: torch.Tensor | QuantizedTensorStorage) -> bool: """Check if tensor is a quantized tensor""" return isinstance(tensor, QuantizedTensorStorage) @@ -285,7 +429,10 @@ def fuse_grouped_mlp_ops( if not fused_op_cls.is_supported(): return ops - if recipe is None or not recipe.mxfp8(): + if recipe is None or not (recipe.mxfp8() or recipe.nvfp4()): + return ops + # NVFP4 fused grouped MLP uses graph-safe grouped quantize, which currently requires RHT. + if recipe.nvfp4() and recipe.disable_rht: return ops if activation_op_types is None: activation_op_types = (ScaledSwiGLU, ScaledClampedQGeGLU) diff --git a/transformer_engine/pytorch/ops/basic/grouped_linear.py b/transformer_engine/pytorch/ops/basic/grouped_linear.py index dc15bc63b..e9787f96b 100644 --- a/transformer_engine/pytorch/ops/basic/grouped_linear.py +++ b/transformer_engine/pytorch/ops/basic/grouped_linear.py @@ -22,7 +22,7 @@ _2X_ACC_DGRAD, _2X_ACC_WGRAD, ) -from ...quantization import FP8GlobalStateManager, Recipe +from ...quantization import FP8GlobalStateManager, QuantizerRole, Recipe from ...quantized_tensor import QuantizedTensorStorage from ...tensor import MXFP8Quantizer, MXFP8Tensor, Quantizer from ...utils import ( @@ -291,6 +291,25 @@ def num_quantizers(self, mode: str) -> int: return self.num_groups return 0 + def get_quantizer_roles(self, mode: str) -> Optional[list[QuantizerRole]]: + name = getattr(self, "name", "") or "" + if mode == "forward": + roles = [] + for _ in range(self.num_groups): + roles.extend( + [ + QuantizerRole(module_type="linear", tensor_type="input", name=name), + QuantizerRole(module_type="linear", tensor_type="weight", name=name), + ] + ) + return roles + if mode == "backward": + return [ + QuantizerRole(module_type="linear", tensor_type="grad_output", name=name) + for _ in range(self.num_groups) + ] + return None + @property def has_bias(self) -> bool: """Whether an additive bias is being applied""" diff --git a/transformer_engine/pytorch/ops/fused/__init__.py b/transformer_engine/pytorch/ops/fused/__init__.py index b29e35814..78f9d880b 100644 --- a/transformer_engine/pytorch/ops/fused/__init__.py +++ b/transformer_engine/pytorch/ops/fused/__init__.py @@ -32,10 +32,10 @@ # Import experimental fusions # Note: Registration logic is non-trivial, so submodule handles it internally. from .forward_grouped_mlp import ( # pylint: disable=wrong-import-position - ForwardGroupedMLP_CuTeGEMMGLU_MXFP8, - ForwardGroupedMLP_CuTeGEMMUnary_MXFP8, + ForwardGroupedMLP_CuTeGEMMGLU, + ForwardGroupedMLP_CuTeGEMMUnary, ) from .backward_grouped_mlp import ( # pylint: disable=wrong-import-position - BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8, - BackwardGroupedMLP_CuTeGEMMDUnary_MXFP8, + BackwardGroupedMLP_CuTeGEMMDGLU, + BackwardGroupedMLP_CuTeGEMMDUnary, ) diff --git a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py index 25ccad137..792b6d781 100644 --- a/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/backward_grouped_mlp.py @@ -14,10 +14,17 @@ import transformer_engine_torch as tex from ...quantization import Recipe +from ...tensor import NVFP4Quantizer, NVFP4Tensor from ...tensor.grouped_tensor import GroupedTensor from ...tensor.mxfp8_tensor import MXFP8Quantizer -from ...utils import clear_tensor_data, get_cached_ones_tensor, get_device_compute_capability -from ...constants import MXFP8_BLOCK_SCALING_SIZE +from ...utils import ( + ceil_div, + clear_tensor_data, + get_cached_ones_tensor, + get_device_compute_capability, + round_up_to_nearest_multiple, +) +from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE from ..basic import GroupedLinear, ScaledSReLU, ScaledClampedQGeGLU from ..fuser import register_backward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext @@ -25,6 +32,9 @@ _cudnn_frontend_geglu_runtime_params, _cudnn_frontend_version_supported, _cudnn_frontend_supports_grouped_gemm_srelu, + _group_quantize_for_grouped_mlp, + _nvfp4_amax, + _nvfp4_single_tensor_from_grouped, fuse_grouped_mlp_ops, get_accumulate_flag_in_param, get_dummy_wgrads_for_params, @@ -34,11 +44,41 @@ view_main_grad_as_grouped_buffer, validate_grouped_mlp_dims, ) -from ...cpp_extensions import general_grouped_gemm_for_grouped_tensor +from ...cpp_extensions import ( + general_gemm, + general_grouped_gemm_for_grouped_tensor, +) from ...module.base import _2X_ACC_WGRAD from ...triton.grouped_dbias_dscales import compute_grouped_dbias_dscales +def _nvfp4_single_group_wgrad_gemm( + grouped_x: GroupedTensor, + grouped_dy: GroupedTensor, + wgrad_output, + *, + weight_shape: tuple[int, int], + accumulate: bool, +) -> None: + """Run one-group NVFP4 wgrad with regular GEMM instead of grouped GEMM.""" + x_single = _nvfp4_single_tensor_from_grouped(grouped_x) + dy_single = _nvfp4_single_tensor_from_grouped(grouped_dy) + if isinstance(wgrad_output, GroupedTensor): + out = wgrad_output.rowwise_data.view(1, *weight_shape)[0] + else: + out = wgrad_output[0] + + general_gemm( + x_single, + dy_single, + out_dtype=out.dtype, + out=out, + layout="NT", + accumulate=accumulate, + use_split_accumulator=_2X_ACC_WGRAD, + ) + + def _cudnn_compute_wgrad( grouped_x: GroupedTensor, grouped_dy: GroupedTensor, @@ -62,8 +102,8 @@ def _cudnn_compute_wgrad( fp8_dtype = torch.float8_e4m3fn - sfa_leading_dim = ((out_features + 127) // 128) * 128 - sfb_leading_dim = ((in_features + 127) // 128) * 128 + sfa_leading_dim = round_up_to_nearest_multiple(out_features, 128) + sfb_leading_dim = round_up_to_nearest_multiple(in_features, 128) if total_tokens == 0: # A workaround for the case with zero-token experts. @@ -220,6 +260,18 @@ def _compute_grad_params( single_grouped_weight=fc_op.single_grouped_weight, current_stream=torch.cuda.current_stream().cuda_stream, ) + elif ( + num_groups == 1 + and isinstance(grouped_x, GroupedTensor) + and isinstance(grouped_dy, GroupedTensor) + and isinstance(grouped_x.quantizer, NVFP4Quantizer) + and isinstance(grouped_dy.quantizer, NVFP4Quantizer) + ): + gemm_fn = functools.partial( + _nvfp4_single_group_wgrad_gemm, + weight_shape=weight_shape, + accumulate=accumulate_into_main_grad, + ) else: gemm_fn = functools.partial( general_grouped_gemm_for_grouped_tensor, @@ -252,8 +304,8 @@ def _compute_grad_params( return w_list + bias_list -class _BackwardGroupedMLP_CuTeGEMMDBase_MXFP8(FusedOperation): - """Base fused backward op for MXFP8 GroupedLinear + activation + GroupedLinear. +class _BackwardGroupedMLP_CuTeGEMMDBase(FusedOperation): + """Base fused backward op for block-scaled GroupedLinear + activation + GroupedLinear. Uses experimental CuTe DSL kernel from cuDNN front-end. @@ -360,7 +412,9 @@ def fuser_backward( grad_output = grad_output.reshape(-1, fc2_weight_shape[0]) out_shape = list(grad_output.size()) num_groups = fc1_op.num_groups - device = fc1_op._get_weight_tensors()[0].device + fc1_weight_param = fc1_op.weight if fc1_op.single_grouped_weight else fc1_op.weight0 + fc2_weight_param = fc2_op.weight if fc2_op.single_grouped_weight else fc2_op.weight0 + device = fc1_weight_param.device dtype = fc1_ctx.dtype # Saved tensors from FC1 forward. @@ -419,10 +473,18 @@ def fuser_backward( output_fc2_dbias = fc2_op.has_bias fc2_dbias_packed = None fc2_dy = None + grad_output_quantizer = getattr(grad_output, "quantizer", None) + fc2_grad_output_quantizer_matches = ( + isinstance(fc2_grad_output_quantizer, MXFP8Quantizer) + and isinstance(grad_output_quantizer, MXFP8Quantizer) + ) or ( + isinstance(fc2_grad_output_quantizer, NVFP4Quantizer) + and isinstance(grad_output_quantizer, NVFP4Quantizer) + ) if ( not output_fc2_dbias and isinstance(grad_output, GroupedTensor) - and isinstance(getattr(grad_output, "quantizer", None), MXFP8Quantizer) + and fc2_grad_output_quantizer_matches ): grouped_fc2_dy = grad_output else: @@ -435,13 +497,26 @@ def fuser_backward( split_sizes, ) else: - grouped_fc2_dy = tex.group_quantize( + grouped_fc2_dy = _group_quantize_for_grouped_mlp( fc2_dy, fc2_grad_output_quantizer, num_groups, split_sizes, + tensor_offsets=base_split_offsets * fc2_weight_shape[0], ) + use_nvfp4 = ( + isinstance(fc2_grad_output_quantizer, NVFP4Quantizer) + or isinstance(fc1_weight_param, NVFP4Tensor) + or isinstance(fc2_weight_param, NVFP4Tensor) + ) + data_dtype = torch.float4_e2m1fn_x2 if use_nvfp4 else torch.float8_e4m3fn + scale_view_dtype = torch.float8_e4m3fn if use_nvfp4 else torch.float8_e8m0fnu + sf_vec_size = NVFP4_BLOCK_SCALING_SIZE if use_nvfp4 else MXFP8_BLOCK_SCALING_SIZE + data_k = out_shape[1] // 2 if use_nvfp4 else out_shape[1] + fc2_weight_k = fc2_weight_shape[1] // 2 if use_nvfp4 else fc2_weight_shape[1] + k_sf_divisor = 2 * sf_vec_size if use_nvfp4 else 4 * sf_vec_size + # Pack data tensors # Note: Fused kernel expects tensor with non-contiguous # logical dims. @@ -451,20 +526,42 @@ def fuser_backward( # Data logical shape: (sum(m), k, 1) # Scale logical shape: (32 (block row), 4 (block row), # sum(m)/128, 4 (block col), k/128, 1) - fc2_dy_data = grouped_fc2_dy.rowwise_data.view(out_shape[0], out_shape[1]) - fc2_dy_data = fc2_dy_data.view(dtype=torch.float8_e4m3fn) + fc2_dy_data = grouped_fc2_dy.rowwise_data.view(dtype=data_dtype) + fc2_dy_data = fc2_dy_data.view(out_shape[0], data_k) fc2_dy_data = fc2_dy_data.unsqueeze(0).permute(1, 2, 0) fc2_dy_scales = grouped_fc2_dy.scale_inv - fc2_dy_scales = fc2_dy_scales.view(dtype=torch.float8_e8m0fnu) - fc2_dy_scales = fc2_dy_scales.view( - 1, - (out_shape[0] + 127) // 128, - (out_shape[1] + 127) // 128, - MXFP8_BLOCK_SCALING_SIZE, - 4, - 4, - ) - fc2_dy_scales = fc2_dy_scales.permute(3, 4, 1, 5, 2, 0) + fc2_dy_scales = fc2_dy_scales.view(dtype=scale_view_dtype) + with_gemm_swizzled_scales = grouped_fc2_dy._with_gemm_swizzled_scales + if use_nvfp4 and with_gemm_swizzled_scales: + fc2_dy_scales = fc2_dy_scales.view( + 1, + ceil_div(out_shape[0], 128), + ceil_div(data_k, k_sf_divisor), + 32, + 4, + 4, + ) + fc2_dy_scales = fc2_dy_scales.permute(3, 4, 1, 5, 2, 0) + elif use_nvfp4: + fc2_dy_scales = fc2_dy_scales.view( + 1, + ceil_div(out_shape[0], 128), + 4, + 32, + ceil_div(data_k, k_sf_divisor), + 4, + ) + fc2_dy_scales = fc2_dy_scales.permute(3, 2, 1, 5, 4, 0) + else: + fc2_dy_scales = fc2_dy_scales.view( + 1, + ceil_div(out_shape[0], 128), + ceil_div(out_shape[1], k_sf_divisor), + 32, + 4, + 4, + ) + fc2_dy_scales = fc2_dy_scales.permute(3, 4, 1, 5, 2, 0) # Kernel scaling factors alpha_tensor = get_cached_ones_tensor(num_groups, dtype, device) @@ -475,25 +572,43 @@ def fuser_backward( scales_tensor = scales_f32.reshape(-1, 1, 1) dscales_tensor = torch.zeros_like(scales_tensor) + fc2_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn + if use_nvfp4: + nvfp4_fp4_max = 6.0 + nvfp4_fp8_max = 448.0 + fc2_alpha_tensor = ( + torch.sqrt( + _nvfp4_amax(grouped_fc2_dy, columnwise=False) + * _nvfp4_amax(grouped_fc2_weight, columnwise=True) + ) + / (nvfp4_fp8_max * nvfp4_fp4_max) + ).expand(num_groups) + fc2_beta_tensor = get_cached_ones_tensor(num_groups, torch.float32, device) + fc2_norm_const_tensor = None + else: + fc2_alpha_tensor = alpha_tensor + fc2_beta_tensor = alpha_tensor + fc2_norm_const_tensor = norm_const_tensor + fc2_dactivation_kwargs = { "a_tensor": fc2_dy_data, "c_tensor": activation_in.unsqueeze(0).permute(1, 2, 0), "sfa_tensor": fc2_dy_scales, "padded_offsets": split_points, - "alpha_tensor": alpha_tensor, + "alpha_tensor": fc2_alpha_tensor, "prob_tensor": scales_tensor, "dprob_tensor": dscales_tensor, "generate_dbias": fc1_op.has_bias, - "norm_const_tensor": norm_const_tensor, - "d_dtype": torch.float8_e4m3fn, + "norm_const_tensor": fc2_norm_const_tensor, + "d_dtype": fc2_d_dtype, "cd_major": "n", - "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, + "sf_vec_size": sf_vec_size, "current_stream": current_stream, - "discrete_col_sfd": True, + "discrete_col_sfd": not use_nvfp4, "use_dynamic_sched": True, } if self._cudnn_dact_func is not None: - fc2_dactivation_kwargs["beta_tensor"] = alpha_tensor + fc2_dactivation_kwargs["beta_tensor"] = fc2_beta_tensor fc2_dactivation_kwargs["act_func"] = self._cudnn_dact_func else: fc2_dactivation_kwargs["use_dsrelu_reuse"] = recompute_fc2_x_from_dsrelu @@ -513,19 +628,23 @@ def fuser_backward( # Data actual shape: (num_groups, k, n) # Data logical shape: (n, k, num_groups) fc2_w_data = fc2_weight_for_gemm.columnwise_data - fc2_w_data = fc2_w_data.view(dtype=torch.float8_e4m3fn) - fc2_w_data = fc2_w_data.view(num_groups, fc2_weight_shape[0], fc2_weight_shape[1]) - fc2_w_data = fc2_w_data.permute(2, 1, 0) - fc2_w_scales = fc2_weight_for_gemm.columnwise_scale_inv.view(dtype=torch.float8_e8m0fnu) + fc2_w_data = fc2_w_data.view(dtype=data_dtype) + fc2_w_data = fc2_w_data.view(num_groups, fc2_weight_shape[0], fc2_weight_k) + fc2_w_data = fc2_w_data.permute(1, 2, 0) if use_nvfp4 else fc2_w_data.permute(2, 1, 0) + fc2_w_scales = fc2_weight_for_gemm.columnwise_scale_inv.view(dtype=scale_view_dtype) fc2_w_scales = fc2_w_scales.view( num_groups, - (fc2_weight_shape[1] + 127) // 128, - (fc2_weight_shape[0] + 127) // 128, - MXFP8_BLOCK_SCALING_SIZE, + ceil_div(fc2_weight_shape[1], k_sf_divisor), + ceil_div(fc2_weight_shape[0], 128), + 32, 4, 4, ) - fc2_w_scales = fc2_w_scales.permute(3, 4, 1, 5, 2, 0) + fc2_w_scales = ( + fc2_w_scales.permute(3, 4, 2, 5, 1, 0) + if use_nvfp4 + else fc2_w_scales.permute(3, 4, 1, 5, 2, 0) + ) fc2_dactivation_kwargs["b_tensor"] = fc2_w_data fc2_dactivation_kwargs["sfb_tensor"] = fc2_w_scales @@ -534,27 +653,43 @@ def fuser_backward( [w._columnwise_data for w in grouped_fc2_weight], device, ) + swizzle_type = ( + "uniform_nvfp4_swizzle" if use_nvfp4 else "uniform_mxfp8_columnwise_swizzle" + ) fc2_sfb_ptrs, _fc2_sfb_buffer = tex.transform_and_copy_data_ptrs_to_device( - "uniform_mxfp8_columnwise_swizzle", + swizzle_type, [w._columnwise_scale_inv for w in grouped_fc2_weight], device, ) fc2_dactivation_kwargs["b_ptrs"] = fc2_b_ptrs fc2_dactivation_kwargs["sfb_ptrs"] = fc2_sfb_ptrs fc2_dactivation_kwargs["n"] = fc2_weight_shape[1] - fc2_dactivation_kwargs["b_dtype"] = torch.float8_e4m3fn - fc2_dactivation_kwargs["b_major"] = "n" + fc2_dactivation_kwargs["b_dtype"] = data_dtype + fc2_dactivation_kwargs["b_major"] = "k" if use_nvfp4 else "n" fc2_dgrad_kernel_out = self.grouped_gemm_dactivation_kernel()(**fc2_dactivation_kwargs) - fc1_dy_row_data = fc2_dgrad_kernel_out["d_row_tensor"] - fc1_dy_row_data = fc1_dy_row_data.view(out_shape[0], fc1_weight_shape[0]) - # View scale in their actual swizzled shape - fc1_dy_row_scale = fc2_dgrad_kernel_out["sfd_row_tensor"].permute(5, 2, 4, 0, 1, 3).view(-1) - fc1_dy_col_data = fc2_dgrad_kernel_out["d_col_tensor"] - fc1_dy_col_data = fc1_dy_col_data.view(out_shape[0], fc1_weight_shape[0]) - # View scale in their actual swizzled shape - fc1_dy_col_scale = fc2_dgrad_kernel_out["sfd_col_tensor"].permute(5, 2, 4, 0, 1, 3).view(-1) + if use_nvfp4: + fc1_dy_bf16 = fc2_dgrad_kernel_out["d_row_tensor"] + fc1_dy_bf16 = fc1_dy_bf16.view(out_shape[0], fc1_weight_shape[0]).contiguous() + fc1_dy_row_data = None + fc1_dy_row_scale = None + fc1_dy_col_data = None + fc1_dy_col_scale = None + else: + fc1_dy_bf16 = None + fc1_dy_row_data = fc2_dgrad_kernel_out["d_row_tensor"] + fc1_dy_row_data = fc1_dy_row_data.view(out_shape[0], fc1_weight_shape[0]) + # View scale in their actual swizzled shape + fc1_dy_row_scale = ( + fc2_dgrad_kernel_out["sfd_row_tensor"].permute(5, 2, 4, 0, 1, 3).view(-1) + ) + fc1_dy_col_data = fc2_dgrad_kernel_out["d_col_tensor"] + fc1_dy_col_data = fc1_dy_col_data.view(out_shape[0], fc1_weight_shape[0]) + # View scale in their actual swizzled shape + fc1_dy_col_scale = ( + fc2_dgrad_kernel_out["sfd_col_tensor"].permute(5, 2, 4, 0, 1, 3).view(-1) + ) grad_scales = fc2_dgrad_kernel_out["dprob_tensor"].view(-1) if recompute_fc2_x_from_dsrelu: @@ -628,21 +763,37 @@ def fuser_backward( # FC1 grad output for dgrad and wgrad GEMMs fc1_dy_tensor_offsets = base_split_offsets * fc1_weight_shape[0] - grouped_fc1_dy = GroupedTensor( - shape=(out_shape[0], fc1_weight_shape[0]), - dtype=dtype, - num_tensors=num_groups, - quantizer=fc1_ctx.grad_output_quantizers[0], - data=fc1_dy_row_data, - columnwise_data=fc1_dy_col_data, - scale_inv=fc1_dy_row_scale, - columnwise_scale_inv=fc1_dy_col_scale, - first_dims=split_sizes, - tensor_offsets=fc1_dy_tensor_offsets, - with_gemm_swizzled_scales=True, - ) + fc1_grad_output_quantizer = fc1_ctx.grad_output_quantizers[0] + if use_nvfp4: + fc1_grad_output_quantizer.set_usage( + rowwise=True, + columnwise=fc1_ctx.weight_requires_grad, + ) + fc1_grad_output_quantizer.optimize_for_gemm = True + grouped_fc1_dy = _group_quantize_for_grouped_mlp( + fc1_dy_bf16, + fc1_grad_output_quantizer, + num_groups, + split_sizes, + tensor_offsets=fc1_dy_tensor_offsets, + ) + else: + grouped_fc1_dy = GroupedTensor( + shape=(out_shape[0], fc1_weight_shape[0]), + dtype=dtype, + num_tensors=num_groups, + quantizer=fc1_grad_output_quantizer, + data=fc1_dy_row_data, + columnwise_data=fc1_dy_col_data, + scale_inv=fc1_dy_row_scale, + columnwise_scale_inv=fc1_dy_col_scale, + first_dims=split_sizes, + tensor_offsets=fc1_dy_tensor_offsets, + with_gemm_swizzled_scales=True, + ) # FC2 wgrad GEMM + wgrad_kernel_fn = None if use_nvfp4 else self.grouped_gemm_wgrad_kernel() fc2_grad_params = _compute_grad_params( fc_op=fc2_op, ctx=fc2_ctx, @@ -655,7 +806,7 @@ def fuser_backward( bias_grads=fc2_bias_grads, bias_grad_packed=fc2_bias_grad_packed, label="FC2", - cudnn_wgrad_kernel_fn=self.grouped_gemm_wgrad_kernel(), + cudnn_wgrad_kernel_fn=wgrad_kernel_fn, offsets=split_points, ) @@ -677,67 +828,110 @@ def fuser_backward( if fc1_ctx.input_requires_grad: in_shape = out_shape[:-1] + [fc1_weight_shape[1]] - fc1_dgrad_a_data = fc2_dgrad_kernel_out["d_row_tensor"] - fc1_dgrad_a_scales = fc2_dgrad_kernel_out["sfd_row_tensor"] - - fc1_dgrad_kwargs = { - "a_tensor": fc1_dgrad_a_data, - "sfa_tensor": fc1_dgrad_a_scales, - "padded_offsets": split_points, - "alpha_tensor": alpha_tensor, - "norm_const_tensor": None, - "prob_tensor": torch.ones((out_shape[0], 1, 1), dtype=torch.float32, device=device), - "acc_dtype": torch.float32, - "d_dtype": dtype, - "cd_major": "n", - "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, - "current_stream": current_stream, - "discrete_col_sfd": True, - "use_dynamic_sched": True, - } - - if fc1_op.single_grouped_weight: - # Clone and swizzle scales for GEMM - fc1_weight_for_gemm = grouped_fc1_weight.copy() - tex.grouped_swizzle_for_gemm(fc1_weight_for_gemm, rowwise=False, columnwise=True) - - fc1_w_data = fc1_weight_for_gemm.columnwise_data - fc1_w_data = fc1_w_data.view(dtype=torch.float8_e4m3fn) - fc1_w_data = fc1_w_data.view(num_groups, fc1_weight_shape[0], fc1_weight_shape[1]) - fc1_w_data = fc1_w_data.permute(2, 1, 0) - fc1_w_scales = fc1_weight_for_gemm.columnwise_scale_inv.view( - dtype=torch.float8_e8m0fnu - ) - fc1_w_scales = fc1_w_scales.view( - num_groups, - (fc1_weight_shape[1] + 127) // 128, - (fc1_weight_shape[0] + 127) // 128, - MXFP8_BLOCK_SCALING_SIZE, - 4, - 4, - ) - fc1_w_scales = fc1_w_scales.permute(3, 4, 1, 5, 2, 0) - - fc1_dgrad_kwargs["b_tensor"] = fc1_w_data - fc1_dgrad_kwargs["sfb_tensor"] = fc1_w_scales + if use_nvfp4: + grad_input = torch.empty(in_shape, dtype=dtype, device=device) + if num_groups == 1: + if fc1_op.single_grouped_weight: + fc1_w_single = grouped_fc1_weight.split_into_quantized_tensors()[0] + else: + fc1_w_single = grouped_fc1_weight[0] + fc1_dy_single = _nvfp4_single_tensor_from_grouped(grouped_fc1_dy) + general_gemm( + fc1_w_single, + fc1_dy_single, + out_dtype=dtype, + out=grad_input, + layout="NN", + ) + else: + fc1_x_tensor_offsets = base_split_offsets * fc1_weight_shape[1] + grouped_grad_input = GroupedTensor( + shape=(out_shape[0], fc1_weight_shape[1]), + dtype=dtype, + num_tensors=num_groups, + quantizer=None, + data=grad_input.view(-1), + first_dims=split_sizes, + tensor_offsets=fc1_x_tensor_offsets, + ) + general_grouped_gemm_for_grouped_tensor( + grouped_fc1_weight, + grouped_fc1_dy, + grouped_grad_input, + layout="NN", + ) else: - fc1_b_ptrs = tex.copy_data_ptrs_to_device( - [w._columnwise_data for w in grouped_fc1_weight], - device, - ) - fc1_sfb_ptrs, _fc1_sfb_buffer = tex.transform_and_copy_data_ptrs_to_device( - "uniform_mxfp8_columnwise_swizzle", - [w._columnwise_scale_inv for w in grouped_fc1_weight], - device, - ) - fc1_dgrad_kwargs["b_ptrs"] = fc1_b_ptrs - fc1_dgrad_kwargs["sfb_ptrs"] = fc1_sfb_ptrs - fc1_dgrad_kwargs["n"] = fc1_weight_shape[1] - fc1_dgrad_kwargs["b_dtype"] = torch.float8_e4m3fn - fc1_dgrad_kwargs["b_major"] = "n" - - fc1_dgrad_kernel_out = self.grouped_gemm_quant_kernel()(**fc1_dgrad_kwargs) - grad_input = fc1_dgrad_kernel_out["d_tensor"].view(in_shape) + fc1_dgrad_a_data = fc2_dgrad_kernel_out["d_row_tensor"] + fc1_dgrad_a_scales = fc2_dgrad_kernel_out["sfd_row_tensor"] + + fc1_dgrad_kwargs = { + "a_tensor": fc1_dgrad_a_data, + "sfa_tensor": fc1_dgrad_a_scales, + "padded_offsets": split_points, + "alpha_tensor": alpha_tensor, + "norm_const_tensor": None, + "prob_tensor": torch.ones( + (out_shape[0], 1, 1), dtype=torch.float32, device=device + ), + "acc_dtype": torch.float32, + "d_dtype": dtype, + "cd_major": "n", + "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, + "current_stream": current_stream, + "discrete_col_sfd": True, + "use_dynamic_sched": True, + } + + if fc1_op.single_grouped_weight: + # Clone and swizzle scales for GEMM + fc1_weight_for_gemm = grouped_fc1_weight.copy() + tex.grouped_swizzle_for_gemm( + fc1_weight_for_gemm, rowwise=False, columnwise=True + ) + + fc1_w_data = fc1_weight_for_gemm.columnwise_data + fc1_w_data = fc1_w_data.view(dtype=torch.float8_e4m3fn) + fc1_w_data = fc1_w_data.view( + num_groups, fc1_weight_shape[0], fc1_weight_shape[1] + ) + fc1_w_data = fc1_w_data.permute(2, 1, 0) + fc1_w_scales = fc1_weight_for_gemm.columnwise_scale_inv.view( + dtype=torch.float8_e8m0fnu + ) + fc1_w_scales = fc1_w_scales.view( + num_groups, + ceil_div(fc1_weight_shape[1], 128), + ceil_div(fc1_weight_shape[0], 128), + MXFP8_BLOCK_SCALING_SIZE, + 4, + 4, + ) + fc1_w_scales = fc1_w_scales.permute(3, 4, 1, 5, 2, 0) + + fc1_dgrad_kwargs["b_tensor"] = fc1_w_data + fc1_dgrad_kwargs["sfb_tensor"] = fc1_w_scales + else: + fc1_b_ptrs = tex.copy_data_ptrs_to_device( + [w._columnwise_data for w in grouped_fc1_weight], + device, + ) + swizzle_type = ( + "uniform_nvfp4_swizzle" if use_nvfp4 else "uniform_mxfp8_columnwise_swizzle" + ) + fc1_sfb_ptrs, _fc1_sfb_buffer = tex.transform_and_copy_data_ptrs_to_device( + swizzle_type, + [w._columnwise_scale_inv for w in grouped_fc1_weight], + device, + ) + + fc1_dgrad_kwargs["b_ptrs"] = fc1_b_ptrs + fc1_dgrad_kwargs["sfb_ptrs"] = fc1_sfb_ptrs + fc1_dgrad_kwargs["n"] = fc1_weight_shape[1] + fc1_dgrad_kwargs["b_dtype"] = torch.float8_e4m3fn + fc1_dgrad_kwargs["b_major"] = "n" + + fc1_dgrad_kernel_out = self.grouped_gemm_quant_kernel()(**fc1_dgrad_kwargs) + grad_input = fc1_dgrad_kernel_out["d_tensor"].view(in_shape) # FC1 wgrad GEMM fc1_grad_params = _compute_grad_params( @@ -752,7 +946,7 @@ def fuser_backward( bias_grads=fc1_bias_grads, bias_grad_packed=fc1_bias_grad_packed, label="FC1", - cudnn_wgrad_kernel_fn=self.grouped_gemm_wgrad_kernel(), + cudnn_wgrad_kernel_fn=wgrad_kernel_fn, offsets=split_points, ) @@ -778,8 +972,8 @@ def fuser_backward( ) -class BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8(_BackwardGroupedMLP_CuTeGEMMDBase_MXFP8): - """Fused backward op for GroupedLinear + scaled GLU + GroupedLinear.""" +class BackwardGroupedMLP_CuTeGEMMDGLU(_BackwardGroupedMLP_CuTeGEMMDBase): + """Fused backward op for block-scaled GroupedLinear + scaled GLU + GroupedLinear.""" @classmethod @functools.lru_cache(maxsize=None) @@ -790,8 +984,8 @@ def grouped_gemm_dactivation_kernel(cls) -> Callable: return grouped_gemm_dglu_wrapper_sm100 -class BackwardGroupedMLP_CuTeGEMMDUnary_MXFP8(_BackwardGroupedMLP_CuTeGEMMDBase_MXFP8): - """Fused backward op for GroupedLinear + scaled unary activation + GroupedLinear.""" +class BackwardGroupedMLP_CuTeGEMMDUnary(_BackwardGroupedMLP_CuTeGEMMDBase): + """Fused backward op for block-scaled GroupedLinear + scaled unary activation + GroupedLinear.""" @classmethod @functools.lru_cache(maxsize=None) @@ -833,7 +1027,7 @@ def fuse_backward_ops( return fuse_grouped_mlp_ops( ops, recipe=recipe, - fused_op_cls=BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8, + fused_op_cls=BackwardGroupedMLP_CuTeGEMMDGLU, ) @@ -845,16 +1039,18 @@ def fuse_backward_srelu_ops( ) -> list[FusibleOperation]: """Apply GroupedLinear + ScaledSReLU + GroupedLinear fusion for backward pass.""" + if recipe is None or not recipe.mxfp8(): + return ops return fuse_grouped_mlp_ops( ops, recipe=recipe, - fused_op_cls=BackwardGroupedMLP_CuTeGEMMDUnary_MXFP8, + fused_op_cls=BackwardGroupedMLP_CuTeGEMMDUnary, activation_op_types=(ScaledSReLU,), ) # Register fusion if available -if BackwardGroupedMLP_CuTeGEMMDGLU_MXFP8.is_supported(): +if BackwardGroupedMLP_CuTeGEMMDGLU.is_supported(): register_backward_fusion(fuse_backward_ops, prepend=True) -if BackwardGroupedMLP_CuTeGEMMDUnary_MXFP8.is_supported(): +if BackwardGroupedMLP_CuTeGEMMDUnary.is_supported(): register_backward_fusion(fuse_backward_srelu_ops, prepend=True) diff --git a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py index a0c5f766c..f4f210857 100644 --- a/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py +++ b/transformer_engine/pytorch/ops/fused/forward_grouped_mlp.py @@ -13,12 +13,18 @@ import torch import transformer_engine_torch as tex +from ...cpp_extensions import general_gemm, general_grouped_gemm_for_grouped_tensor from ...quantization import Recipe -from ...tensor import Quantizer -from ...utils import get_cached_ones_tensor, get_device_compute_capability, mark_grouped_tensor +from ...tensor import NVFP4Quantizer, NVFP4Tensor, Quantizer +from ...utils import ( + ceil_div, + get_cached_ones_tensor, + get_device_compute_capability, + mark_grouped_tensor, +) from ...tensor.grouped_tensor import GroupedTensor from ...tensor.mxfp8_tensor import MXFP8Quantizer -from ...constants import MXFP8_BLOCK_SCALING_SIZE +from ...constants import MXFP8_BLOCK_SCALING_SIZE, NVFP4_BLOCK_SCALING_SIZE from ..basic import GroupedLinear, ScaledSReLU, ScaledClampedQGeGLU from ..fuser import register_forward_fusion from ..op import FusedOperation, FusibleOperation, OperationContext @@ -26,7 +32,10 @@ _cudnn_frontend_geglu_runtime_params, _cudnn_frontend_version_supported, _cudnn_frontend_supports_grouped_gemm_srelu, + _group_quantize_for_grouped_mlp, _nvidia_cudnn_frontend_supports_wgrad, + _nvfp4_amax, + _nvfp4_single_tensor_from_grouped, fuse_grouped_mlp_ops, is_glu_activation, is_quantized_tensor, @@ -67,8 +76,8 @@ def _grouped_gemm_dsrelu_backward_supported() -> bool: return grouped_gemm_dsrelu_wrapper_sm100 is not None -class _ForwardGroupedMLP_CuTeGEMMBase_MXFP8(FusedOperation): - """Base fused op for MXFP8 GroupedLinear + activation + GroupedLinear. +class _ForwardGroupedMLP_CuTeGEMMBase(FusedOperation): + """Base fused op for block-scaled GroupedLinear + activation + GroupedLinear. Uses experimental CuTe DSL kernel from cuDNN front-end. @@ -202,6 +211,7 @@ def fuser_forward( split_sizes = split_sizes.to(dtype=torch.int64, device=device) base_split_offsets = tex.splits_to_offsets(split_sizes, 1) split_points = base_split_offsets[1:].to(dtype=torch.int) + fc1_x_tensor_offsets = base_split_offsets * fc1_weight_shape[1] fc2_x_tensor_offsets = base_split_offsets * fc2_weight_shape[1] # Extract per-row activation probabilities from the middle op. @@ -224,7 +234,7 @@ def fuser_forward( if fc1_op.weight.rowwise_data is None: raise RuntimeError("FC1 grouped weight has no rowwise_data to quantize.") fc1_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - grouped_fc1_weight = tex.group_quantize( + grouped_fc1_weight = _group_quantize_for_grouped_mlp( fc1_op.weight.rowwise_data.view(fc1_op.weight.logical_shape), fc1_weight_quantizer, num_groups, @@ -256,7 +266,7 @@ def fuser_forward( if fc2_op.weight.rowwise_data is None: raise RuntimeError("FC2 grouped weight has no rowwise_data to quantize.") fc2_weight_quantizer.set_usage(rowwise=True, columnwise=input_requires_grad) - grouped_fc2_weight = tex.group_quantize( + grouped_fc2_weight = _group_quantize_for_grouped_mlp( fc2_op.weight.rowwise_data.view(fc2_op.weight.logical_shape), fc2_weight_quantizer, num_groups, @@ -276,25 +286,45 @@ def fuser_forward( grouped_fc2_weight = quantized_fc2_weights # Some wrapper-copy paths may drop grouped storage metadata; enforce defaults. - if getattr(grouped_fc1_weight, "_with_gemm_swizzled_scales", None) is None and isinstance( - grouped_fc1_weight, GroupedTensor + if isinstance(grouped_fc1_weight, GroupedTensor) and not hasattr( + grouped_fc1_weight, "_with_gemm_swizzled_scales" ): grouped_fc1_weight._with_gemm_swizzled_scales = False - if getattr(grouped_fc2_weight, "_with_gemm_swizzled_scales", None) is None and isinstance( - grouped_fc2_weight, GroupedTensor + if isinstance(grouped_fc2_weight, GroupedTensor) and not hasattr( + grouped_fc2_weight, "_with_gemm_swizzled_scales" ): grouped_fc2_weight._with_gemm_swizzled_scales = False # Group-quantize input tensor and convert dtypes if needed fc1_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) fc1_input_quantizer.optimize_for_gemm = True - if isinstance(input_, GroupedTensor) and isinstance( - getattr(input_, "quantizer", None), MXFP8Quantizer + input_quantizer = getattr(input_, "quantizer", None) + if isinstance(input_, GroupedTensor) and ( + isinstance(fc1_input_quantizer, MXFP8Quantizer) + and isinstance(input_quantizer, MXFP8Quantizer) + or isinstance(fc1_input_quantizer, NVFP4Quantizer) + and isinstance(input_quantizer, NVFP4Quantizer) ): grouped_fc1_x = input_ else: fc1_x = maybe_dequantize(input_, dtype) - grouped_fc1_x = tex.group_quantize(fc1_x, fc1_input_quantizer, num_groups, split_sizes) + grouped_fc1_x = _group_quantize_for_grouped_mlp( + fc1_x, + fc1_input_quantizer, + num_groups, + split_sizes, + tensor_offsets=fc1_x_tensor_offsets, + ) + + use_nvfp4 = isinstance(fc1_input_quantizer, NVFP4Quantizer) or isinstance( + fc1_weight_param, NVFP4Tensor + ) + data_dtype = torch.float4_e2m1fn_x2 if use_nvfp4 else torch.float8_e4m3fn + scale_view_dtype = torch.float8_e4m3fn if use_nvfp4 else torch.float8_e8m0fnu + sf_vec_size = NVFP4_BLOCK_SCALING_SIZE if use_nvfp4 else MXFP8_BLOCK_SCALING_SIZE + data_in_k = in_shape[1] // 2 if use_nvfp4 else in_shape[1] + fc1_weight_k = fc1_weight_shape[1] // 2 if use_nvfp4 else fc1_weight_shape[1] + k_sf_divisor = 2 * sf_vec_size if use_nvfp4 else 4 * sf_vec_size # Pack data tensors # Note: Fused kernel expects tensor with non-contiguous @@ -305,20 +335,42 @@ def fuser_forward( # Data logical shape: (sum(m), k, 1) # Scale logical shape: (32 (block row), 4 (block row), # sum(m)/128, 4 (block col), k/128, 1) - fc1_x_data = grouped_fc1_x.rowwise_data.view(in_shape[0], in_shape[1]) - fc1_x_data = fc1_x_data.view(dtype=torch.float8_e4m3fn) + fc1_x_data = grouped_fc1_x.rowwise_data.view(dtype=data_dtype) + fc1_x_data = fc1_x_data.view(in_shape[0], data_in_k) fc1_x_data = fc1_x_data.unsqueeze(0).permute(1, 2, 0) fc1_x_scales = grouped_fc1_x.scale_inv - fc1_x_scales = fc1_x_scales.view(dtype=torch.float8_e8m0fnu) - fc1_x_scales = fc1_x_scales.view( - 1, - (in_shape[0] + 127) // 128, - (in_shape[1] + 127) // 128, - MXFP8_BLOCK_SCALING_SIZE, - 4, - 4, - ) - fc1_x_scales = fc1_x_scales.permute(3, 4, 1, 5, 2, 0) + fc1_x_scales = fc1_x_scales.view(dtype=scale_view_dtype) + with_gemm_swizzled_scales = grouped_fc1_x._with_gemm_swizzled_scales + if use_nvfp4 and with_gemm_swizzled_scales: + fc1_x_scales = fc1_x_scales.view( + 1, + ceil_div(in_shape[0], 128), + ceil_div(data_in_k, k_sf_divisor), + 32, + 4, + 4, + ) + fc1_x_scales = fc1_x_scales.permute(3, 4, 1, 5, 2, 0) + elif use_nvfp4: + fc1_x_scales = fc1_x_scales.view( + 1, + ceil_div(in_shape[0], 128), + 4, + 32, + ceil_div(data_in_k, k_sf_divisor), + 4, + ) + fc1_x_scales = fc1_x_scales.permute(3, 2, 1, 5, 4, 0) + else: + fc1_x_scales = fc1_x_scales.view( + 1, + ceil_div(in_shape[0], 128), + ceil_div(in_shape[1], k_sf_divisor), + 32, + 4, + 4, + ) + fc1_x_scales = fc1_x_scales.permute(3, 4, 1, 5, 2, 0) alpha_tensor = get_cached_ones_tensor(num_groups, dtype, device) norm_const_tensor = get_cached_ones_tensor(1, torch.float32, device) @@ -327,21 +379,37 @@ def fuser_forward( fc1_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc1_op) fc2_bias_packed = _pack_grouped_linear_bias_for_cudnn(fc2_op) + fc1_d_dtype = torch.bfloat16 if use_nvfp4 else torch.float8_e4m3fn + fc1_prob_tensor = ( + scales.detach().to(dtype=torch.float32 if use_nvfp4 else dtype).reshape(-1, 1, 1) + ) + fc1_norm_const_tensor = None if use_nvfp4 else norm_const_tensor + if use_nvfp4: + nvfp4_fp4_max = 6.0 + nvfp4_fp8_max = 448.0 + fc1_alpha_tensor = ( + _nvfp4_amax(grouped_fc1_x, columnwise=False) + * _nvfp4_amax(grouped_fc1_weight, columnwise=False) + / (nvfp4_fp4_max**2 * nvfp4_fp8_max**2) + ).to(torch.float32) + else: + fc1_alpha_tensor = alpha_tensor + fc1_activation_kwargs = { "a_tensor": fc1_x_data, "sfa_tensor": fc1_x_scales, "padded_offsets": split_points, - "alpha_tensor": alpha_tensor, + "alpha_tensor": fc1_alpha_tensor, "bias_tensor": fc1_bias_packed, - "norm_const_tensor": norm_const_tensor, - "prob_tensor": scales.detach().to(dtype=dtype).reshape(-1, 1, 1), + "norm_const_tensor": fc1_norm_const_tensor, + "prob_tensor": fc1_prob_tensor, "acc_dtype": torch.float32, "c_dtype": torch.bfloat16, - "d_dtype": torch.float8_e4m3fn, + "d_dtype": fc1_d_dtype, "cd_major": "n", - "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, + "sf_vec_size": sf_vec_size, "current_stream": current_stream, - "discrete_col_sfd": True, + "discrete_col_sfd": not use_nvfp4, "use_dynamic_sched": True, } if self._cudnn_act_func is not None: @@ -363,15 +431,15 @@ def fuser_forward( # Data actual shape: (num_groups, n, k) # Data logical shape: (n, k, num_groups) fc1_w_data = fc1_weight_for_gemm.rowwise_data - fc1_w_data = fc1_w_data.view(dtype=torch.float8_e4m3fn) - fc1_w_data = fc1_w_data.view(num_groups, fc1_weight_shape[0], fc1_weight_shape[1]) + fc1_w_data = fc1_w_data.view(dtype=data_dtype) + fc1_w_data = fc1_w_data.view(num_groups, fc1_weight_shape[0], fc1_weight_k) fc1_w_data = fc1_w_data.permute(1, 2, 0) - fc1_w_scales = fc1_weight_for_gemm.scale_inv.view(dtype=torch.float8_e8m0fnu) + fc1_w_scales = fc1_weight_for_gemm.scale_inv.view(dtype=scale_view_dtype) fc1_w_scales = fc1_w_scales.view( num_groups, - (fc1_weight_shape[0] + 127) // 128, - (fc1_weight_shape[1] + 127) // 128, - MXFP8_BLOCK_SCALING_SIZE, + ceil_div(fc1_weight_shape[0], 128), + ceil_div(fc1_weight_shape[1], k_sf_divisor), + 32, 4, 4, ) @@ -385,15 +453,16 @@ def fuser_forward( [w._rowwise_data for w in grouped_fc1_weight], device, ) + swizzle_type = "uniform_nvfp4_swizzle" if use_nvfp4 else "uniform_mxfp8_rowwise_swizzle" fc1_sfb_ptrs, _fc1_sfb_buffer = tex.transform_and_copy_data_ptrs_to_device( - "uniform_mxfp8_rowwise_swizzle", + swizzle_type, [w._rowwise_scale_inv for w in grouped_fc1_weight], device, ) fc1_activation_kwargs["b_ptrs"] = fc1_b_ptrs fc1_activation_kwargs["sfb_ptrs"] = fc1_sfb_ptrs fc1_activation_kwargs["n"] = fc1_weight_shape[0] - fc1_activation_kwargs["b_dtype"] = torch.float8_e4m3fn + fc1_activation_kwargs["b_dtype"] = data_dtype fc1_activation_kwargs["b_major"] = "k" fc1_kernel_out = self.grouped_gemm_activation_kernel()(**fc1_activation_kwargs) @@ -409,96 +478,173 @@ def fuser_forward( # k/128, 4 (block row), sum(m_splits)/128, 1) activation_in = fc1_kernel_out["c_tensor"] activation_in = activation_in.view(in_shape[0], fc1_weight_shape[0]) - fc2_in_row_data = fc1_kernel_out["d_tensor"] - fc2_in_row_data = fc2_in_row_data.view(in_shape[0], fc2_weight_shape[1]) - fc2_in_row_scale = fc1_kernel_out["sfd_row_tensor"] - fc2_in_row_scale = fc2_in_row_scale.permute(5, 2, 4, 0, 1, 3) - - fc2_in_col_data = fc1_kernel_out["d_col_tensor"] - fc2_in_col_data = fc2_in_col_data.view(in_shape[0], fc2_weight_shape[1]) - fc2_in_col_scale = fc1_kernel_out["sfd_col_tensor"] - fc2_in_col_scale = fc2_in_col_scale.permute(5, 2, 4, 0, 1, 3) - # Repack columnwise scales on GPU to preserve group ordering. - - # FC2 inputs scales are already swizzled/optimized for GEMM - grouped_fc2_x = GroupedTensor( - shape=(in_shape[0], fc2_weight_shape[1]), - dtype=dtype, - num_tensors=num_groups, - quantizer=fc2_input_quantizer, - data=fc2_in_row_data.reshape(-1), - columnwise_data=fc2_in_col_data.reshape(-1), - scale_inv=fc2_in_row_scale.reshape(-1), - columnwise_scale_inv=fc2_in_col_scale.reshape(-1), - first_dims=split_sizes, - tensor_offsets=fc2_x_tensor_offsets, - with_gemm_swizzled_scales=True, - ) # FC2 GEMM fc2_out_shape = in_shape[:-1] + [fc2_weight_shape[0]] fc2_scales = basic_op_extra_inputs[2][1] if fc2_op._scale_bias else None - fc2_scales_tensor = ( - fc2_scales.detach().to(dtype=torch.float32).reshape(-1, 1, 1) - if fc2_scales is not None - else torch.ones((in_shape[0], 1, 1), dtype=torch.float32, device=device) - ) - fc2_quant_kwargs = { - "a_tensor": fc1_kernel_out["d_tensor"], - "sfa_tensor": fc1_kernel_out["sfd_row_tensor"], - "padded_offsets": split_points, - "alpha_tensor": alpha_tensor, - "bias_tensor": fc2_bias_packed, - "norm_const_tensor": None, - "prob_tensor": fc2_scales_tensor, - "acc_dtype": torch.float32, - "d_dtype": dtype, - "cd_major": "n", - "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, - "current_stream": current_stream, - "use_dynamic_sched": True, - } - if fc2_op.single_grouped_weight: - # Clone and swizzle scales for GEMM (original stays unmodified for save_for_backward) - fc2_weight_for_gemm = grouped_fc2_weight.copy() - tex.grouped_swizzle_for_gemm(fc2_weight_for_gemm, rowwise=True, columnwise=False) - - fc2_w_data = fc2_weight_for_gemm.rowwise_data - fc2_w_data = fc2_w_data.view(dtype=torch.float8_e4m3fn) - fc2_w_data = fc2_w_data.view(num_groups, fc2_weight_shape[0], fc2_weight_shape[1]) - fc2_w_data = fc2_w_data.permute(1, 2, 0) - - fc2_w_scales = fc2_weight_for_gemm.scale_inv.view(dtype=torch.float8_e8m0fnu) - fc2_w_scales = fc2_w_scales.view( + if use_nvfp4: + fc2_bias_for_gemm = None + fc2_bias_scale = None + if fc2_bias_packed is not None: + fc2_bias_for_gemm = fc2_op._get_grouped_bias_for_gemm(dtype) + if fc2_scales is not None: + fc2_bias_scale = fc2_scales.reshape(-1) + if fc2_bias_scale.dtype != torch.float32: + fc2_bias_scale = fc2_bias_scale.to(dtype=torch.float32) + + fc2_in = fc1_kernel_out["d_tensor"] + fc2_in = fc2_in.view(in_shape[0], fc2_weight_shape[1]).contiguous() + fc2_input_quantizer.set_usage(rowwise=True, columnwise=weight_requires_grad) + fc2_input_quantizer.optimize_for_gemm = True + grouped_fc2_x = _group_quantize_for_grouped_mlp( + fc2_in, + fc2_input_quantizer, num_groups, - (fc2_weight_shape[0] + 127) // 128, - (fc2_weight_shape[1] + 127) // 128, - MXFP8_BLOCK_SCALING_SIZE, - 4, - 4, + split_sizes, + tensor_offsets=fc2_x_tensor_offsets, ) - fc2_w_scales = fc2_w_scales.permute(3, 4, 1, 5, 2, 0) - fc2_quant_kwargs["b_tensor"] = fc2_w_data - fc2_quant_kwargs["sfb_tensor"] = fc2_w_scales + + fc2_out_buf = torch.empty(fc2_out_shape, dtype=dtype, device=device) + if ( + num_groups == 1 + and grouped_fc2_x.columnwise_data is not None + and grouped_fc2_x.columnwise_scale_inv is not None + ): + if fc2_op.single_grouped_weight: + fc2_w_single = grouped_fc2_weight.split_into_quantized_tensors()[0] + else: + fc2_w_single = grouped_fc2_weight[0] + fc2_x_single = _nvfp4_single_tensor_from_grouped( + grouped_fc2_x, + fc2_input_quantizer, + fp4_dtype=fc2_w_single._fp4_dtype, + ) + general_gemm( + fc2_w_single, + fc2_x_single, + out_dtype=dtype, + out=fc2_out_buf, + layout="TN", + use_split_accumulator=False, + ) + if fc2_bias_packed is not None: + token_bias = ( + fc2_bias_packed.transpose(0, 1).contiguous().expand(in_shape[0], -1) + ) + if fc2_scales is not None: + fc2_out_buf = fc2_out_buf + token_bias * fc2_scales.view(-1, 1) + else: + fc2_out_buf = fc2_out_buf + token_bias + else: + fc2_out_offsets = base_split_offsets * fc2_weight_shape[0] + fc2_out_grouped = GroupedTensor( + shape=(in_shape[0], fc2_weight_shape[0]), + dtype=dtype, + num_tensors=num_groups, + quantizer=None, + data=fc2_out_buf.view(-1), + first_dims=split_sizes, + tensor_offsets=fc2_out_offsets, + ) + general_grouped_gemm_for_grouped_tensor( + grouped_fc2_weight, + grouped_fc2_x, + fc2_out_grouped, + layout="TN", + bias=fc2_bias_for_gemm, + bias_scale=fc2_bias_scale, + ) + fc2_out = fc2_out_buf else: - fc2_b_ptrs = tex.copy_data_ptrs_to_device( - [w._rowwise_data for w in grouped_fc2_weight], - device, + fc2_in_row_data = fc1_kernel_out["d_tensor"] + fc2_in_row_data = fc2_in_row_data.view(in_shape[0], fc2_weight_shape[1]) + fc2_in_row_scale = fc1_kernel_out["sfd_row_tensor"] + fc2_in_row_scale = fc2_in_row_scale.permute(5, 2, 4, 0, 1, 3) + + fc2_in_col_data = fc1_kernel_out["d_col_tensor"] + fc2_in_col_data = fc2_in_col_data.view(in_shape[0], fc2_weight_shape[1]) + fc2_in_col_scale = fc1_kernel_out["sfd_col_tensor"] + fc2_in_col_scale = fc2_in_col_scale.permute(5, 2, 4, 0, 1, 3) + + grouped_fc2_x = GroupedTensor( + shape=(in_shape[0], fc2_weight_shape[1]), + dtype=dtype, + num_tensors=num_groups, + quantizer=fc2_input_quantizer, + data=fc2_in_row_data.reshape(-1), + columnwise_data=fc2_in_col_data.reshape(-1), + scale_inv=fc2_in_row_scale.reshape(-1), + columnwise_scale_inv=fc2_in_col_scale.reshape(-1), + first_dims=split_sizes, + tensor_offsets=fc2_x_tensor_offsets, + with_gemm_swizzled_scales=True, ) - fc2_sfb_ptrs, _fc2_sfb_buffer = tex.transform_and_copy_data_ptrs_to_device( - "uniform_mxfp8_rowwise_swizzle", - [w._rowwise_scale_inv for w in grouped_fc2_weight], - device, + + fc2_scales_tensor = ( + fc2_scales.detach().to(dtype=torch.float32).reshape(-1, 1, 1) + if fc2_scales is not None + else torch.ones((in_shape[0], 1, 1), dtype=torch.float32, device=device) ) - fc2_quant_kwargs["b_ptrs"] = fc2_b_ptrs - fc2_quant_kwargs["sfb_ptrs"] = fc2_sfb_ptrs - fc2_quant_kwargs["n"] = fc2_weight_shape[0] - fc2_quant_kwargs["b_dtype"] = torch.float8_e4m3fn - fc2_quant_kwargs["b_major"] = "k" + fc2_quant_kwargs = { + "a_tensor": fc1_kernel_out["d_tensor"], + "sfa_tensor": fc1_kernel_out["sfd_row_tensor"], + "padded_offsets": split_points, + "alpha_tensor": alpha_tensor, + "bias_tensor": fc2_bias_packed, + "norm_const_tensor": None, + "prob_tensor": fc2_scales_tensor, + "acc_dtype": torch.float32, + "d_dtype": dtype, + "cd_major": "n", + "sf_vec_size": MXFP8_BLOCK_SCALING_SIZE, + "current_stream": current_stream, + "use_dynamic_sched": True, + } + + if fc2_op.single_grouped_weight: + # Clone and swizzle scales for GEMM (original stays unmodified for save_for_backward) + fc2_weight_for_gemm = grouped_fc2_weight.copy() + tex.grouped_swizzle_for_gemm(fc2_weight_for_gemm, rowwise=True, columnwise=False) + + fc2_w_data = fc2_weight_for_gemm.rowwise_data + fc2_w_data = fc2_w_data.view(dtype=torch.float8_e4m3fn) + fc2_w_data = fc2_w_data.view(num_groups, fc2_weight_shape[0], fc2_weight_shape[1]) + fc2_w_data = fc2_w_data.permute(1, 2, 0) + + fc2_w_scales = fc2_weight_for_gemm.scale_inv.view(dtype=torch.float8_e8m0fnu) + fc2_w_scales = fc2_w_scales.view( + num_groups, + ceil_div(fc2_weight_shape[0], 128), + ceil_div(fc2_weight_shape[1], 128), + MXFP8_BLOCK_SCALING_SIZE, + 4, + 4, + ) + fc2_w_scales = fc2_w_scales.permute(3, 4, 1, 5, 2, 0) + fc2_quant_kwargs["b_tensor"] = fc2_w_data + fc2_quant_kwargs["sfb_tensor"] = fc2_w_scales + else: + fc2_b_ptrs = tex.copy_data_ptrs_to_device( + [w._rowwise_data for w in grouped_fc2_weight], + device, + ) + swizzle_type = ( + "uniform_nvfp4_swizzle" if use_nvfp4 else "uniform_mxfp8_rowwise_swizzle" + ) + fc2_sfb_ptrs, _fc2_sfb_buffer = tex.transform_and_copy_data_ptrs_to_device( + swizzle_type, + [w._rowwise_scale_inv for w in grouped_fc2_weight], + device, + ) + fc2_quant_kwargs["b_ptrs"] = fc2_b_ptrs + fc2_quant_kwargs["sfb_ptrs"] = fc2_sfb_ptrs + fc2_quant_kwargs["n"] = fc2_weight_shape[0] + fc2_quant_kwargs["b_dtype"] = torch.float8_e4m3fn + fc2_quant_kwargs["b_major"] = "k" - fc2_kernel_out = self.grouped_gemm_quant_kernel()(**fc2_quant_kwargs) - fc2_out = fc2_kernel_out["d_tensor"].permute(2, 0, 1).view(fc2_out_shape).contiguous() + fc2_kernel_out = self.grouped_gemm_quant_kernel()(**fc2_quant_kwargs) + fc2_out = fc2_kernel_out["d_tensor"].permute(2, 0, 1).view(fc2_out_shape).contiguous() # Save state for backward pass if requires_grad: @@ -517,11 +663,13 @@ def fuser_forward( ) saved_grouped_fc2_x = None if recompute_srelu_fc2_x else grouped_fc2_x - # Save the input ``GroupedTensor``s themselves for the activations. - for grouped_fc_x in (grouped_fc1_x, saved_grouped_fc2_x): - if grouped_fc_x is not None: - grouped_fc_x.rowwise_data = None - grouped_fc_x.scale_inv = None + # MXFP8 wgrad only needs columnwise tiles. NVFP4 generic GEMM fallbacks + # need the full grouped tensor state, including rowwise data and amax. + if not use_nvfp4: + for grouped_fc_x in (grouped_fc1_x, saved_grouped_fc2_x): + if grouped_fc_x is not None: + grouped_fc_x.rowwise_data = None + grouped_fc_x.scale_inv = None # FC1 saved-tensor layout. # [split_sizes, base_split_offsets, split_points, @@ -586,8 +734,8 @@ def fuser_forward( return fc2_out, [(), (), ()] -class ForwardGroupedMLP_CuTeGEMMGLU_MXFP8(_ForwardGroupedMLP_CuTeGEMMBase_MXFP8): - """Fused op for MXFP8 GroupedLinear + scaled GLU + GroupedLinear.""" +class ForwardGroupedMLP_CuTeGEMMGLU(_ForwardGroupedMLP_CuTeGEMMBase): + """Fused op for block-scaled GroupedLinear + scaled GLU + GroupedLinear.""" @classmethod @functools.lru_cache(maxsize=None) @@ -598,8 +746,8 @@ def grouped_gemm_activation_kernel(cls) -> Callable: return grouped_gemm_glu_wrapper_sm100 -class ForwardGroupedMLP_CuTeGEMMUnary_MXFP8(_ForwardGroupedMLP_CuTeGEMMBase_MXFP8): - """Fused op for MXFP8 GroupedLinear + scaled unary activation + GroupedLinear.""" +class ForwardGroupedMLP_CuTeGEMMUnary(_ForwardGroupedMLP_CuTeGEMMBase): + """Fused op for block-scaled GroupedLinear + scaled unary activation + GroupedLinear.""" @classmethod @functools.lru_cache(maxsize=None) @@ -641,7 +789,7 @@ def fuse_forward_ops( return fuse_grouped_mlp_ops( ops, recipe=recipe, - fused_op_cls=ForwardGroupedMLP_CuTeGEMMGLU_MXFP8, + fused_op_cls=ForwardGroupedMLP_CuTeGEMMGLU, ) @@ -653,16 +801,18 @@ def fuse_forward_srelu_ops( ) -> list[FusibleOperation]: """Apply GroupedLinear + ScaledSReLU + GroupedLinear fusion for forward pass.""" + if recipe is None or not recipe.mxfp8(): + return ops return fuse_grouped_mlp_ops( ops, recipe=recipe, - fused_op_cls=ForwardGroupedMLP_CuTeGEMMUnary_MXFP8, + fused_op_cls=ForwardGroupedMLP_CuTeGEMMUnary, activation_op_types=(ScaledSReLU,), ) # Register fusion if available -if ForwardGroupedMLP_CuTeGEMMGLU_MXFP8.is_supported(): +if ForwardGroupedMLP_CuTeGEMMGLU.is_supported(): register_forward_fusion(fuse_forward_ops, prepend=True) -if ForwardGroupedMLP_CuTeGEMMUnary_MXFP8.is_supported(): +if ForwardGroupedMLP_CuTeGEMMUnary.is_supported(): register_forward_fusion(fuse_forward_srelu_ops, prepend=True) diff --git a/transformer_engine/pytorch/utils.py b/transformer_engine/pytorch/utils.py index 250daec67..fd8f817b3 100644 --- a/transformer_engine/pytorch/utils.py +++ b/transformer_engine/pytorch/utils.py @@ -626,8 +626,15 @@ def get_sm_count() -> int: return torch.cuda.get_device_properties(torch.cuda.current_device()).multi_processor_count +def ceil_div(numerator, denominator): + """Integer ceiling division: ``ceil(numerator / denominator)``.""" + if denominator == 0: + raise ValueError("denominator cannot be zero.") + return (numerator + denominator - 1) // denominator + + def round_up_to_nearest_multiple(value, multiple): - """Round up `value` to the next mutiple of `multiple`""" + """Round up `value` to the next multiple of `multiple`""" if multiple == 0: raise ValueError("multiple cannot be zero.") return ((value + multiple - 1) // multiple) * multiple From 1609c890ee7da24393ce62a2680e0ebb0b4eb5f6 Mon Sep 17 00:00:00 2001 From: jomitchellnv <148147880+jomitchellnv@users.noreply.github.com> Date: Mon, 1 Jun 2026 13:16:29 -0700 Subject: [PATCH 096/180] Adds GEMM Profiling Guide to TE (#2863) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * adds blog post Signed-off-by: Jonathan Mitchell * Address review comments on GEMM profiling guide Benchmark tool: - Always benchmark Dgrad separately (remove --verify-dgrad flag) - Pass measured Dgrad data to plot instead of 2x Fprop approximation - Add FP8 CurrentScaling and DelayedScaling benchmark support - Add FP8Block to shape mode (was missing, only in model-config mode) - Add --no-fp8-current and --no-fp8-delayed CLI flags Documentation: - Restructure: concise speedups.rst in features/, full tutorial in examples/ - Add device-specific precision recipes (Hopper vs Blackwell) - Add Hopper (H200) benchmark results alongside Blackwell (B300) - Remove misleading FP8 Block vs MXFP8 comparison (different target devices) - Rename "How Shapes Are Derived" to appendix, promote key sections - Convert benchmark tool references to GitHub links - Refresh all benchmark numbers with FP8 Current/Delayed columns Signed-off-by: Jonathan Mitchell * fixes failing test Signed-off-by: Jonathan Mitchell * cleanup per comments Signed-off-by: Jonathan Mitchell * greptile Signed-off-by: Jonathan Mitchell * Address review comments on speedups.rst - Define autocast vs pre-quantized modes upfront before the figures - Remove the --pre-quantize flag reference and the standalone note - Replace unclear quantization-overhead jargon with plain language - Condense the verbose "Speedup Is Shape-Dependent" section - Reword "Fprop vs Dgrad comparisons" to per-operation breakdowns - Fix benchmark_gemm.py: skip FP8 DelayedScaling in pre-quantized mode (it has no pre-quantized variant and silently fell back to the autocast path, producing a misleading bar in the pre-quantized plots) Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jonathan Mitchell * Regenerate GEMM speedup figures with DelayedScaling fix Re-ran the model-config benchmark on B300 (SM100) and H200 (SM90) with the pre-quantized DelayedScaling fix applied, and synced the numbers in speedups.rst: - B300 autocast: now includes FP8Block (1.30x); FP8Current 1.41x, FP8Delayed 1.61x, MXFP8 1.44x, NVFP4 2.03x - B300 pre-quantized: FP8Delayed bar removed, FP8Block (1.82x) added; NVFP4 3.55x - H200 autocast: FP8Current 1.57x, FP8Delayed 1.69x, FP8Block 1.41x - H200 pre-quantized: FP8Delayed removed; FP8Block dropped (no Hopper prequant support); raw FP8 1.92x Co-Authored-By: Claude Opus 4.8 (1M context) Signed-off-by: Jonathan Mitchell * Apply suggestion from @pggPL Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> --------- Signed-off-by: Jonathan Mitchell Signed-off-by: Jonathan Mitchell Signed-off-by: Jonathan Mitchell Signed-off-by: Jonathan Mitchell Signed-off-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> Co-authored-by: Jonathan Mitchell Co-authored-by: Jonathan Mitchell Co-authored-by: Jonathan Mitchell Co-authored-by: Claude Opus 4.8 (1M context) Co-authored-by: Paweł Gadziński <62263673+pggPL@users.noreply.github.com> --- benchmarks/gemm/benchmark_gemm.py | 1883 +++++++++++++++++ .../gemm_profiling/gemm_profiling.rst | 589 ++++++ .../img/b300_model_config_speedup.png | Bin 0 -> 93892 bytes .../b300_model_config_speedup_prequant.png | Bin 0 -> 94754 bytes .../img/h200_model_config_speedup.png | Bin 0 -> 89350 bytes .../h200_model_config_speedup_prequant.png | Bin 0 -> 85789 bytes .../img/b300_model_config_speedup.png | Bin 0 -> 96655 bytes .../b300_model_config_speedup_prequant.png | Bin 0 -> 93593 bytes .../img/h200_model_config_speedup.png | Bin 0 -> 88890 bytes .../h200_model_config_speedup_prequant.png | Bin 0 -> 82439 bytes .../features/low_precision_training/index.rst | 3 +- .../low_precision_training/speedups.rst | 114 + docs/index.rst | 1 + 13 files changed, 2589 insertions(+), 1 deletion(-) create mode 100644 benchmarks/gemm/benchmark_gemm.py create mode 100644 docs/examples/gemm_profiling/gemm_profiling.rst create mode 100644 docs/examples/gemm_profiling/img/b300_model_config_speedup.png create mode 100644 docs/examples/gemm_profiling/img/b300_model_config_speedup_prequant.png create mode 100644 docs/examples/gemm_profiling/img/h200_model_config_speedup.png create mode 100644 docs/examples/gemm_profiling/img/h200_model_config_speedup_prequant.png create mode 100644 docs/features/low_precision_training/gemm_profiling/img/b300_model_config_speedup.png create mode 100644 docs/features/low_precision_training/gemm_profiling/img/b300_model_config_speedup_prequant.png create mode 100644 docs/features/low_precision_training/gemm_profiling/img/h200_model_config_speedup.png create mode 100644 docs/features/low_precision_training/gemm_profiling/img/h200_model_config_speedup_prequant.png create mode 100644 docs/features/low_precision_training/speedups.rst diff --git a/benchmarks/gemm/benchmark_gemm.py b/benchmarks/gemm/benchmark_gemm.py new file mode 100644 index 000000000..2382cc339 --- /dev/null +++ b/benchmarks/gemm/benchmark_gemm.py @@ -0,0 +1,1883 @@ +# Copyright (c) 2022-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# +# See LICENSE for license information. + + +"""Unified GEMM benchmark for BF16, FP8 (Current/Delayed/Block), MXFP8, and NVFP4 precisions. + +Compares matrix-multiplication throughput across precisions using +Transformer Engine on NVIDIA GPUs. Supports two timing back-ends, +pre-quantized and autocast quantization modes, arbitrary MxKxN matrix +shapes, Nsight Systems profiling integration, and bar-chart output. + +Timing back-ends +---------------- +* **cuda-events** -- CUDA event pairs with a leading-kernel trick to + hide CPU dispatch latency. Measures the full GPU-side duration of + the timed loop (includes quantisation when using autocast mode). +* **profiler** -- ``torch.profiler`` (CUPTI) kernel timestamps. + Only the matched GEMM compute kernels (gemm, nvjet, xmma, cutlass) + are summed, giving a kernel-only measurement. + +Usage examples:: + + # Kernel-only timing via torch.profiler: + python benchmarks/gemm/benchmark_gemm.py --timing profiler --pre-quantize -o kernel.png + + # End-to-end timing via CUDA events: + python benchmarks/gemm/benchmark_gemm.py --timing cuda-events -o e2e.png + + # Custom non-square shapes: + python benchmarks/gemm/benchmark_gemm.py --shapes 88064x2560x10240,88064x10240x2560 + + # Nsight profiling of a single shape: + nsys profile --capture-range=cudaProfilerApi \\ + python benchmarks/gemm/benchmark_gemm.py --profile --profile-shape 4096 + + # Model config mode (derives all 12 GEMM shapes from hyperparameters): + python benchmarks/gemm/benchmark_gemm.py \\ + --hidden_size 4096 --intermediate_size 16384 \\ + --num_attention_heads 32 --num_hidden_layers 24 \\ + --micro_batch_size 31 --sequence_length 512 +""" + +import argparse +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Optional + +import matplotlib.pyplot as plt +import numpy as np +import torch +from torch.profiler import ProfilerActivity, profile + +try: + import transformer_engine.pytorch as te + import transformer_engine_torch as tex + from transformer_engine.common.recipe import ( + DelayedScaling, + Float8BlockScaling, + Float8CurrentScaling, + Format, + MXFP8BlockScaling, + NVFP4BlockScaling, + ) + + TE_AVAILABLE = True +except ImportError: + TE_AVAILABLE = False + + +GEMM_KERNEL_PATTERNS = ("gemm", "nvjet", "xmma", "cutlass") + +PRECISION_COLORS = { + "BF16": "#808080", + "FP8Current": "#2E8B57", + "FP8Delayed": "#20B2AA", + "FP8Block": "#006400", + "MXFP8": "#4B0082", + "NVFP4": "#B22222", +} + + +# --------------------------------------------------------------------------- +# Data structures +# --------------------------------------------------------------------------- +@dataclass +class GEMMResult: + """Single GEMM benchmark measurement.""" + + tflops: float + avg_time_ms: float + shape: tuple[int, int, int] + precision: str + + +@dataclass +class ModelConfig: + """Transformer model hyperparameters for GEMM shape derivation.""" + + hidden_size: int + intermediate_size: int + num_attention_heads: int + num_hidden_layers: int + micro_batch_size: int + sequence_length: int + + +# --------------------------------------------------------------------------- +# Hardware helpers +# --------------------------------------------------------------------------- +def is_blackwell_available() -> bool: + """Return True when the current device is Blackwell (SM100+) for NVFP4 support.""" + if not torch.cuda.is_available(): + return False + major, _ = torch.cuda.get_device_capability() + return major >= 10 + + +def compute_gemm_flops(M: int, K: int, N: int) -> int: + """Theoretical FLOP count for C = A @ B: 2 * M * N * K.""" + return 2 * M * N * K + + +# --------------------------------------------------------------------------- +# torch.profiler helpers (kernel-only timing) +# --------------------------------------------------------------------------- +def _is_gemm_kernel(name: str) -> bool: + """Return True when *name* looks like a GEMM compute kernel.""" + low = name.lower() + return any(p in low for p in GEMM_KERNEL_PATTERNS) + + +def _extract_gemm_kernel_time_us( + prof_result: profile, + num_iters: int, + verbose: bool = False, +) -> float: + """Average GEMM-kernel time in microseconds from profiler events.""" + total_us = 0.0 + count = 0 + seen: dict[str, float] = {} + + for evt in prof_result.events(): + if evt.device_type == torch.autograd.DeviceType.CUDA and _is_gemm_kernel(evt.name): + total_us += evt.device_time + count += 1 + seen[evt.name] = seen.get(evt.name, 0.0) + evt.device_time + + if verbose and seen: + print(f" Matched GEMM kernels ({count} invocations):") + for kname, kus in seen.items(): + print(f" {kname}: {kus:.0f} us total") + + if count == 0: + if verbose: + print(" WARNING: No GEMM kernels found. All CUDA events:") + for evt in prof_result.events(): + if evt.device_type == torch.autograd.DeviceType.CUDA: + print(f" {evt.name}: {evt.device_time:.0f} us") + return 0.0 + + return total_us / num_iters + + +# --------------------------------------------------------------------------- +# Timing wrappers +# --------------------------------------------------------------------------- +def _time_with_profiler( + run_fn, + num_iters: int, + flops: int, + verbose: bool = False, +) -> tuple[float, float]: + """Return (tflops, avg_ms) using torch.profiler kernel extraction.""" + with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof: + for _ in range(num_iters): + run_fn() + torch.cuda.synchronize() + + avg_us = _extract_gemm_kernel_time_us(prof, num_iters, verbose=verbose) + avg_s = avg_us / 1e6 + tflops = (flops / avg_s) / 1e12 if avg_s > 0 else 0.0 + return tflops, avg_us / 1000.0 + + +def _time_with_cuda_events( + run_fn, + num_iters: int, + flops: int, + leading_fn=None, +) -> tuple[float, float]: + """Return (tflops, avg_ms) using CUDA events with optional leading kernel.""" + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + + if leading_fn is not None: + leading_fn() + + start.record() + for _ in range(num_iters): + run_fn() + end.record() + torch.cuda.synchronize() + + avg_ms = start.elapsed_time(end) / num_iters + avg_s = avg_ms / 1000.0 + tflops = (flops / avg_s) / 1e12 if avg_s > 0 else 0.0 + return tflops, avg_ms + + +# --------------------------------------------------------------------------- +# BF16 benchmark +# --------------------------------------------------------------------------- +def benchmark_bf16( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> GEMMResult: + """Benchmark BF16 torch.matmul.""" + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + A = torch.randn(M, K, dtype=torch.bfloat16, device=device) + B = torch.randn(K, N, dtype=torch.bfloat16, device=device) + + for _ in range(num_warmup): + torch.matmul(A, B) + torch.cuda.synchronize() + + def _run(): + torch.matmul(A, B) + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + A_lg = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + B_lg = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + tflops, avg_ms = _time_with_cuda_events( + _run, num_iters, flops, leading_fn=lambda: torch.matmul(A_lg, B_lg) + ) + del A_lg, B_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="BF16") + + +# --------------------------------------------------------------------------- +# FP8 tensor-wise scaling benchmarks (CurrentScaling / DelayedScaling) +# --------------------------------------------------------------------------- +def benchmark_fp8_current( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """FP8 GEMM with Float8CurrentScaling recipe via te.Linear autocast.""" + if not TE_AVAILABLE: + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + linear = te.Linear(K, N, bias=False, params_dtype=torch.bfloat16).to(device) + x = torch.randn(M, K, dtype=torch.bfloat16, device=device) + recipe = Float8CurrentScaling() + + with te.autocast(enabled=True, recipe=recipe): + for _ in range(num_warmup): + linear(x) + torch.cuda.synchronize() + + def _run(): + linear(x) + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + lin_lg = te.Linear(4096, 4096, bias=False, params_dtype=torch.bfloat16).to(device) + x_lg = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + tflops, avg_ms = _time_with_cuda_events( + _run, num_iters, flops, leading_fn=lambda: lin_lg(x_lg) + ) + del lin_lg, x_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="FP8Current") + + +def benchmark_fp8_current_prequantized( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """Pre-quantized FP8 GEMM with Float8CurrentScaling via tex.generic_gemm.""" + if not TE_AVAILABLE: + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + try: + quantizer = te.Float8CurrentScalingQuantizer(tex.DType.kFloat8E4M3, device=device) + + A_q = quantizer.quantize(torch.randn(K, M, dtype=torch.bfloat16, device=device)) + B_q = quantizer.quantize(torch.randn(K, N, dtype=torch.bfloat16, device=device)) + D = torch.empty(N, M, dtype=torch.bfloat16, device=device) + ws_size = 32 * 1024 * 1024 + ws = torch.empty(ws_size, dtype=torch.uint8, device=device) + + def _run(): + tex.generic_gemm( + A_q, + False, + B_q, + True, + D, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + for _ in range(num_warmup): + _run() + torch.cuda.synchronize() + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + A_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + B_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + D_lg = torch.empty(4096, 4096, dtype=torch.bfloat16, device=device) + + def _lead(): + tex.generic_gemm( + A_lg_q, + False, + B_lg_q, + True, + D_lg, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + tflops, avg_ms = _time_with_cuda_events(_run, num_iters, flops, leading_fn=_lead) + del A_lg_q, B_lg_q, D_lg + + return GEMMResult( + tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="FP8Current" + ) + except Exception as e: + print(f"Warning: FP8 CurrentScaling prequantized benchmark failed: {e}") + return None + + +def benchmark_fp8_delayed( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """FP8 GEMM with DelayedScaling recipe via te.Linear autocast.""" + if not TE_AVAILABLE: + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + linear = te.Linear(K, N, bias=False, params_dtype=torch.bfloat16).to(device) + x = torch.randn(M, K, dtype=torch.bfloat16, device=device) + recipe = DelayedScaling() + + with te.autocast(enabled=True, recipe=recipe): + for _ in range(num_warmup): + linear(x) + torch.cuda.synchronize() + + def _run(): + linear(x) + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + lin_lg = te.Linear(4096, 4096, bias=False, params_dtype=torch.bfloat16).to(device) + x_lg = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + tflops, avg_ms = _time_with_cuda_events( + _run, num_iters, flops, leading_fn=lambda: lin_lg(x_lg) + ) + del lin_lg, x_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="FP8Delayed") + + +# --------------------------------------------------------------------------- +# MXFP8 benchmarks +# --------------------------------------------------------------------------- +def benchmark_fp8( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """MXFP8 GEMM via te.Linear autocast.""" + if not TE_AVAILABLE: + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + linear = te.Linear(K, N, bias=False, params_dtype=torch.bfloat16).to(device) + x = torch.randn(M, K, dtype=torch.bfloat16, device=device) + recipe = MXFP8BlockScaling(fp8_format=Format.E4M3) + + with te.autocast(enabled=True, recipe=recipe): + for _ in range(num_warmup): + linear(x) + torch.cuda.synchronize() + + def _run(): + linear(x) + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + lin_lg = te.Linear(4096, 4096, bias=False, params_dtype=torch.bfloat16).to(device) + x_lg = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + tflops, avg_ms = _time_with_cuda_events( + _run, num_iters, flops, leading_fn=lambda: lin_lg(x_lg) + ) + del lin_lg, x_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="MXFP8") + + +def benchmark_fp8_prequantized( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """Pre-quantized MXFP8 GEMM via tex.generic_gemm (raw kernel throughput).""" + if not TE_AVAILABLE: + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + try: + quantizer = te.MXFP8Quantizer(tex.DType.kFloat8E4M3) + + # tex.generic_gemm uses column-major convention: A=(K,M), B=(K,N), + # D=(N,M) with transa=False, transb=True for a logical C(M,N) GEMM. + A_q = quantizer.quantize(torch.randn(K, M, dtype=torch.bfloat16, device=device)) + B_q = quantizer.quantize(torch.randn(K, N, dtype=torch.bfloat16, device=device)) + D = torch.empty(N, M, dtype=torch.bfloat16, device=device) + ws_size = 32 * 1024 * 1024 + ws = torch.empty(ws_size, dtype=torch.uint8, device=device) + + def _run(): + tex.generic_gemm( + A_q, + False, + B_q, + True, + D, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + for _ in range(num_warmup): + _run() + torch.cuda.synchronize() + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + A_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + B_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + D_lg = torch.empty(4096, 4096, dtype=torch.bfloat16, device=device) + + def _lead(): + tex.generic_gemm( + A_lg_q, + False, + B_lg_q, + True, + D_lg, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + tflops, avg_ms = _time_with_cuda_events(_run, num_iters, flops, leading_fn=_lead) + del A_lg_q, B_lg_q, D_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="MXFP8") + except Exception as e: + print(f"Warning: FP8 prequantized benchmark failed: {e}") + return None + + +# --------------------------------------------------------------------------- +# Float8 Block-Scaling benchmarks +# --------------------------------------------------------------------------- +def benchmark_fp8_block( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """FP8 GEMM with Float8BlockScaling recipe via te.Linear autocast.""" + if not TE_AVAILABLE: + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + linear = te.Linear(K, N, bias=False, params_dtype=torch.bfloat16).to(device) + x = torch.randn(M, K, dtype=torch.bfloat16, device=device) + recipe = Float8BlockScaling(fp8_format=Format.E4M3) + + with te.autocast(enabled=True, recipe=recipe): + for _ in range(num_warmup): + linear(x) + torch.cuda.synchronize() + + def _run(): + linear(x) + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + lin_lg = te.Linear(4096, 4096, bias=False, params_dtype=torch.bfloat16).to(device) + x_lg = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + tflops, avg_ms = _time_with_cuda_events( + _run, num_iters, flops, leading_fn=lambda: lin_lg(x_lg) + ) + del lin_lg, x_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="FP8Block") + + +def benchmark_fp8_block_prequantized( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """Pre-quantized FP8 GEMM with Float8BlockScaling via tex.generic_gemm.""" + if not TE_AVAILABLE: + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + try: + quantizer = te.Float8BlockQuantizer( + tex.DType.kFloat8E4M3, + rowwise=True, + columnwise=True, + ) + + A_q = quantizer.quantize(torch.randn(K, M, dtype=torch.bfloat16, device=device)) + B_q = quantizer.quantize(torch.randn(K, N, dtype=torch.bfloat16, device=device)) + D = torch.empty(N, M, dtype=torch.bfloat16, device=device) + ws_size = 32 * 1024 * 1024 + ws = torch.empty(ws_size, dtype=torch.uint8, device=device) + + def _run(): + tex.generic_gemm( + A_q, + False, + B_q, + True, + D, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + for _ in range(num_warmup): + _run() + torch.cuda.synchronize() + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + A_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + B_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + D_lg = torch.empty(4096, 4096, dtype=torch.bfloat16, device=device) + + def _lead(): + tex.generic_gemm( + A_lg_q, + False, + B_lg_q, + True, + D_lg, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + tflops, avg_ms = _time_with_cuda_events(_run, num_iters, flops, leading_fn=_lead) + del A_lg_q, B_lg_q, D_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="FP8Block") + except Exception as e: + print(f"Warning: FP8 Block-Scaling prequantized benchmark failed: {e}") + return None + + +# --------------------------------------------------------------------------- +# NVFP4 benchmarks (Blackwell SM100+ only) +# --------------------------------------------------------------------------- +def benchmark_fp4( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """NVFP4 GEMM via te.Linear autocast (Blackwell only).""" + if not TE_AVAILABLE or not is_blackwell_available(): + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + linear = te.Linear(K, N, bias=False, params_dtype=torch.bfloat16).to(device) + x = torch.randn(M, K, dtype=torch.bfloat16, device=device) + recipe = NVFP4BlockScaling(fp4_format=Format.E2M1) + + with te.autocast(enabled=True, recipe=recipe): + for _ in range(num_warmup): + linear(x) + torch.cuda.synchronize() + + def _run(): + linear(x) + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + lin_lg = te.Linear(4096, 4096, bias=False, params_dtype=torch.bfloat16).to(device) + x_lg = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + tflops, avg_ms = _time_with_cuda_events( + _run, num_iters, flops, leading_fn=lambda: lin_lg(x_lg) + ) + del lin_lg, x_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="NVFP4") + + +def benchmark_fp4_prequantized( + M: int, + K: int, + N: int, + num_warmup: int = 10, + num_iters: int = 100, + timing: str = "cuda-events", + verbose: bool = False, +) -> Optional[GEMMResult]: + """Pre-quantized NVFP4 GEMM via tex.generic_gemm (Blackwell only).""" + if not TE_AVAILABLE or not is_blackwell_available(): + return None + + device = torch.device("cuda") + flops = compute_gemm_flops(M, K, N) + + try: + quantizer = te.NVFP4Quantizer(tex.DType.kFloat4E2M1) + + # tex.generic_gemm uses column-major convention: A=(K,M), B=(K,N), + # D=(N,M) with transa=False, transb=True for a logical C(M,N) GEMM. + A_q = quantizer.quantize(torch.randn(K, M, dtype=torch.bfloat16, device=device)) + B_q = quantizer.quantize(torch.randn(K, N, dtype=torch.bfloat16, device=device)) + D = torch.empty(N, M, dtype=torch.bfloat16, device=device) + ws_size = 32 * 1024 * 1024 + ws = torch.empty(ws_size, dtype=torch.uint8, device=device) + + def _run(): + tex.generic_gemm( + A_q, + False, + B_q, + True, + D, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + for _ in range(num_warmup): + _run() + torch.cuda.synchronize() + + if timing == "profiler": + tflops, avg_ms = _time_with_profiler(_run, num_iters, flops, verbose=verbose) + else: + A_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + B_lg_q = quantizer.quantize( + torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + ) + D_lg = torch.empty(4096, 4096, dtype=torch.bfloat16, device=device) + + def _lead(): + tex.generic_gemm( + A_lg_q, + False, + B_lg_q, + True, + D_lg, + None, + tex.DType.kBFloat16, + None, + tex.DType.kBFloat16, + False, + None, + False, + ws, + ws_size, + False, + False, + ) + + tflops, avg_ms = _time_with_cuda_events(_run, num_iters, flops, leading_fn=_lead) + del A_lg_q, B_lg_q, D_lg + + return GEMMResult(tflops=tflops, avg_time_ms=avg_ms, shape=(M, K, N), precision="NVFP4") + except Exception as e: + print(f"Warning: FP4 prequantized benchmark failed: {e}") + return None + + +# --------------------------------------------------------------------------- +# Shape helpers +# --------------------------------------------------------------------------- +def get_default_shapes() -> list[tuple[int, int, int]]: + """Default set of square matrix shapes for benchmarking.""" + return [ + (256, 256, 256), + (512, 512, 512), + (768, 768, 768), + (1024, 1024, 1024), + (1536, 1536, 1536), + (2048, 2048, 2048), + (3072, 3072, 3072), + (4096, 4096, 4096), + (6144, 6144, 6144), + (8192, 8192, 8192), + (16384, 16384, 16384), + ] + + +def parse_shapes_arg(shapes_arg: str) -> list[tuple[int, int, int]]: + """Parse ``--shapes`` into a list of (M, K, N) tuples. + + Accepts either square sizes (``1024,2048,4096``) or explicit + triplets (``8192x5120x10240,8192x10240x5120``), or a mix. + + Raises: + ValueError: On malformed input. + """ + items = [s.strip() for s in shapes_arg.split(",") if s.strip()] + if not items: + raise ValueError("Empty --shapes argument.") + + shapes: list[tuple[int, int, int]] = [] + for item in items: + if "x" in item: + parts = [p.strip() for p in item.lower().split("x")] + if len(parts) != 3: + raise ValueError(f"Invalid shape '{item}'. Expected 'MxKxN'.") + shapes.append((int(parts[0]), int(parts[1]), int(parts[2]))) + else: + size = int(item) + shapes.append((size, size, size)) + return shapes + + +def compute_gemm_shapes( + config: ModelConfig, +) -> tuple[ + list[tuple[str, int, int, int]], + list[tuple[str, int, int, int]], + list[tuple[str, int, int, int]], +]: + """Derive Fprop, Dgrad, and Wgrad GEMM shapes from a transformer model config. + + For forward Y = X @ W with shape (M, K, N): + - Dgrad: dX = dY @ Wᵀ → (M, N, K) (K and N swap) + - Wgrad: dW = Xᵀ @ dY → (K, M, N) (M moves to contraction axis) + + Returns: + (fprop_shapes, dgrad_shapes, wgrad_shapes) where each is a list of + (label, M, K, N) tuples. + """ + H = config.hidden_size + I = config.intermediate_size + M = config.micro_batch_size * config.sequence_length + + if H % config.num_attention_heads != 0: + raise ValueError( + f"hidden_size ({H}) must be divisible by " + f"num_attention_heads ({config.num_attention_heads})" + ) + + N_qkv = 3 * H + + fprop_shapes = [ + ("QKV Proj", M, H, N_qkv), + ("Attn Out", M, H, H), + ("MLP Up", M, H, I), + ("MLP Down", M, I, H), + ] + + dgrad_shapes = [ + ("QKV Proj (Dgrad)", M, N_qkv, H), + ("Attn Out (Dgrad)", M, H, H), + ("MLP Up (Dgrad)", M, I, H), + ("MLP Down (Dgrad)", M, H, I), + ] + + wgrad_shapes = [ + ("QKV Proj (Wgrad)", H, M, N_qkv), + ("Attn Out (Wgrad)", H, M, H), + ("MLP Up (Wgrad)", H, M, I), + ("MLP Down (Wgrad)", I, M, H), + ] + + return fprop_shapes, dgrad_shapes, wgrad_shapes + + +# --------------------------------------------------------------------------- +# GPU warmup +# --------------------------------------------------------------------------- +def warmup_gpu(duration_seconds: float = 5.0) -> None: + """Run sustained matmuls to stabilize GPU clocks before benchmarking.""" + print(f"Warming up GPU for {duration_seconds:.1f} seconds...") + device = torch.device("cuda") + A = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + B = torch.randn(4096, 4096, dtype=torch.bfloat16, device=device) + + torch.cuda.synchronize() + t0 = time.time() + while time.time() - t0 < duration_seconds: + for _ in range(10): + torch.matmul(A, B) + torch.cuda.synchronize() + + del A, B + torch.cuda.empty_cache() + print("GPU warmup complete.\n") + + +# --------------------------------------------------------------------------- +# Main orchestrator +# --------------------------------------------------------------------------- +def run_benchmarks( + shapes: list[tuple[int, int, int]], + num_warmup: int = 10, + num_iters: int = 100, + include_fp8_current: bool = True, + include_fp8_delayed: bool = True, + include_fp8: bool = True, + include_fp8_block: bool = True, + include_fp4: bool = True, + gpu_warmup_seconds: float = 5.0, + pre_quantize: bool = False, + timing: str = "cuda-events", + profile_shape: Optional[int] = None, +) -> dict[str, list[float]]: + """Run GEMM benchmarks for every shape and enabled precision. + + Returns: + Dict mapping precision name to a list of TFLOPS values, one per shape. + """ + results: dict[str, list[float]] = { + "BF16": [], + "FP8Current": [], + "FP8Delayed": [], + "FP8Block": [], + "MXFP8": [], + "NVFP4": [], + } + time_results: dict[str, list[float]] = { + "BF16": [], + "FP8Current": [], + "FP8Delayed": [], + "FP8Block": [], + "MXFP8": [], + "NVFP4": [], + } + + has_blackwell = is_blackwell_available() + run_fp8_current = include_fp8_current and TE_AVAILABLE + # DelayedScaling has no pre-quantized variant: the recipe differs from CurrentScaling + # only in how the scaling factor is computed each step (via amax history), which is + # exactly the work pre-quantized mode skips. Enabling it here would silently fall back + # to the autocast path and plot a misleading bar, so omit it when pre-quantizing. + run_fp8_delayed = include_fp8_delayed and TE_AVAILABLE and not pre_quantize + run_fp8 = include_fp8 and TE_AVAILABLE + run_fp8_block = include_fp8_block and TE_AVAILABLE + run_fp4 = include_fp4 and TE_AVAILABLE and has_blackwell + + gpu_name = torch.cuda.get_device_name(0) + timing_label = ( + "torch.profiler (CUPTI kernel timestamps)" if timing == "profiler" else "CUDA events" + ) + + print(f"\nGEMM Benchmark on {gpu_name}") + print(f"Timing method: {timing_label}") + print(f"Warmup iterations: {num_warmup}, Timed iterations: {num_iters}") + if pre_quantize: + print("Mode: Pre-quantized inputs (raw kernel throughput)") + else: + print("Mode: Autocast (includes quantization overhead)") + if not has_blackwell and include_fp4: + print("Note: NVFP4 requires Blackwell (SM100+), skipping FP4 benchmarks") + + if profile_shape is not None: + shapes = [(profile_shape, profile_shape, profile_shape)] + print(f"\n*** PROFILING MODE: shape {profile_shape}x{profile_shape}x{profile_shape} ***") + print( + "*** Run with: nsys profile --capture-range=cudaProfilerApi python